Fix bug in CropAndResize for empty tensors (#1588)

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -3791,6 +3791,18 @@ def func(input_x, boxes, box_ind, corp_size):
                             {_INPUT: input_x_val, _INPUT1: boxes_val, _INPUT2: box_ind_val, _INPUT3: corp_size_val},
                             rtol=1e-04, atol=1e-03)
 
+    @check_opset_min_version(11, "CropAndResize")
+    def test_crop_and_resize_empty_tensor(self):
+        def func(input_x, boxes, box_ind, corp_size):
+            return tf.image.crop_and_resize(input_x, boxes, box_ind, corp_size, name=_TFOUTPUT, extrapolation_value=1.0)
+        input_x_val = np.random.randint(low=0, high=256, size=[0, 36, 36, 3]).astype(np.float32)  # NHWC
+        boxes_val = np.array([]).astype(np.float32).reshape([0, 4])
+        box_ind_val = np.array([]).astype(np.int32)
+        corp_size_val = np.array([40, 40]).astype(np.int32)
+        self._run_test_case(func, [_OUTPUT],
+                            {_INPUT: input_x_val, _INPUT1: boxes_val, _INPUT2: box_ind_val, _INPUT3: corp_size_val},
+                            rtol=1e-04, atol=1e-03)
+
     def test_batch_to_space3d(self):
         block_size = [2, 2]
         crop = [[0, 1], [2, 1]]

tf2onnx/onnx_opset/nn.py

@@ -318,9 +318,10 @@ def build_dynamic_target_size(ctx, transposed_intput, target_hw):
     shape_of_transposed_input = ctx.make_node("Shape", [transposed_intput])
     first_half_of_shape = GraphBuilder(ctx).make_slice(
         {"data": shape_of_transposed_input.output[0], "ends": [2], "starts": [0]})
-    target_size_int64 = ctx.make_node("Cast", [target_hw], attr={'to': TensorProto.INT64})
+    if ctx.get_dtype(target_hw) != TensorProto.INT64:
+        target_hw = ctx.make_node("Cast", [target_hw], attr={'to': TensorProto.INT64}).output[0]
     # We build a tensor containing [n c nh nw]
-    final_target_size = ctx.make_node("Concat", [first_half_of_shape, target_size_int64.output[0]], {'axis': 0})
+    final_target_size = ctx.make_node("Concat", [first_half_of_shape, target_hw], {'axis': 0})
     return final_target_size
 
 
@@ -1183,9 +1184,13 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
             "method").s == b"nearest" else "linear"
         extrapolation_value = float(node.get_attr("extrapolation_value", "0").f)
         input_x = node.input[0]
+        x_shape = ctx.make_node("Shape", [input_x]).output[0]
+        num_channels = GraphBuilder(ctx).make_slice({"data": x_shape, "starts": [3], "ends": [4], "axes": [0]})
         boxes = node.input[1]
         box_ind = node.input[2]
         crop_size = node.input[3]
+        if ctx.get_dtype(crop_size) != TensorProto.INT64:
+            crop_size = ctx.make_node("Cast", [crop_size], attr={'to': TensorProto.INT64}).output[0]
         trip_name = utils.make_name(node.name + "_i")
         cond_name = utils.make_name(node.name + "_cond")
         cond_out_name = utils.make_name(node.name + "cond_out")
@@ -1233,6 +1238,10 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
         branches = {"body": g}
         inner_loop = ctx.make_node("Loop", [trip_node.output[0], cond_const.output[0]], name=node.name,
                                    outputs=node.output, branches=branches)
+        const_neg_one = ctx.make_const(utils.make_name("const_neg_one"), np.array([-1], np.int64)).output[0]
+        final_shape = ctx.make_node("Concat", [const_neg_one, crop_size, num_channels], attr={'axis': 0}).output[0]
+        # This reshape fixes the case when there are no iterations and the scan output is empty.
+        ctx.insert_new_node_on_output("Reshape", inner_loop.output[0], inputs=[inner_loop.output[0], final_shape])
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):