Perform per class NMS.

keras_retinanet/backend/tensorflow_backend.py

@@ -18,10 +18,6 @@ import tensorflow
 import keras
 
 
-def top_k(*args, **kwargs):
-    return tensorflow.nn.top_k(*args, **kwargs)
-
-
 def resize_images(*args, **kwargs):
     return tensorflow.image.resize_images(*args, **kwargs)
 
@@ -34,6 +30,10 @@ def range(*args, **kwargs):
     return tensorflow.range(*args, **kwargs)
 
 
+def scatter_nd(*args, **kwargs):
+    return tensorflow.scatter_nd(*args, **kwargs)
+
+
 def gather_nd(*args, **kwargs):
     return tensorflow.gather_nd(*args, **kwargs)
 

keras_retinanet/layers/_misc.py

@@ -76,43 +76,60 @@ class Anchors(keras.layers.Layer):
 
 
 class NonMaximumSuppression(keras.layers.Layer):
-    def __init__(self, nms_threshold=0.5, top_k=None, max_boxes=300, *args, **kwargs):
-        self.nms_threshold = nms_threshold
-        self.top_k         = top_k
-        self.max_boxes     = max_boxes
+    def __init__(self, nms_threshold=0.5, score_threshold=0.05, max_boxes=300, *args, **kwargs):
+        self.nms_threshold   = nms_threshold
+        self.score_threshold = score_threshold
+        self.max_boxes       = max_boxes
         super(NonMaximumSuppression, self).__init__(*args, **kwargs)
 
     def call(self, inputs, **kwargs):
-        boxes, classification, detections = inputs
-
         # TODO: support batch size > 1.
-        boxes          = boxes[0]
-        classification = classification[0]
-        detections     = detections[0]
+        boxes           = inputs[0][0]
+        classification  = inputs[1][0]
+        other           = [i[0] for i in inputs[2:]]  # can be any user-specified additional data
+        indices         = backend.range(keras.backend.shape(classification)[0])
+        selected_scores = []
+
+        # perform per class NMS
+        for c in range(int(classification.shape[1])):
+            scores = classification[:, c]
+
+            # threshold based on score
+            score_indices = backend.where(keras.backend.greater(scores, self.score_threshold))
+            score_indices = keras.backend.cast(score_indices, 'int32')
+            boxes_        = backend.gather_nd(boxes, score_indices)
+            scores        = keras.backend.gather(scores, score_indices)[:, 0]
+
+            # perform NMS
+            nms_indices = backend.non_max_suppression(boxes_, scores, max_output_size=self.max_boxes, iou_threshold=self.nms_threshold)
+
+            # filter set of original indices
+            selected_indices = keras.backend.gather(score_indices, nms_indices)
+
+            # mask original classification column, setting all suppressed values to 0
+            scores = keras.backend.gather(scores, nms_indices)
+            scores = backend.scatter_nd(selected_indices, scores, keras.backend.shape(classification[:, c]))
+            scores = keras.backend.expand_dims(scores, axis=1)
 
-        scores = keras.backend.max(classification, axis=1)
+            selected_scores.append(scores)
 
-        # selecting best anchors theoretically improves speed at the cost of minor performance
-        if self.top_k:
-            scores, indices = backend.top_k(scores, self.top_k, sorted=False)
-            boxes           = keras.backend.gather(boxes, indices)
-            classification  = keras.backend.gather(classification, indices)
-            detections      = keras.backend.gather(detections, indices)
+        # reconstruct the (suppressed) classification scores
+        classification = keras.backend.concatenate(selected_scores, axis=1)
 
-        indices = backend.non_max_suppression(boxes, scores, max_output_size=self.max_boxes, iou_threshold=self.nms_threshold)
+        # reconstruct into the expected output
+        detections = keras.backend.concatenate([boxes, classification] + other, axis=1)
 
-        detections = keras.backend.gather(detections, indices)
         return keras.backend.expand_dims(detections, axis=0)
 
     def compute_output_shape(self, input_shape):
-        return (input_shape[2][0], None, input_shape[2][2])
+        return (input_shape[0][0], input_shape[0][1], sum([i[2] for i in input_shape]))
 
     def get_config(self):
         config = super(NonMaximumSuppression, self).get_config()
         config.update({
-            'nms_threshold' : self.nms_threshold,
-            'top_k'         : self.top_k,
-            'max_boxes'     : self.max_boxes,
+            'nms_threshold'   : self.nms_threshold,
+            'score_threshold' : self.score_threshold,
+            'max_boxes'       : self.max_boxes,
         })
 
         return config

keras_retinanet/models/retinanet.py

@@ -208,12 +208,13 @@ def retinanet_bbox(inputs, num_classes, nms=True, name='retinanet-bbox', *args,
     classification = model.outputs[2]
 
     # apply predicted regression to anchors
-    boxes      = layers.RegressBoxes(name='boxes')([anchors, regression])
-    detections = keras.layers.Concatenate(axis=2)([boxes, classification] + model.outputs[3:])
+    boxes = layers.RegressBoxes(name='boxes')([anchors, regression])
 
     # additionally apply non maximum suppression
     if nms:
-        detections = layers.NonMaximumSuppression(name='nms')([boxes, classification, detections])
+        detections = layers.NonMaximumSuppression(name='nms')([boxes, classification] + model.outputs[3:])
+    else:
+        detections = keras.layers.Concatenate(axis=2)([boxes, classification] + model.outputs[3:])
 
     # construct the model
     return keras.models.Model(inputs=inputs, outputs=model.outputs[1:] + [detections], name=name)