detection_tracking.py

import cv2
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort

class DetectionTracking:
    """
    This class is responsible for detection and tracking of objects in a frame.

    It uses the YOLO model for object detection and DeepSort for tracking the detected objects.

    Attributes:
        model: The YOLO model used for object detection.
        tracker: The DeepSort tracker used for tracking the detected objects.
    """

    def __init__(self):
        """
        Initializes the DetectionTracking with a YOLO model and a DeepSort tracker.
        """
        self.model = YOLO("models/yolov8x.pt")
        self.tracker = DeepSort(max_iou_distance=0.5 ,max_age=15, n_init=3, embedder='torchreid', embedder_model_name='resnet50')

    def process_frame(self, frame):
        """
        Processes a frame using the YOLO model and the DeepSort tracker.

        The YOLO model is used to detect objects in the frame. The detected objects are then tracked using the DeepSort tracker.

        Args:
            frame: The frame to be processed.

        Returns:
            The tracks of the detected objects.
        """
        CONFIDENCE_THRESHOLD = 0.5  # Increased confidence threshold
        detections = self.model(frame, classes=[0])[0]

        deepsort_detections = []
        for data in detections.boxes.data.tolist():
            x_min, y_min, x_max, y_max, confidence, class_id = data
            if float(confidence) < CONFIDENCE_THRESHOLD:
                continue

            bbox = [int(x_min), int(y_min), int(x_max) - int(x_min), int(y_max) - int(y_min)]
            deepsort_detections.append((bbox, confidence, int(class_id)))

        tracks = self.tracker.update_tracks(deepsort_detections, frame=frame)
        return tracks