Add New Tracker

Author: adipandas

motrackers/__init__.py

@@ -1,10 +1,11 @@
 """
-Multi Object Trackers
+Multi-object Trackers in Python
 
 Author: Aditya M. Deshpande
-
-Website:  http://adipandas.github.io/
+Blog: http://adipandas.github.io/
+Github: adipandas
 """
 
 from motrackers.simple_tracker import SimpleTracker
 from motrackers.simple_tracker2 import SimpleTracker2
+from motrackers.iou_tracker import IOUTracker

motrackers/iou_tracker.py

@@ -0,0 +1,90 @@
+"""
+Implementation of this algorithm is heavily based on the following:
+https://github.com/bochinski/iou-tracker
+"""
+
+import numpy as np
+from motrackers.utils import get_centroids, iou
+from motrackers import SimpleTracker2
+
+
+class IOUTracker(SimpleTracker2):
+    def __init__(self, max_lost=2, iou_threshold=0.5, min_detection_confidence=0.4, max_detection_confidence=0.7):
+        self.iou_threshold = iou_threshold
+        self.max_detection_confidence = max_detection_confidence
+        self.min_detection_confidence = min_detection_confidence
+        super(IOUTracker, self).__init__(max_lost=max_lost)
+
+    def update(self, bboxes: list, class_ids: list, detection_scores: list):
+        """
+        Update the tracker based on the new bboxes as input.
+
+        Parameters
+        ----------
+        bboxes : list
+                 List of bounding boxes detected in the current frame/timestep. Each element of the list represent
+                 coordinates of bounding box as tuple (top-left-x, top-left-y, bottom-right-x, bottom-right-y).
+        class_ids : list
+                    List of class_ids (int) corresponding to labels of the detected object. Default is `None`.
+        detection_scores: list
+                         List of detection scores / probability of each detected object or objectness.
+
+        Returns
+        -------
+        outputs : list
+                 List of tracks being currently tracked by the tracker.
+                 Each element of this list contains the tuple in
+                 format (frame#, trackid, class_id, centroid, bbox, info_dict).
+                 class_id is the id for label of the detection.
+                 centroid represents the pixel coordinates of the centroid of bounding box, i.e., (x, y).
+                 bbox is the bounding box coordinates as (x_top_left, y_top_left, x_bottom_right, y_bottom_right).
+                 info_dict is the dictionary of information which may be useful from the tracker (example:
+                 number of times tracker was lost while tracking.).
+
+        """
+
+        assert len(bboxes) == len(class_ids), "Containers must be of same length. len(bboxes)={}," \
+                                              " len(class_ids)={}".format(len(bboxes), len(class_ids))
+
+        assert len(bboxes) == len(detection_scores), "Containers must be of same length. len(bboxes)={}," \
+                                                     " len(class_ids)={}".format(len(bboxes), len(detection_scores))
+
+        self.frame_count += 1
+
+        new_bboxes = np.array(bboxes, dtype='int')
+        new_class_ids = np.array(class_ids, dtype='int')
+        new_detection_scores = np.array(detection_scores)
+
+        new_centroids = get_centroids(new_bboxes)
+
+        new_detections = list(zip(
+            range(len(bboxes)), new_bboxes, new_class_ids, new_centroids, new_detection_scores
+        ))
+
+        track_ids = list(self.tracks.keys())
+
+        updated_tracks = []
+        for track_id in track_ids:
+            if len(new_detections) > 0:
+                idx, bb, cid, ctrd, scr = max(new_detections, key=lambda x: iou(self.tracks[track_id].bbox, x[1]))
+
+                if iou(self.tracks[track_id].bbox, bb) > self.iou_threshold and self.tracks[track_id].class_id == cid:
+                    max_score = max(self.tracks[track_id].info['max_score'], scr)
+                    self._update_track(track_id, ctrd, bb, score=scr, max_score=max_score)
+
+                    updated_tracks.append(track_id)
+
+                    del new_detections[idx]
+
+            if len(updated_tracks) == 0 or track_id is not updated_tracks[-1]:
+                self.tracks[track_id].lost += 1
+
+                if self.tracks[track_id].lost > self.max_lost and \
+                        self.tracks[track_id].info['max_score'] >= self.max_detection_confidence:
+                    self._remove_track(track_id)
+
+        for idx, bb, cid, ctrd, scr in new_detections:
+            self._add_track(ctrd, bb, cid, score=scr, max_score=scr)
+
+        outputs = self._get_tracks(self.tracks)
+        return outputs

motrackers/simple_tracker2.py

@@ -1,7 +1,7 @@
 from collections import OrderedDict
 import numpy as np
 from scipy.spatial import distance
-from motrackers.utils.misc import get_centroid
+from motrackers.utils.misc import get_centroids
 from motrackers.track import Track
 
 
@@ -24,7 +24,7 @@ def __init__(self, max_lost=5):
         self.max_lost = max_lost
         self.frame_count = 0
 
-    def _add_object(self, centroid, bbox, class_id):
+    def _add_track(self, centroid, bbox, class_id, **kwargs):
         """
         Add a newly detected object to the queue
 
@@ -46,6 +46,9 @@ class label
                                                 centroid=centroid,
                                                 bbox=bbox,
                                                 class_id=class_id)
+        for key, value in kwargs.items():
+            self.tracks[self.next_track_id].info[key] = value
+
         self.next_track_id += 1
 
     def _remove_track(self, track_id):
@@ -63,7 +66,14 @@ def _remove_track(self, track_id):
         """
         del self.tracks[track_id]
 
-    def get_tracks(self, tracks):
+    def _update_track(self, track_id, centroid, bbox, **kwargs):
+        self.tracks[track_id].centroid = centroid
+        self.tracks[track_id].bbox = bbox
+        self.tracks[track_id].lost = 0
+        for key, value in kwargs.items():
+            self.tracks[track_id].info[key] = value
+
+    def _get_tracks(self, tracks):
         """
         Output the information of tracks
 
@@ -95,7 +105,7 @@ def get_tracks(self, tracks):
 
         return outputs
 
-    def update(self, bboxes: list, class_ids: list):
+    def update(self, bboxes: list, class_ids: list, detection_scores: list):
         """
         Update the tracker based on the new bboxes as input.
 
@@ -106,6 +116,8 @@ def update(self, bboxes: list, class_ids: list):
                  coordinates of bounding box as tuple (top-left-x, top-left-y, bottom-right-x, bottom-right-y).
         class_ids : list
                     List of class_ids (int) corresponding to labels of the detected object. Default is `None`.
+        detection_scores: list
+                         List of detection scores / probability of each detected object or objectness.
 
         Returns
         -------
@@ -122,39 +134,45 @@ def update(self, bboxes: list, class_ids: list):
         """
         self.frame_count += 1
 
+        new_bboxes = np.array(bboxes, dtype='int')
+        new_class_ids = np.array(class_ids, dtype='int')
+        new_detection_scores = np.array(detection_scores)
+
+        new_centroids = get_centroids(new_bboxes)
+
+        new_detections = list(zip(
+            range(len(bboxes)), new_bboxes, new_class_ids, new_centroids, new_detection_scores
+        ))
+
         if len(bboxes) == 0:        # if no object detected
             lost_ids = list(self.tracks.keys())
             for track_id in lost_ids:
                 self.tracks[track_id].lost += 1
                 if self.tracks[track_id].lost > self.max_lost:
                     self._remove_track(track_id)
 
-            outputs = self.get_tracks(self.tracks)
+            outputs = self._get_tracks(self.tracks)
             return outputs
 
-        new_class_ids = np.array(class_ids, dtype='int')
-        new_centroids = np.zeros((len(bboxes), 2), dtype="int")
-        for (i, bbox) in enumerate(bboxes):
-            new_centroids[i] = get_centroid(bbox)
-
-        if len(self.tracks):
-            track_ids = list(self.tracks.keys())
+        track_ids = list(self.tracks.keys())
+        if len(track_ids):
             old_centroids = np.array([self.tracks[tid].centroid for tid in track_ids])
-            D = distance.cdist(old_centroids, new_centroids)   # (row, col) = distance between old (row) and new (col)
-            row_idxs = D.min(axis=1).argsort()                 # old tracks sorted as per min distance from new
-            col_idxs = D.argmin(axis=1)[row_idxs]              # new tracks sorted as per min distance from old
+            D = distance.cdist(old_centroids, new_centroids)  # (row, col) = distance between old (row) and new (col)
+
+            row_idxs = D.min(axis=1).argsort()          # old tracks sorted as per min distance from new
+            col_idxs = D.argmin(axis=1)[row_idxs]       # new tracks sorted as per min distance from old
 
             assigned_rows, assigned_cols = set(), set()
             for (row_idx, col_idx) in zip(row_idxs, col_idxs):
                 if row_idx in assigned_rows or col_idx in assigned_cols:
                     continue
 
                 track_id = track_ids[row_idx]
+                
+                col_idx, bbox, class_id, centroid, detection_score = new_detections[col_idx]
 
-                if self.tracks[track_id].class_id == new_class_ids[col_idx]:
-                    self.tracks[track_id].centroid = new_centroids[col_idx]
-                    self.tracks[track_id].bbox = bboxes[col_idx]
-                    self.tracks[track_id].lost = 0
+                if self.tracks[track_id].class_id == class_id:
+                    self._update_track(track_id, centroid, bbox, score=detection_score)
                     assigned_rows.add(row_idx)
                     assigned_cols.add(col_idx)
 
@@ -170,10 +188,10 @@ def update(self, bboxes: list, class_ids: list):
                         self._remove_track(track_id)
             else:
                 for col_idx in unassigned_cols:
-                    self._add_object(new_centroids[col_idx], bboxes[col_idx], class_ids[col_idx])
+                    self._add_track(new_centroids[col_idx], bboxes[col_idx], class_ids[col_idx])
         else:
             for i in range(0, len(bboxes)):
-                self._add_object(new_centroids[i], bboxes[i], class_ids[i])
+                self._add_track(new_centroids[i], bboxes[i], class_ids[i])
 
-        outputs = self.get_tracks(self.tracks)
+        outputs = self._get_tracks(self.tracks)
         return outputs

motrackers/track.py

@@ -8,10 +8,14 @@ def __init__(self, track_id, centroid, bbox=None, class_id=None):
 
         Parameters
         ----------
-        track_id :
-        centroid :
-        bbox :
-        class_id :
+        track_id : int
+                   Track id.
+        centroid : tuple
+                   Centroid of the track pixel coordinate (x, y).
+        bbox : tuple, list, numpy.ndarray
+               Bounding box of the track.
+        class_id : int
+                   Class label id.
         """
         self.id = track_id
         self.class_id = class_id
@@ -22,4 +26,8 @@ def __init__(self, track_id, centroid, bbox=None, class_id=None):
         self.bbox = bbox
         self.lost = 0
 
-        self.info = {}
+        self.info = dict(
+            max_score=0.0,
+            lost=0,
+            score=0.0,
+        )

motrackers/utils/__init__.py

@@ -2,3 +2,5 @@
 from .filechooser_utils import select_videofile
 from .filechooser_utils import select_yolo_model
 from .filechooser_utils import select_tfmobilenet
+from .misc import iou
+from .misc import get_centroid, get_centroids

motrackers/utils/misc.py

@@ -1,15 +1,83 @@
+import numpy as np
 
 
 def get_centroid(bounding_box):
     """
     Calculate the centroid of bounding box.
 
-    :param bounding_box: list of bounding box coordinates of top-left and bottom-right (xlt, ylt, xrb, yrb)
-    :return: bounding box centroid coordinates (x, y)
-    """
+    Parameters
+    ----------
+    bounding_box : list
+                   list of bounding box coordinates of top-left and bottom-right (xlt, ylt, xrb, yrb)
+
+    Returns
+    -------
+    centroid: tuple
+              Bounding box centroid pixel coordinates (x, y).
 
+    """
     xlt, ylt, xrb, yrb = bounding_box
     centroid_x = int((xlt + xrb) / 2.0)
     centroid_y = int((ylt + yrb) / 2.0)
 
     return centroid_x, centroid_y
+
+
+def get_centroids(bboxes):
+    assert bboxes.shape[1] == 4
+
+    x = np.mean(bb[:, [0, 2]], axis=1, keepdims=True, dtype='int')
+    y = np.mean(bb[:, [1, 3]], axis=1, keepdims=True, dtype='int')
+    centroids = np.concatenate([x, y], axis=1)
+    return centroids
+
+
+def iou(bbox1, bbox2):
+    """
+    Calculates the intersection-over-union of two bounding boxes.
+    Source: https://github.com/bochinski/iou-tracker/blob/master/util.py
+
+    Parameters
+    ----------
+    bbox1 : numpy.array, list of floats
+            bounding box in format (x-top-left, y-top-left, x-bottom-right, y-bottom-right) of length 4.
+    bbox2 : numpy.array, list of floats
+            bounding box in format (x-top-left, y-top-left, x-bottom-right, y-bottom-right) of length 4.
+
+    Returns
+    -------
+    iou: float
+         intersection-over-onion of bbox1, bbox2.
+    """
+
+    bbox1 = [float(x) for x in bbox1]
+    bbox2 = [float(x) for x in bbox2]
+
+    (x0_1, y0_1, x1_1, y1_1), (x0_2, y0_2, x1_2, y1_2) = bbox1, bbox2
+
+    # get the overlap rectangle
+    overlap_x0 = max(x0_1, x0_2)
+    overlap_y0 = max(y0_1, y0_2)
+    overlap_x1 = min(x1_1, x1_2)
+    overlap_y1 = min(y1_1, y1_2)
+
+    # check if there is an overlap
+    if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:
+        return 0.0
+
+    # if yes, calculate the ratio of the overlap to each ROI size and the unified size
+    size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
+    size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
+    size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
+    size_union = size_1 + size_2 - size_intersection
+
+    iou_ = size_intersection / size_union
+
+    return iou_
+
+
+if __name__ == '__main__':
+    bb = np.random.random_integers(0, 100, size=(20,)).reshape((5, 4))
+    c = get_centroids(bb)
+    print(bb)
+    print(c)