Spaces:
Runtime error
Runtime error
File size: 1,949 Bytes
bf69e73 |
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 |
from deep_sort.tracker import Tracker as DeepSortTracker
from tools import generate_detections as gdet
from deep_sort import nn_matching
from deep_sort.detection import Detection
import numpy as np
class Tracker:
tracker = None
encoder = None
tracks = None
def __init__(self):
max_cosine_distance = 0.4
nn_budget = None
encoder_model_filename = 'model_data/mars-small128.pb'
metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
self.tracker = DeepSortTracker(metric)
self.encoder = gdet.create_box_encoder(encoder_model_filename, batch_size=1)
def update(self, frame, detections):
if len(detections) == 0:
self.tracker.predict()
self.tracker.update([])
self.update_tracks()
return
bboxes = np.asarray([d[:-2] for d in detections])
bboxes[:, 2:] = bboxes[:, 2:] - bboxes[:, 0:2]
scores = [d[-1] for d in detections]
class_ids = [d[-2] for d in detections]
features = self.encoder(frame, bboxes)
dets = []
for bbox_id, bbox in enumerate(bboxes):
dets.append(Detection(bbox, scores[bbox_id], class_ids[bbox_id], features[bbox_id]))
self.tracker.predict()
self.tracker.update(dets)
self.update_tracks()
def update_tracks(self):
tracks = []
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
class_id = track.get_class()
id = track.track_id
tracks.append(Track(id, bbox, class_id))
self.tracks = tracks
class Track:
track_id = None
bbox = None
class_id = None
def __init__(self, id, bbox, class_id):
self.track_id = id
self.bbox = bbox
self.class_id = class_id
|