instance mask cropping + bug fix on original ltwh update

Author: Ling Evan

deep_sort_realtime/deep_sort/detection.py

@@ -16,6 +16,8 @@ class Detection(object):
         A feature vector that describes the object contained in this image.
     class_name : Optional str
         Detector predicted class name.
+    instance_mask : Optional 
+        Instance mask corresponding to bounding box
     others : Optional any
         Other supplementary fields associated with detection that wants to be stored as a "memory" to be retrieve through the track downstream.
 
@@ -30,12 +32,13 @@ class Detection(object):
 
     """
 
-    def __init__(self, ltwh, confidence, feature, class_name=None, others=None):
+    def __init__(self, ltwh, confidence, feature, class_name=None, instance_mask=None, others=None):
         # def __init__(self, ltwh, feature):
         self.ltwh = np.asarray(ltwh, dtype=np.float)
         self.confidence = float(confidence)
         self.feature = np.asarray(feature, dtype=np.float32)
         self.class_name = class_name
+        self.instance_mask = instance_mask
         self.others = others
 
     def get_ltwh(self):

deep_sort_realtime/deep_sort/track.py

@@ -44,6 +44,8 @@ class Track:
         Classname of matched detection
     det_conf : Optional float
         Confidence associated with matched detection
+    instance_mask : Optional 
+        Instance mask associated with matched detection
     others : Optional any
         Any supplementary fields related to matched detection
 
@@ -80,6 +82,7 @@ def __init__(
         original_ltwh=None,
         det_class=None,
         det_conf=None,
+        instance_mask=None,
         others=None,
     ):
         self.mean = mean
@@ -100,6 +103,7 @@ def __init__(
         self.original_ltwh = original_ltwh
         self.det_class = det_class
         self.det_conf = det_conf
+        self.instance_mask = instance_mask
         self.others = others
 
     def to_tlwh(self, orig=False, orig_strict=False):
@@ -189,12 +193,24 @@ def get_det_class(self):
         """
         return self.det_class
 
+    def get_instance_mask(self):
+        '''
+        Get instance mask associated with detection. Will be None is there are no associated detection this round
+        '''
+        return self.instance_mask
+    
     def get_det_supplementary(self):
         """
         Get supplementary info associated with the detection. Will be None is there are no associated detection this round.
         """
         return self.others
 
+    def get_feature(self):
+        '''
+        Get latest appearance feature
+        '''
+        return self.features[-1]
+
     def predict(self, kf):
         """Propagate the state distribution to the current time step using a
         Kalman filter prediction step.
@@ -210,6 +226,7 @@ def predict(self, kf):
         self.time_since_update += 1
         self.original_ltwh = None
         self.det_conf = None
+        self.instance_mask = None
         self.others = None
 
     def update(self, kf, detection):
@@ -231,6 +248,7 @@ def update(self, kf, detection):
         self.features.append(detection.feature)
         self.det_conf = detection.confidence
         self.det_class = detection.class_name
+        self.instance_mask = detection.instance_mask
         self.others = detection.others
 
         self.hits += 1

deep_sort_realtime/deep_sort/tracker.py

@@ -194,8 +194,10 @@ def _initiate_track(self, detection):
                 self.max_age,
                 # mean, covariance, self._next_id, self.n_init, self.max_age,
                 feature=detection.feature,
+                original_ltwh=detection.get_ltwh(),
                 det_class=detection.class_name,
                 det_conf=detection.confidence,
+                instance_mask=detection.instance_mask,
                 others=detection.others,
             )
         )

deep_sort_realtime/deepsort_tracker.py

@@ -144,7 +144,7 @@ def __init__(
         logger.info(f'- in-build embedder : {"No" if self.embedder is None else "Yes"}')
         logger.info(f'- polygon detections : {"No" if polygon is False else "Yes"}')
 
-    def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, others=None):
+    def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, others=None, instance_masks=None):
 
         """Run multi-target tracker on a particular sequence.
 
@@ -162,6 +162,8 @@ def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, oth
             Provide today's date, for naming of tracks
         others: Optional[ List ] = None
             Other things associated to detections to be stored in tracks, usually, could be corresponding segmentation mask, other associated values, etc. Currently others is ignored with polygon is True.
+        instance_masks: Optional [ List ] = None
+            Instance masks corresponding to detections. If given, they are used to filter out background and only use foreground for apperance embedding. Expects numpy boolean mask matrix.
 
         Returns
         -------
@@ -184,10 +186,10 @@ def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, oth
             raw_detections = [d for d in raw_detections if d[0][2] > 0 and d[0][3] > 0]
 
             if embeds is None:
-                embeds = self.generate_embeds(frame, raw_detections)
+                embeds = self.generate_embeds(frame, raw_detections, instance_masks=instance_masks)
 
             # Proper deep sort detection objects that consist of bbox, confidence and embedding.
-            detections = self.create_detections(raw_detections, embeds, others=others)
+            detections = self.create_detections(raw_detections, embeds, instance_masks=instance_masks, others=others)
         else:
             polygons, bounding_rects = self.process_polygons(raw_detections[0])
 
@@ -218,15 +220,24 @@ def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, oth
     def refresh_track_ids(self):
         self.tracker._next_id
 
-    def generate_embeds(self, frame, raw_dets):
-        crops = self.crop_bb(frame, raw_dets)
-        return self.embedder.predict(crops)
+    def generate_embeds(self, frame, raw_dets, instance_masks=None):
+        crops, cropped_inst_masks = self.crop_bb(frame, raw_dets, instance_masks=instance_masks)
+        if cropped_inst_masks is not None:
+            masked_crops = []
+            for crop, mask in zip(crops, cropped_inst_masks):
+                masked_crop = np.zeros_like(crop)
+                masked_crop = masked_crop + np.array([123.675, 116.28, 103.53], dtype=crop.dtype)
+                masked_crop[mask] = crop[mask]
+                masked_crops.append(masked_crop)
+            return self.embedder.predict(masked_crops)
+        else:
+            return self.embedder.predict(crops)
 
     def generate_embeds_poly(self, frame, polygons, bounding_rects):
         crops = self.crop_poly_pad_black(frame, polygons, bounding_rects)
         return self.embedder.predict(crops)
 
-    def create_detections(self, raw_dets, embeds, others=None):
+    def create_detections(self, raw_dets, embeds, instance_masks=None, others=None):
         detection_list = []
         for i, (raw_det, embed) in enumerate(zip(raw_dets, embeds)):
             detection_list.append(
@@ -235,6 +246,7 @@ def create_detections(self, raw_dets, embeds, others=None):
                     raw_det[1], 
                     embed, 
                     class_name=raw_det[2] if len(raw_det)==3 else None,
+                    instance_mask = instance_masks[i] if isinstance(instance_masks, Iterable) else instance_masks,
                     others = others[i] if isinstance(others, Iterable) else others,
                 )
             )  # raw_det = [bbox, conf_score, class]
@@ -265,10 +277,14 @@ def process_polygons(raw_polygons):
         return polygons, bounding_rects
 
     @staticmethod
-    def crop_bb(frame, raw_dets):
+    def crop_bb(frame, raw_dets, instance_masks=None):
         crops = []
         im_height, im_width = frame.shape[:2]
-        for detection in raw_dets:
+        if instance_masks is not None: 
+            masks = []
+        else:
+            masks = None
+        for i, detection in enumerate(raw_dets):
             l, t, w, h = [int(x) for x in detection[0]]
             r = l + w
             b = t + h
@@ -277,7 +293,10 @@ def crop_bb(frame, raw_dets):
             crop_t = max(0, t)
             crop_b = min(im_height, b)
             crops.append(frame[crop_t:crop_b, crop_l:crop_r])
-        return crops
+            if instance_masks is not None: 
+                masks.append( instance_masks[i][crop_t:crop_b, crop_l:crop_r] )
+        
+        return crops, masks
 
     @staticmethod
     def crop_poly_pad_black(frame, polygons, bounding_rects):