Implement non-standard nms for tflite (#1318)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_tflite_postprocess.py

@@ -45,9 +45,17 @@ def test_postprocess_model3(self):
     def test_postprocess_model4(self):
         self._test_postprocess(num_classes=5, num_boxes=99, detections_per_class=2, max_detections=20, extra_class=True)
 
-    def _test_postprocess(self, num_classes, num_boxes, detections_per_class, max_detections, extra_class=False):
+    @requires_tflite("TFLite_Detection_PostProcess")
+    @check_opset_min_version(11, "Pad")
+    def test_postprocess_model5(self):
+        self._test_postprocess(num_classes=1, num_boxes=100, detections_per_class=0,
+                               max_detections=50, use_regular_nms=False)
+
+    def _test_postprocess(self, num_classes, num_boxes, detections_per_class,
+                          max_detections, extra_class=False, use_regular_nms=True):
         model = self.make_postprocess_model(num_classes=num_classes, detections_per_class=detections_per_class,
-                                            max_detections=max_detections, x_scale=11.0, w_scale=6.0)
+                                            max_detections=max_detections, x_scale=11.0, w_scale=6.0,
+                                            use_regular_nms=use_regular_nms)
 
         np.random.seed(42)
         box_encodings_val = np.random.random_sample([1, num_boxes, 4]).astype(np.float32)

tf2onnx/tflite_handlers/tfl_postprocess.py

@@ -69,7 +69,12 @@ def version_11(cls, ctx, node, **kwargs):
         score_threshold = np.array(node.get_attr_value('nms_score_threshold'), np.float32)
         score_threshold_const = ctx.make_const(utils.make_name('score_threshold'), score_threshold).output[0]
 
-        boxes_per_class = np.array(node.get_attr_value('detections_per_class', 100), np.int64)
+        if node.get_attr_value('use_regular_nms', False):
+            boxes_per_class = np.array(node.get_attr_value('detections_per_class', 100), np.int64)
+        else:
+            # When tflite uses FastNMS, detections_per_class is ignored.
+            logging.warning("NMS node %s uses fast NMS. ONNX will approximate with standard NMS.", node.name)
+            boxes_per_class = np.array(max_detections, np.int64)
         max_boxes_per_class_const = ctx.make_const(utils.make_name('max_boxes_per_class'), boxes_per_class).output[0]
 
         # scores.shape = [batch_dim, classes_num, box_num]