I would like to ask for advice on difficulties with object tracking.

[pewang2187]

Body

The objects I'm tracking all look the same, but I want to identify them differently, but I can't figure it out. I'd like to ask for advice.

import cv2
import torch
from ultralytics import YOLO
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort
import numpy as np

YOLOv8 모델 로드

model = YOLO("C:/Users/Pc/Downloads/person.v7i.yolov8/best.pt")

DeepSORT 설정

cfg = get_config()
cfg.merge_from_file("C:/Users/Pc/Downloads/person.v7i.yolov8/deep_sort_pytorch/configs/deep_sort.yaml")
deepsort = DeepSort(
"C:/Users/Pc/Downloads/person.v7i.yolov8/deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7",
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True
)

고정 색상 매핑 (ID 1~5)

TRACK_COLORS = {
1: (255, 0, 0), # 파랑
2: (0, 255, 0), # 초록
3: (0, 0, 255), # 빨강
4: (255, 255, 0), # 노랑
5: (255, 0, 255) # 핑크
}

비디오 로드

cap = cv2.VideoCapture("C:/Users/Pc/Downloads/person.v7i.yolov8/way.mp4")
if not cap.isOpened():
print("Error: Cannot open video file.")
exit()

객체 ID와 색상을 매핑하기 위한 딕셔너리

track_id_to_color = {}

while True:
ret, frame = cap.read()
if not ret:
print("End of video.")
break

# YOLOv8 객체 탐지
results = model(frame, conf=0.5)
bbox_xywh = []
confs = []
classes = []

for r in results:
    for box in r.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        conf = float(box.conf[0])
        cls = int(box.cls[0])

        if model.names[cls] == "person":
            cx = (x1 + x2) / 2
            cy = (y1 + y2) / 2
            w = x2 - x1
            h = y2 - y1

            bbox_xywh.append([cx, cy, w, h])
            confs.append(conf)
            classes.append(cls)

            # YOLO 바운딩 박스 (초록색)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, f"YOLO Conf: {conf:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# DeepSORT 업데이트
if len(bbox_xywh) > 0:
    xywhs = torch.Tensor(bbox_xywh)
    confss = torch.Tensor(confs)
    classes_tensor = torch.Tensor(classes).int()

    outputs = deepsort.update(xywhs, confss, classes_tensor, frame)

    # DeepSORT 객체 추적 (고유 색상 적용)
    if outputs is not None and len(outputs) > 0:
        for output in outputs:
            try:
                if len(output) >= 5:
                    x1, y1, x2, y2, track_id = map(int, output[:5])

                    # 객체 ID가 색상 매핑에 없으면 기본 색상으로 설정
                    if track_id not in track_id_to_color:
                        track_id_to_color[track_id] = TRACK_COLORS.get(
                            (track_id % 5) + 1, (0, 255, 255))  # 순환 색상

                    color = track_id_to_color[track_id]
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(frame, f"Track ID: {track_id}", (x1, y1 - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                else:
                    print(f"Invalid DeepSORT Output Format: {output}")
            except Exception as e:
                print(f"Error processing output: {output}, Error: {e}")
else:
    print("No valid detections for DeepSORT.")

cv2.putText(frame, "YOLO: Green (Detection)", (20, 20),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
cv2.putText(frame, "DeepSORT: Unique Colors (Tracking)", (20, 50),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

cv2.imshow("YOLOv8 + DeepSORT", frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
    break

cap.release()
cv2.destroyAllWindows()

Guidelines

• I have read and understood this category's guidelines before making this post.