Skip to content

Commit

Permalink
Merge pull request #1049 from rolson24/fix-pr-bug
Browse files Browse the repository at this point in the history
fix enumerate call
  • Loading branch information
SkalskiP authored Mar 28, 2024
2 parents 0a5fb7b + a2aca62 commit 908a1b4
Showing 1 changed file with 8 additions and 15 deletions.
23 changes: 8 additions & 15 deletions supervision/tracker/byte_tracker/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -272,11 +272,8 @@ def callback(frame: np.ndarray, index: int) -> np.ndarray:
"""

tensors = detections2boxes(detections=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])
Expand All @@ -286,20 +283,16 @@ def callback(frame: np.ndarray, index: int) -> np.ndarray:
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]
)
detections.tracker_id = np.full(len(detections), -1, dtype=int)
for i_detection, i_track in matches:
detections.tracker_id[i_detection] = int(tracks[i_track].track_id)

return detections[detections.tracker_id != -1]

else:
final_detections.tracker_id = np.array([], dtype=int)
detections.tracker_id = np.array([], dtype=int)

return final_detections
return detections

def reset(self):
"""
Expand Down

0 comments on commit 908a1b4

Please sign in to comment.