Как получить классы вместе с идентификаторами при использовании yolov8+SORT для обнаружения и отслеживания объектов?

Question

Я новичок в йоло и СОРТИРОВКЕ. Я пытаюсь определить, сколько людей не имеют этих защитных средств. Но простое использование yolo выдаст выходные данные для каждого кадра и не будет отслеживать людей. Затем я попробовал реализовать 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 .