Merge pull request #860 from RandySheriffH/rashuai/NMSV4V5

Convert NMS V4&V5

Author: RandySheriffH

tests/test_backend.py

@@ -2580,10 +2580,52 @@ def test_non_max_suppression(self):
         box_num = 10
         boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
         scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
         def func(boxes, scores):
             res1 = tf.image.non_max_suppression(boxes, scores, max_output_size=int(box_num / 2))
             res2 = tf.image.non_max_suppression(boxes, scores, max_output_size=0)
             return tf.identity(res1, name=_TFOUTPUT), tf.identity(res2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_opset_min_version(11, "NonMaxSuppressionV4")
+    def test_non_max_suppression_v4(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_padded(boxes, scores, max_output_size=int(box_num * 2),
+                                                             pad_to_max_output_size=True)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_opset_min_version(11, "NonMaxSuppressionV4")
+    def test_non_max_suppression_v4_no_padding(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_padded(boxes, scores, max_output_size=int(box_num),
+                                                             pad_to_max_output_size=False)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_tf_min_version("1.15")
+    @check_opset_min_version(11, "NonMaxSuppressionV5")
+    def test_non_max_suppression_v5(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_with_scores(boxes, scores, max_output_size=int(box_num / 2),
+                                                                  soft_nms_sigma=0.0)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
 
     def _conv1d_test(self, x_val, w, stride=None, padding="VALID", rtol=1e-07):

tf2onnx/onnx_opset/tensor.py

@@ -1452,7 +1452,8 @@ def version_10(cls, ctx, node, **kwargs):
             raise ValueError("dtype " + str(node_dtype) + " is not supported in onnx for now")
 
 
-@tf_op(["NonMaxSuppressionV2", "NonMaxSuppressionV3"], onnx_op="NonMaxSuppression")
+@tf_op(["NonMaxSuppressionV2", "NonMaxSuppressionV3", "NonMaxSuppressionV4", "NonMaxSuppressionV5"],
+       onnx_op="NonMaxSuppression")
 class NonMaxSuppression:
     @classmethod
     def version_10(cls, ctx, node, **kwargs):
@@ -1464,18 +1465,41 @@ def version_10(cls, ctx, node, **kwargs):
         # onnx output is [num_selected_boxes, 3], the meaning of last dim is [batch_index, class_index, box_index]
         # while tf's output is [num_selected_boxes]
         ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[0], axes=[0])
-        ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], axes=[0, 1])
+        input_score = ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], axes=[0, 1])
         ctx.insert_new_node_on_input(node, "Cast", node.input[2], to=onnx_pb.TensorProto.INT64)
         # replace original node with nonmaxsurppress + slice + squeeze + cast
-        dtypes = [ctx.get_dtype(node.output[0])]
-        shapes = [ctx.get_shape(node.output[0])]
+        dtypes = [[ctx.get_dtype(output)] for output in node.output]
+        shapes = [[ctx.get_shape(output)] for output in node.output]
+        max_output_size = node.input[2]
+        utils.make_sure(len(node.inputs) <= 5 or int(node.inputs[5].get_tensor_value(False)) == 0,
+                        "soft_nms_sigma must be 0")
         ctx.remove_node(node.name)
-        new_nonmaxsurppress = ctx.make_node(node.type, node.input).output[0]
+        new_nonmaxsurppress = ctx.make_node(node.type, node.input[: 5]).output[0]
         slice_op = GraphBuilder(ctx).make_slice({"data": new_nonmaxsurppress,
                                                  "axes": [1], "ends": [3], "starts": [2]})
         squeeze_op = ctx.make_node("Squeeze", [slice_op], attr={"axes": [1]})
-        ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
-                      name=node.name, outputs=node.output, dtypes=dtypes, shapes=shapes)
+        if len(node.input) > 5:  # v5, called by ..._with_scores(), pad_to_max_output_size always False
+            ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[0]], dtypes=dtypes[0], shapes=shapes[0])
+            ctx.make_node("Gather", inputs=[input_score.input[0], squeeze_op.output[0]],
+                          outputs=[node.output[1]], dtypes=dtypes[1], shapes=shapes[1])
+        elif "pad_to_max_output_size" in node.attr:  # V4
+            shape_op = ctx.make_node("Shape", inputs=[squeeze_op.output[0]])
+            const_zero = ctx.make_const(utils.make_name("const_zero"), np.array([0], dtype=np.int64))
+            sub_op = ctx.make_node("Sub", inputs=[max_output_size, shape_op.output[0]])
+            raw_pad = ctx.make_node("Concat", inputs=[const_zero.output[0], sub_op.output[0]], attr={'axis': 0})
+            raw_pad_float = ctx.make_node("Cast", inputs=[raw_pad.output[0]], attr={"to": onnx_pb.TensorProto.FLOAT})
+            relu_op = ctx.make_node("Relu", inputs=[raw_pad_float.output[0]])
+            pad_val = ctx.make_node("Cast", inputs=[relu_op.output[0]], attr={"to": onnx_pb.TensorProto.INT64})
+            pad_op = ctx.make_node("Pad", inputs=[squeeze_op.output[0], pad_val.output[0]])
+            ctx.make_node("Cast", inputs=pad_op.output, name="cast_A", attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[0]], dtypes=dtypes[0], shapes=shapes[0])
+            reduce_op = ctx.make_node("ReduceSum", inputs=shape_op.output, attr={"axes": [0], "keepdims": 0})
+            ctx.make_node("Cast", inputs=[reduce_op.output[0]], name="cast_B", attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[1]], dtypes=dtypes[1], shapes=shapes[1])
+        else:
+            ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
+                          name=node.name, outputs=node.output, dtypes=dtypes[0], shapes=shapes[0])
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):