File size: 764 Bytes
5f1587b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Pipeline:
    def __init__(self, detector, tracker):
        """
        Initialize the Pipeline class with a detector and tracker.

        Args:
            detector (object): The object detection model.
            tracker (object): The object tracking model.
        """
        self.detector = detector
        self.tracker = tracker

    def load_state_dict(self, onnx_path):
        self.detector.load_state_dict(onnx_path)

    def __call__(self, frame):
        """
        Run the detection and tracking on the input image.

        Args:
            frame (np.ndarray): The input image to process.

        Returns:
            supervision.Detections: Detections object after tracking.
        """
        return self.tracker(self.detector(frame))