From 1eb877c4d924ce620bf5acbdf7ce46bfb7e8f917 Mon Sep 17 00:00:00 2001
From: Raif Olson <raif.t.olson.24@dartmouth.edu>
Date: Sat, 23 Mar 2024 18:53:05 -0500
Subject: [PATCH] Fix issue #754

---
 supervision/tracker/byte_tracker/core.py | 49 ++++++++++++++----------
 1 file changed, 29 insertions(+), 20 deletions(-)

diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py
index 89494ee6..9156fb58 100644
--- a/supervision/tracker/byte_tracker/core.py
+++ b/supervision/tracker/byte_tracker/core.py
@@ -3,6 +3,7 @@ from typing import List, Tuple
 import numpy as np
 
 from supervision.detection.core import Detections
+from supervision.detection.utils import box_iou_batch
 from supervision.tracker.byte_tracker import matching
 from supervision.tracker.byte_tracker.basetrack import BaseTrack, TrackState
 from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
@@ -270,27 +271,35 @@ class ByteTrack:
             ```
         """
 
-        tracks = self.update_with_tensors(
-            tensors=detections2boxes(detections=detections)
-        )
-        detections = Detections.empty()
-        if len(tracks) > 0:
-            detections.xyxy = np.array(
-                [track.tlbr for track in tracks], dtype=np.float32
-            )
-            detections.class_id = np.array(
-                [int(t.class_ids) for t in tracks], dtype=int
-            )
-            detections.tracker_id = np.array(
-                [int(t.track_id) for t in tracks], dtype=int
-            )
-            detections.confidence = np.array(
-                [t.score for t in tracks], dtype=np.float32
-            )
-        else:
-            detections.tracker_id = np.array([], dtype=int)
+        tensors = detections2boxes(detections=detections)
 
-        return detections
+        tracks = self.update_with_tensors(tensors=tensors)
+
+        final_detections = Detections.empty()
+
+        if len(tracks) > 0:
+            detection_bounding_boxes = np.asarray([det[:4] for det in tensors])
+            track_bounding_boxes = np.asarray([track.tlbr for track in tracks])
+
+            ious = box_iou_batch(detection_bounding_boxes, track_bounding_boxes)
+
+            iou_costs = 1 - ious
+
+            matches, _, _ = matching.linear_assignment(iou_costs, 0.5)
+            for i, idet, itrack in enumerate(matches):
+                if i == 0:
+                    final_detections = detections[[idet]]
+                    final_detections.tracker_id[0] = int(tracks[itrack].track_id)
+                else:
+                    current_detection = detections[[idet]]
+                    current_detection.tracker_id[0] = int(tracks[itrack].track_id)
+                    final_detections = Detections.merge(
+                        [final_detections, current_detection]
+                    )
+        else:
+            final_detections.tracker_id = np.array([], dtype=int)
+
+        return final_detections
 
     def reset(self):
         """