Update and fix SORT

Author: adipandas

examples/example_scripts/mot_TF_SSDMobileNet.py

@@ -23,14 +23,16 @@ def main(video_path, weights_path, config_path, use_gpu, tracker):
             print("Cannot read the video feed.")
             break
 
+        image = cv.resize(image, (700, 500))
+
         bboxes, confidences, class_ids = model.detect(image)
         tracks = tracker.update(bboxes, confidences, class_ids)
         updated_image = model.draw_bboxes(image.copy(), bboxes, confidences, class_ids)
 
         updated_image = draw_tracks(updated_image, tracks)
 
         cv.imshow("image", updated_image)
-        if cv.waitKey(0) & 0xFF == ord('q'):
+        if cv.waitKey(1) & 0xFF == ord('q'):
             break
 
     cap.release()
@@ -75,7 +77,7 @@ def main(video_path, weights_path, config_path, use_gpu, tracker):
     elif args.tracker == 'CentroidKF_Tracker':
         tracker = CentroidKF_Tracker(max_lost=0, tracker_output_format='mot_challenge')
     elif args.tracker == 'SORT':
-        tracker = SORT(max_lost=2, tracker_output_format='mot_challenge', iou_threshold=0.5, time_step=1)
+        tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3)
     else:
         raise NotImplementedError
 

examples/example_scripts/mot_YOLOv3.py

@@ -0,0 +1,90 @@
+import cv2 as cv
+from motrackers.detectors import YOLOv3
+from motrackers import CentroidTracker, CentroidKF_Tracker, SORT
+from motrackers.utils import draw_tracks
+
+
+def main(video_path, weights_path, config_path, labels_path, use_gpu, tracker):
+    model = YOLOv3(
+        weights_path=weights_path,
+        configfile_path=config_path,
+        labels_path=labels_path,
+        confidence_threshold=0.5,
+        nms_threshold=0.2,
+        draw_bboxes=True,
+        use_gpu=use_gpu
+    )
+
+    cap = cv.VideoCapture(video_path)
+    while True:
+        ok, image = cap.read()
+
+        if not ok:
+            print("Cannot read the video feed.")
+            break
+
+        image = cv.resize(image, (700, 500))
+
+        bboxes, confidences, class_ids = model.detect(image)
+        tracks = tracker.update(bboxes, confidences, class_ids)
+        updated_image = model.draw_bboxes(image.copy(), bboxes, confidences, class_ids)
+
+        updated_image = draw_tracks(updated_image, tracks)
+
+        cv.imshow("image", updated_image)
+        if cv.waitKey(1) & 0xFF == ord('q'):
+            break
+
+    cap.release()
+    cv.destroyAllWindows()
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description='Object detections in input video using YOLOv3 trained on COCO dataset.'
+    )
+
+    parser.add_argument(
+        '--video', '-v', type=str, default="./../video_data/cars.mp4", help='Input video path.')
+
+    parser.add_argument(
+        '--weights', '-w', type=str,
+        default="./../pretrained_models/yolo_weights/yolov3.weights",
+        help='path to weights file of YOLOv3 (`.weights` file.)'
+    )
+
+    parser.add_argument(
+        '--config', '-c', type=str,
+        default="./../pretrained_models/yolo_weights/yolov3.cfg",
+        help='path to config file of YOLOv3 (`.cfg` file.)'
+    )
+
+    parser.add_argument(
+        '--labels', '-l', type=str,
+        default="./../pretrained_models/yolo_weights/coco.names",
+        help='path to labels file of coco dataset (`.names` file.)'
+    )
+
+    parser.add_argument(
+        '--gpu', type=bool,
+        default=False, help='Flag to use gpu to run the deep learning model. Default is `False`'
+    )
+
+    parser.add_argument(
+        '--tracker', type=str, default='CentroidTracker',
+        help="Tracker used to track objects. Options include ['CentroidTracker', 'CentroidKF_Tracker', 'SORT']")
+
+    args = parser.parse_args()
+
+    if args.tracker == 'CentroidTracker':
+        tracker = CentroidTracker(max_lost=3, tracker_output_format='mot_challenge')
+    elif args.tracker == 'CentroidKF_Tracker':
+        tracker = CentroidKF_Tracker(max_lost=3, tracker_output_format='mot_challenge')
+    elif args.tracker == 'SORT':
+        tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3, time_step=1)
+    else:
+        raise NotImplementedError
+
+    main(args.video, args.weights, args.config, args.labels, args.gpu, tracker)

motrackers/kalman_tracker.py

@@ -224,9 +224,14 @@ def __init__(self, bbox, process_noise_scale=1.0, measurement_noise_scale=1.0, t
             [0, 1, 0, 0, 0, 1, 0],
             [0, 0, 1, 0, 0, 0, 1]]) * process_noise_scale
 
+        process_noise_covariance[-1, -1] *= 0.01
+        process_noise_covariance[4:, 4:] *= 0.01
+
         measurement_noise_covariance = np.eye(4) * measurement_noise_scale
+        measurement_noise_covariance[2:, 2:] *= 0.01
 
         prediction_covariance = np.ones_like(transition_matrix) * 10.
+        prediction_covariance[4:, 4:] *= 100.
 
         initial_state = np.array([bbox[0], bbox[1], bbox[2], bbox[3], 0., 0., 0.])
 

motrackers/sort_tracker.py

@@ -1,7 +1,7 @@
 import numpy as np
 from scipy.optimize import linear_sum_assignment
 from motrackers.utils.misc import iou_xywh as iou
-from motrackers.track import KFTrackSORT
+from motrackers.track import KFTrackSORT, KFTrack4DSORT
 from motrackers.centroid_kf_tracker import CentroidKF_Tracker
 
 
@@ -36,7 +36,6 @@ def assign_tracks2detection_iou(bbox_tracks, bbox_detections, iou_threshold=0.3)
         bbox_detections = bbox_detections[None, :]
 
     iou_matrix = np.zeros((bbox_tracks.shape[0], bbox_detections.shape[0]), dtype=np.float32)
-
     for t in range(bbox_tracks.shape[0]):
         for d in range(bbox_detections.shape[0]):
             iou_matrix[t, d] = iou(bbox_tracks[t, :], bbox_detections[d, :])
@@ -107,27 +106,42 @@ def __init__(
         )
 
     def _add_track(self, frame_id, bbox, detection_confidence, class_id, **kwargs):
-        self.tracks[self.next_track_id] = KFTrackSORT(
+        # self.tracks[self.next_track_id] = KFTrackSORT(
+        #     self.next_track_id, frame_id, bbox, detection_confidence, class_id=class_id,
+        #     data_output_format=self.tracker_output_format, process_noise_scale=self.process_noise_scale,
+        #     measurement_noise_scale=self.measurement_noise_scale, **kwargs
+        # )
+        self.tracks[self.next_track_id] = KFTrack4DSORT(
             self.next_track_id, frame_id, bbox, detection_confidence, class_id=class_id,
             data_output_format=self.tracker_output_format, process_noise_scale=self.process_noise_scale,
-            measurement_noise_scale=self.measurement_noise_scale, **kwargs
-        )
+            measurement_noise_scale=self.measurement_noise_scale, kf_time_step=1, **kwargs)
         self.next_track_id += 1
 
     def update(self, bboxes, detection_scores, class_ids):
         self.frame_count += 1
 
         bbox_detections = np.array(bboxes, dtype='int')
 
+        # track_ids_all = list(self.tracks.keys())
+        # bbox_tracks = []
+        # track_ids = []
+        # for track_id in track_ids_all:
+        #     bb = self.tracks[track_id].predict()
+        #     if np.any(np.isnan(bb)):
+        #         self._remove_track(track_id)
+        #     else:
+        #         track_ids.append(track_id)
+        #         bbox_tracks.append(bb)
+
         track_ids = list(self.tracks.keys())
         bbox_tracks = []
         for track_id in track_ids:
-            bbox_tracks.append(self.tracks[track_id].predict())
-        bbox_tracks = np.array(bbox_tracks)
+            bb = self.tracks[track_id].predict()
+            bbox_tracks.append(bb)
 
+        bbox_tracks = np.array(bbox_tracks)
         matches, unmatched_detections, unmatched_tracks = assign_tracks2detection_iou(
-            bbox_tracks, bbox_detections, iou_threshold=0.3
-        )
+            bbox_tracks, bbox_detections, iou_threshold=0.3)
 
         for i in range(matches.shape[0]):
             t, d = matches[i, :]

motrackers/track.py

@@ -1,5 +1,5 @@
 import numpy as np
-from motrackers.kalman_tracker import KFTrackerSORT, KFTracker2D
+from motrackers.kalman_tracker import KFTracker2D, KFTrackerSORT, KFTracker4D
 
 
 class Track:
@@ -183,9 +183,16 @@ def __init__(self, track_id, frame_id, bbox, detection_confidence, class_id=None
                          iou_score=iou_score, data_output_format=data_output_format, **kwargs)
 
     def predict(self):
+        if (self.kf.x[6] + self.kf.x[2]) <= 0:
+            self.kf.x[6] *= 0.0
+
         x = self.kf.predict()
+
+        if x[2] * x[3] < 0:
+            return np.array([np.nan, np.nan, np.nan, np.nan])
+
         w = np.sqrt(x[2] * x[3])
-        h = x[2] / w
+        h = x[2] / float(w)
         bb = np.array([x[0]-0.5*w, x[1]-0.5*h, w, h])
         return bb
 
@@ -196,6 +203,30 @@ def update(self, frame_id, bbox, detection_confidence, class_id=None, lost=0, io
         self.kf.update(z)
 
 
+class KFTrack4DSORT(Track):
+    """
+    Track based on Kalman filter tracker used for SORT MOT-Algorithm.
+    """
+    def __init__(self, track_id, frame_id, bbox, detection_confidence, class_id=None, lost=0, iou_score=0.,
+                 data_output_format='mot_challenge', process_noise_scale=1.0, measurement_noise_scale=1.0,
+                 kf_time_step=1, **kwargs):
+        self.kf = KFTracker4D(
+            bbox.copy(), process_noise_scale=process_noise_scale, measurement_noise_scale=measurement_noise_scale,
+            time_step=kf_time_step)
+        super().__init__(track_id, frame_id, bbox, detection_confidence, class_id=class_id, lost=lost,
+                         iou_score=iou_score, data_output_format=data_output_format, **kwargs)
+
+    def predict(self):
+        x = self.kf.predict()
+        bb = np.array([x[0], x[3], x[6], x[9]])
+        return bb
+
+    def update(self, frame_id, bbox, detection_confidence, class_id=None, lost=0, iou_score=0., **kwargs):
+        super().update(
+            frame_id, bbox, detection_confidence, class_id=class_id, lost=lost, iou_score=iou_score, **kwargs)
+        self.kf.update(bbox.copy())
+
+
 class KFTrackCentroid(Track):
     """
     Track based on Kalman filter used for Centroid Tracking of bounding box in MOT.

motrackers/tracker.py

@@ -115,7 +115,8 @@ def _get_tracks(tracks):
 
         outputs = []
         for trackid, track in tracks.items():
-            outputs.append(track.output())
+            if not track.lost:
+                outputs.append(track.output())
         return outputs
 
     @staticmethod

motrackers/utils/misc.py

@@ -98,8 +98,8 @@ def iou_xywh(bbox1, bbox2):
     iou: float
          intersection-over-onion of bbox1, bbox2.
     """
-    bbox1 = bbox1[0], bbox1[1], bbox1[2]+bbox1[0], bbox1[3]+bbox1[0]
-    bbox2 = bbox2[0], bbox2[1], bbox2[2] + bbox2[0], bbox2[3] + bbox2[0]
+    bbox1 = bbox1[0], bbox1[1], bbox1[0]+bbox1[2], bbox1[1]+bbox1[3]
+    bbox2 = bbox2[0], bbox2[1], bbox2[0]+bbox2[2], bbox2[1]+bbox2[3]
 
     iou_ = iou(bbox1, bbox2)