Spaces:
Runtime error
Runtime error
File size: 2,586 Bytes
20c36d1 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import numpy as np
import torch
from deep_sort_realtime.deepsort_tracker import DeepSort
from datetime import datetime
class DeepSORTTracker:
def __init__(self, max_age=5):
"""Initialize DeepSORT tracker."""
self.max_age = max_age
self.tracker = self._create_tracker()
def _create_tracker(self):
"""Create a new instance of DeepSort tracker."""
return DeepSort(
max_age=self.max_age,
embedder='mobilenet', # Using default MobileNetV2 embedder
today=datetime.now().date() # For track naming and daily ID reset
)
def reset(self):
"""Reset the tracker state by creating a new instance."""
print("Resetting DeepSORT tracker...")
self.tracker = self._create_tracker()
def update(self, frame, detections):
"""Update tracking with new detections.
Args:
frame: Current video frame (numpy array)
detections: List of (box, keyword) tuples where box is [x1, y1, x2, y2] normalized
Returns:
List of (box, keyword, track_id) tuples
"""
if not detections:
return []
height, width = frame.shape[:2]
# Convert normalized coordinates to absolute and format detections
detection_list = []
for box, keyword in detections:
x1 = int(box[0] * width)
y1 = int(box[1] * height)
x2 = int(box[2] * width)
y2 = int(box[3] * height)
w = x2 - x1
h = y2 - y1
# Format: ([left,top,w,h], confidence, detection_class)
detection_list.append(([x1, y1, w, h], 1.0, keyword))
# Update tracker
tracks = self.tracker.update_tracks(detection_list, frame=frame)
# Convert back to normalized coordinates with track IDs
tracked_objects = []
for track in tracks:
if not track.is_confirmed():
continue
ltrb = track.to_ltrb() # Get [left,top,right,bottom] format
x1, y1, x2, y2 = ltrb
# Normalize coordinates
x1 = max(0.0, min(1.0, x1 / width))
y1 = max(0.0, min(1.0, y1 / height))
x2 = max(0.0, min(1.0, x2 / width))
y2 = max(0.0, min(1.0, y2 / height))
tracked_objects.append(([x1, y1, x2, y2], track.det_class, track.track_id))
return tracked_objects |