Update deep sort to have class name.

Author: theAIGuysCode

Diff for: deep_sort/detection.py

@@ -21,17 +21,22 @@ class Detection(object):
         Bounding box in format `(top left x, top left y, width, height)`.
     confidence : ndarray
         Detector confidence score.
+    class_name : ndarray
+        Detector class.
     feature : ndarray | NoneType
         A feature vector that describes the object contained in this image.
 
     """
 
-    def __init__(self, tlwh, confidence, cls, feature):
+    def __init__(self, tlwh, confidence, class_name, feature):
         self.tlwh = np.asarray(tlwh, dtype=np.float)
         self.confidence = float(confidence)
-        self.cls = cls
+        self.class_name = class_name
         self.feature = np.asarray(feature, dtype=np.float32)
 
+    def get_class(self):
+        return self.class_name
+
     def to_tlbr(self):
         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
         `(top left, bottom right)`.

Diff for: deep_sort/track.py

@@ -64,7 +64,7 @@ class Track:
     """
 
     def __init__(self, mean, covariance, track_id, n_init, max_age,
-                 feature=None):
+                 feature=None, class_name=None):
         self.mean = mean
         self.covariance = covariance
         self.track_id = track_id
@@ -79,6 +79,7 @@ def __init__(self, mean, covariance, track_id, n_init, max_age,
 
         self._n_init = n_init
         self._max_age = max_age
+        self.class_name = class_name
 
     def to_tlwh(self):
         """Get current position in bounding box format `(top left x, top left y,
@@ -108,6 +109,9 @@ def to_tlbr(self):
         ret = self.to_tlwh()
         ret[2:] = ret[:2] + ret[2:]
         return ret
+    
+    def get_class(self):
+        return self.class_name
 
     def predict(self, kf):
         """Propagate the state distribution to the current time step using a

Diff for: deep_sort/tracker.py

@@ -132,7 +132,8 @@ def gated_metric(tracks, dets, track_indices, detection_indices):
 
     def _initiate_track(self, detection):
         mean, covariance = self.kf.initiate(detection.to_xyah())
+        class_name = detection.get_class()
         self.tracks.append(Track(
             mean, covariance, self._next_id, self.n_init, self.max_age,
-            detection.feature))
+            detection.feature, class_name))
         self._next_id += 1