yolo로 '사람 추적' 및 '바운딩 박스' 좌표 저장하기

1. 우선 conda activate를 해준다.

• ( pip install ultralytics opencv-python )

 

 

2. 다음과 같이 파이썬 파일을 생성 및 실행  ( 이는 초안...)

from ultralytics import YOLO
import csv
import numpy as np

# ==========================
# 설정
# ==========================
VIDEO_PATH = "output.mp4"       # 추적할 영상
MODEL_PATH = "yolov8n.pt"       # YOLO 가중치 (다른 걸 써도 됨)
CSV_PATH   = "output_boxes.csv" # 바운딩 박스 좌표를 저장할 CSV 파일 이름

CONF_THRES = 0.4                # confidence threshold
TRACKER_CFG = "bytetrack.yaml"  # 기본 ByteTrack 설정


def main():
    # 1) YOLO 모델 로드
    model = YOLO(MODEL_PATH)

    # 2) CSV에 기록할 리스트 (나중에 한 번에 저장)
    rows = []  # 각 줄: [frame_idx, track_id, class_id, conf, x1, y1, x2, y2]

    # 3) frame index 초기화
    frame_idx = 0

    # 4) YOLO tracking 실행 (stream=True 로 프레임 단위로 결과 받기)
    results = model.track(
        source=VIDEO_PATH,
        conf=CONF_THRES,
        tracker=TRACKER_CFG,
        classes=[0],      # 사람만 추적 (COCO class 0 = person). 전체 클래스를 원하면 이 줄 삭제
        stream=True,      # 프레임 단위로 결과 뽑기
        save=True,        # 추적 결과 영상도 저장 (runs/track/... 아래)
        show=False,       # 실시간으로 보여주고 싶으면 True
        persist=True      # 동일한 ID 유지
    )

    for r in results:
        boxes = r.boxes

        if boxes is None or len(boxes) == 0:
            # 이 프레임에는 검출된 객체가 없음
            frame_idx += 1
            continue

        # GPU -> CPU, numpy 로 변환
        xyxy = boxes.xyxy.cpu().numpy()   # (N, 4)  [x1, y1, x2, y2]
        confs = boxes.conf.cpu().numpy()  # (N,)
        clss  = boxes.cls.cpu().numpy()   # (N,)
        ids   = boxes.id                  # (N,) or None

        if ids is not None:
            ids = ids.cpu().numpy()
        else:
            # 추적 ID가 아직 없는 경우 -1로 넣어둠
            ids = np.full(len(xyxy), -1)

        # 각 객체별로 CSV에 들어갈 row 생성
        for (x1, y1, x2, y2), conf, cls_id, tid in zip(xyxy, confs, clss, ids):
            row = [
                frame_idx,           # 몇 번째 프레임인지
                int(tid),            # track_id (ByteTrack에서 부여한 ID)
                int(cls_id),         # class id (사람이면 0)
                float(conf),         # confidence
                float(x1), float(y1),
                float(x2), float(y2)
            ]
            rows.append(row)

        frame_idx += 1

    # 5) CSV 파일로 저장
    with open(CSV_PATH, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["frame", "track_id", "class_id", "conf", "x1", "y1", "x2", "y2"])
        writer.writerows(rows)

    print(f"바운딩 박스 좌표를 '{CSV_PATH}' 에 저장했습니다.")
    print("추적 결과 영상은 'runs/track/' (또는 'runs/detect/track/') 폴더 아래에 생성됩니다.")


if __name__ == "__main__":
    main()

 

 

3. 다음과 같이 파이썬 파일을 생성 및 실행  ( 이는  x1 ,y1, x2, y2, center_X , center_y , width , height  뽑아냄)

from ultralytics import YOLO
import csv
import numpy as np

# ==========================
# 설정
# ==========================
VIDEO_PATH = "output.mp4"       # 추적할 영상
MODEL_PATH = "yolov8n.pt"       # YOLO 가중치 (다른 걸 써도 됨)
CSV_PATH   = "output_boxes.csv" # 바운딩 박스 좌표를 저장할 CSV 파일 이름

CONF_THRES = 0.4                # confidence threshold
TRACKER_CFG = "bytetrack.yaml"  # 기본 ByteTrack 설정
SCENE_NAME = "school"


def main():
    # 1) YOLO 모델 로드
    model = YOLO(MODEL_PATH)

    # 2) CSV에 기록할 리스트 (나중에 한 번에 저장)
    rows = []  # 각 줄: [frame_idx, track_id, class_id, conf, x1, y1, x2, y2]

    # 3) frame index 초기화
    frame_idx = 0

    # 6) YOLO tracking 실행 (stream=True 로 프레임 단위로 결과 받기)
    results = model.track(
        source=VIDEO_PATH,
        conf=CONF_THRES,
        tracker=TRACKER_CFG,
        classes=[0],      # 사람만 추적 (COCO class 0 = person). 전체 클래스를 원하면 이 줄 삭제
        stream=True,      # 프레임 단위로 결과 뽑기... 사실상 이를 false를 사용하는 경우는 거의 없다.
        save=True,        # 추적 결과 영상도 저장 (runs/track/... 아래)
        show=False,       # 실시간으로 보여주고 싶으면 True. 즉, True는 새로운 창이 뜨면서 해당 박스가 실시간으로 처리 되는것을 보여줌.
        persist=True      # 동일한 ID 유지
    )

# yolo모델을 가지고 model.track을 한경우 results라는 객체를 생성한다.
# 이는 해당프레임에서 찾아낸 모든 정보가 담긴 객체이다. 
# 아래에서는 results에서 하나씩 꺼내서 쓰며, r.boxes를 해서 바운딩 박스 정보를 가져옴.
# boxes.xyxy는 좌상단(x1,y1) , 우하단(x2,y2) 좌표가 담겨짐.
# boxes.id는 각 물체에 부여된 고유 ID
# boxes.cinf는 확신하는 정도 ( 0.0 ~ 1.0 )
# boxes.cls는 클래스 번호... 사람은 0이다.
# boxes.xywh는 중심점(x,y)와 width height좌표
    for r in results:
        boxes = r.boxes

        if boxes is None or len(boxes) == 0:
            # 이 프레임에는 검출된 객체가 없음
            frame_idx += 1
            continue

        # GPU -> CPU, numpy 로 변환
        xyxy = boxes.xyxy.cpu().numpy()   # (N, 4)  [x1, y1, x2, y2]
        xywh = boxes.xywh.cpu().numpy()   # (N, 4)  [ 박스의 중심좌표 , 박스의 너비 , 박스의 높이 ]
        confs = boxes.conf.cpu().numpy()  # (N,)
        clss  = boxes.cls.cpu().numpy()   # (N,)
        ids   = boxes.id                  # (N,) or None

        if ids is not None:
            ids = ids.cpu().numpy()
        else:
            # 추적 ID가 아직 없는 경우 -1로 넣어둠
            ids = np.full(len(xyxy), -1)

        # 각 객체별로 CSV에 들어갈 row 생성
        for (x1, y1, x2, y2), (x,y,w,h) ,conf, cls_id, tid in zip(xyxy, xywh , confs, clss, ids):
            row = [
                frame_idx,           # 몇 번째 프레임인지
                int(tid),            # track_id (ByteTrack에서 부여한 ID)
                int(cls_id),         # class id (사람이면 0)
                float(conf),         # confidence
                float(x1), float(y1),
                float(x2), float(y2),
                float(x) , float(y),
                SCENE_NAME
            ]
            rows.append(row)

        frame_idx += 1

    # 5) CSV 파일로 저장
    with open(CSV_PATH, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["frame", "track_id", "class_id", "conf", "x1", "y1", "x2", "y2","x","y","sceneId"])
        writer.writerows(rows)

    print(f"바운딩 박스 좌표를 '{CSV_PATH}' 에 저장했습니다.")
    print("추적 결과 영상은 'runs/track/' (또는 'runs/detect/track/') 폴더 아래에 생성됩니다.")


if __name__ == "__main__":
    main()

track_and_save_boxes.py

0.00MB