diff --git a/.gitignore b/.gitignore index 90a0bb3..46711f3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ output/ mot_benchmark +__pycache__ diff --git a/requirements.txt b/requirements.txt index b70813b..f3e3cfc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ +wheel==0.34.2 +numpy filterpy==1.4.5 scikit-image==0.17.2 lap==0.4.0 diff --git a/sort.py b/sort.py index 4b2ca8a..f901e02 100644 --- a/sort.py +++ b/sort.py @@ -20,7 +20,6 @@ import os import numpy as np import matplotlib -matplotlib.use('TkAgg') import matplotlib.pyplot as plt import matplotlib.patches as patches from skimage import io @@ -96,6 +95,7 @@ class KalmanBoxTracker(object): This class represents the internal state of individual tracked objects observed as bbox. """ count = 0 + countReset = 10000 def __init__(self,bbox): """ Initialises a tracker using initial bounding box. @@ -114,11 +114,14 @@ def __init__(self,bbox): self.kf.x[:4] = convert_bbox_to_z(bbox) self.time_since_update = 0 self.id = KalmanBoxTracker.count - KalmanBoxTracker.count += 1 + KalmanBoxTracker.count = (KalmanBoxTracker.count+1) % KalmanBoxTracker.countReset self.history = [] self.hits = 0 self.hit_streak = 0 self.age = 0 + self.original_id = bbox[5] + self.original_conf = bbox[4] + def update(self,bbox): """ @@ -128,6 +131,8 @@ def update(self,bbox): self.history = [] self.hits += 1 self.hit_streak += 1 + self.original_id = bbox[5] + self.original_conf = bbox[4] self.kf.update(convert_bbox_to_z(bbox)) def predict(self): @@ -207,7 +212,7 @@ def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3): self.trackers = [] self.frame_count = 0 - def update(self, dets=np.empty((0, 5))): + def update(self, dets=np.empty((0, 6))): """ Params: dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...] @@ -218,12 +223,12 @@ def update(self, dets=np.empty((0, 5))): """ self.frame_count += 1 # get predicted locations from existing trackers. - trks = np.zeros((len(self.trackers), 5)) + trks = np.zeros((len(self.trackers), 7)) to_del = [] ret = [] for t, trk in enumerate(trks): pos = self.trackers[t].predict()[0] - trk[:] = [pos[0], pos[1], pos[2], pos[3], 0] + trk[:] = [pos[0], pos[1], pos[2], pos[3], 0, 0, 0] if np.any(np.isnan(pos)): to_del.append(t) trks = np.ma.compress_rows(np.ma.masked_invalid(trks)) @@ -243,14 +248,14 @@ def update(self, dets=np.empty((0, 5))): for trk in reversed(self.trackers): d = trk.get_state()[0] if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits): - ret.append(np.concatenate((d,[trk.id+1])).reshape(1,-1)) # +1 as MOT benchmark requires positive + ret.append(np.concatenate((d,[trk.id+1], [trk.original_id], [trk.original_conf])).reshape(1,-1)) # +1 as MOT benchmark requires positive i -= 1 # remove dead tracklet if(trk.time_since_update > self.max_age): self.trackers.pop(i) if(len(ret)>0): return np.concatenate(ret) - return np.empty((0,5)) + return np.empty((0,7)) def parse_args(): """Parse input arguments.""" @@ -309,7 +314,7 @@ def parse_args(): plt.title(seq + ' Tracked Targets') start_time = time.time() - trackers = mot_tracker.update(dets) + trackers = mot_tracker.update(np.hstack((dets, np.zeros((dets.shape[0], 1))))) cycle_time = time.time() - start_time total_time += cycle_time