使用yolov8+SORT进行对象检测和跟踪时如何获取类和ID？

Question

我是 yolo 和 SORT 的新手。 我正在尝试找出有多少人没有这些防护装备。但仅使用 yolo 就会给出每一帧的输出，并且不会跟踪人员。 然后我尝试实现 SORT，但现在陷入了如何检测 SORT 检测到的 ID 类别的困境。我尝试保存类数组，然后迭代 SORT 结果，但在 SORT 处理它之前和之后，类的排列方式不同。

class_names = ['Hardhat', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest', 'Person', 'Safety Cone', 'Safety Vest', 'machinery', 'vehicle']

def process_video(video_path: str, model: YOLO):
    tracker = Sort(max_age=2000,min_hits= 3,iou_threshold=0.01)
    cap = cv2.VideoCapture(video_path)
    
    while True:
        # read a frame from the video
        ret, img = cap.read()
        if not ret:
            break

        # process the frame with YOLO
        results = model(img, stream=True)

        detections = np.empty((0,5))
        for r in results:
            boxes = r.boxes
            i = 0
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                w,h = x2-x1 , y2-y1
                # getting confidence level
                conf = math.ceil((box.conf[0] * 100))
                # class names
                cls = class_names[int(box.cls[0])]
                if cls in ['Person','NO-Hardhat', 'NO-Mask', 'NO-Safety Vest']:
                    currentArray = np.array([x1,y1,x2,y2,conf])  #check conf[0]
                    detections = np.vstack((detections,currentArray))
        resultsTracker = tracker.update(detections)

        for result in resultsTracker:
                x1, y1, x2, y2, id = result
                x1, y1, x2, y2, id = int(x1), int(y1), int(x2), int(y2), int(id)
                (text_width, text_height), _ = cv2.getTextSize(f'{id}', cv2.FONT_HERSHEY_PLAIN, fontScale=0.5, thickness=1)
                text_offset_x = x1
                text_offset_y = y1 - text_height
                cv2.rectangle(img,(x1,y1),(x2,y2),(000,000,255),1)
                cv2.putText(img, f'{id} ', (text_offset_x, text_offset_y+6), cv2.FONT_HERSHEY_PLAIN, fontScale=0.5, color=(255, 255, 255), thickness=1)
  
    
        cv2.imshow('video',img)
        if cv2.waitKey(1) & 0xFF == ord('q'): # press 'q' to quit
            break
        elif cv2.waitKey(0):
            pass
    cap.release()
    cv2.destroyAllWindows()

如果我的方法看起来错误，请纠正我。欢迎任何建议。谢谢你

Answer 1

您是否尝试过使用

model.track()

代替

model()

？这意味着您不必管理跟踪本身，YOLOv8 会为您做这件事。因此，您无需将两个列表匹配在一起，而是可以从单个结果列表中进行跟踪和分类。请参阅https://docs.ultralytics.com/modes/track/#tracking