main.py

import cv2

import numpy as np

from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort

from homography_tracker import MultiCameraTracker
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)


class MultiCameraObjectTracking:
    def __init__(self, v0 : str, v1 : str, calib :str, 
                 weights : str, conf : float, espc : int):
        self._video0 = cv2.VideoCapture(v0)
        self._video1 = cv2.VideoCapture(v1)
        assert self._video0.isOpened(), f"Could not open video1 source {v0}"
        assert self._video1.isOpened(), f"Could not open video1 source {v1}"

        # self._video1.set(cv2.CAP_PROP_POS_FRAMES, 17)
        self._centroides = {}
        self._espc = espc
        self._conf = conf

        calibration = np.linalg.inv(np.load(calib))
        homographies = [np.eye(3), calibration]
        

        self._detector = YOLO(weights)
        self._tracker0 =  DeepSort(max_age=100)
        self._tracker1 =  DeepSort(max_age=100)
        self._gb_tracker = MultiCameraTracker(homographies, iou_thres=0.20)


    @property
    def num_frames(self):
        num_frames1 = self._video0.get(cv2.CAP_PROP_FRAME_COUNT)
        num_frames2 = self._video1.get(cv2.CAP_PROP_FRAME_COUNT)
        num_frames = min(num_frames2, num_frames1)
        return int(num_frames)
    
    @property
    def frames(self):
        frame0 = self._video0.read()[1]
        frame1 = self._video1.read()[1]

        # return [frame1[:, :, ::-1], frame2[:, :, ::-1]]
        return frame0, frame1
    

    def release(self):
        self._video0.release()
        self._video1.release()
        


    def detections(self, frame, conf : float = None, espc : int = None, show_it : bool = True):
        # assert ((isinstance(self._espc, list) and not self._espc) or self._espc < 0 or espc is None, 
        #         f"Instances a")

        results = []

        for dt in self._detector(frame)[0].boxes.data.tolist():
            conf = float(dt[4])
            class_id = int(dt[5])

            if  conf < (conf or self._conf) and (isinstance(espc, list) and not class_id in espc) or \
                (isinstance(self._espc, list) and not class_id in self._espc) or \
                (isinstance(espc, int) and class_id != espc) or \
                (isinstance(self._espc, int) and class_id != self._espc):
               continue


            xmin, ymin, xmax, ymax = int(dt[0]), int(dt[1]), int(dt[2]), int(dt[3])
            results.append([[xmin, ymin, xmax - xmin, ymax - ymin], conf, class_id])
        

        return results
    
    def preprocessed_track(self, id, detections, frame):
        _tracks = self.__getattribute__(f"_tracker{id}").update_tracks(detections, frame= frame)
        print()
        tracks = []
        for track in _tracks:
            # if the track is not confirmed, ignore it
            if not track.is_confirmed():
                continue

            # print(track.__dict__)
            # get the track id and the bounding box
            track_id = track.track_id
            ltrb = track.to_ltrb()

            tracks.append([int(ltrb[0]), int(ltrb[1]), int(ltrb[2]), int(ltrb[3]), int(track_id)])

        return tracks

    def draw_tracks(self, image, tracks, ids_dict, src, classes=None):
        """
        Draw bounding boxes on an image and print tracking IDs for each box.

        Args:
            image: An array representing the image to draw on.
            boxes: A list of bounding boxes, where each box is a tuple of (x, y, w, h).
            ids: A list of tracking IDs, where each ID corresponds to a box in the boxes list.
        """

        def color_from_id(id):
            np.random.seed(id)
            return np.random.randint(0, 255, size=3).tolist()
        
        def draw_history(image, box, centroids, color):
            """
            Draw a bounding box and its historical centroids on an image.

            Args:
                image: An array representing the image to draw on.
                box: A tuple of (x, y, w, h) representing the bounding box to draw.
                centroids: A list of tuples representing the historical centroids of the bounding box.
            """
            # Convert the image to RGB color space
            vis = np.array(image)

            # Draw the bounding box on the image
            x1, y1, x2, y2 = np.int_(box)
            thickness = 2
            cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness)

            centroids = np.int_(centroids)
            # Draw the historical centroids on the vis
            for i, centroid in enumerate(centroids):
                if i == 0:
                    # Draw the current centroid as a circle
                    cv2.circle(vis, centroid, 2, color, thickness=-1)
                else:
                    # Draw the historical centroids as lines connecting them
                    prev_centroid = centroids[i - 1]
                    cv2.line(vis, prev_centroid, centroid, color, thickness=2)

            return vis
        
        # Convert the image to RGB color space
        vis = np.array(image)
        bboxes = tracks[:, :4]
        ids = tracks[:, 4]
        # labels = tracks[:, 5]
        self._centroides[src] = self._centroides.get(src, {})

        # Loop over each bounding box and draw it on the image
        for i, box in enumerate(bboxes):
            id = ids_dict[ids[i]]
            color = color_from_id(id)

            # Get the box coordinates
            x1, y1, x2, y2 = np.int_(box)

            # Draw the box on the image
            if self._centroides == None:
                vis = cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness=2)
            else:
                self._centroides[src][id] = self._centroides[src].get(id, [])
                self._centroides[src][id].append(((x1 + x2) // 2, (y1 + y2) // 2))
                vis = draw_history(vis, box, self._centroides[src][id], color)

            # Print the tracking ID next to the box
            if classes == None:
                text = f"person {id}"
            
            vis = cv2.putText(
                vis, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness=2
            )

        return vis
        
        




def main(opts):
    mvot = MultiCameraObjectTracking(opts.video0, opts.video1, opts.homography, opts.weights,
                                   opts.confidence, opts.especific_class)
    
    frame_time = 1 / mvot.num_frames
    for idx in range(mvot.num_frames):
        # Get frames
        f0, f1= mvot.frames
        


        d0 = mvot.detections(f0)
        d1 = mvot.detections(f1)

        t0 = mvot.preprocessed_track(0, d0, frame=f0)
        t1 = mvot.preprocessed_track(1, d1, frame=f1)
        # print(t0, t1)
        id0, id1 = mvot._gb_tracker.update([t0, t1])

        # print(id0, id1)

        # f0 = mvot.view_tracking(t0, id0, f0)
        # f1 = mvot.view_tracking(t1, id1, f1)
        if t0:
            f0 = mvot.draw_tracks(f0, np.array(t0), id0, 0)
        if t1:
            f1 = mvot.draw_tracks(f1, np.array(t1), id1, 1)
        

        vis = np.hstack([f0[:, :, ::-1], f1[:, :, ::-1]])
        
        cv2.namedWindow("Vis", cv2.WINDOW_NORMAL)
        cv2.imshow("Vis", vis)
        key = cv2.waitKey(1)

        if key == ord("q"):
            break

    mvot.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--video0", type=str, default="./multiple_views/4p-c0.avi")
    parser.add_argument("--video1", type=str, default="./multiple_views/4p-c1.avi")
    parser.add_argument("--homography", type=str, default="./matrix.npy")
    parser.add_argument("--weights", type=str, default="yolo11n.pt")
    parser.add_argument("--confidence", type=float, default=0.8)
    parser.add_argument("--especific_class", type=float, default=0) #People class has ID 0

    opts = parser.parse_args()

    main(opts)