Upgrade conversion of Core ML Crop (#249)

Author: wschin

onnxmltools/convert/common/_apply_operation.py

@@ -7,6 +7,7 @@
 # best to use those functions because they can produce ONNX operators according to the ONNX version specified in the
 # `container` argument. Notice that those function behaviors are defined in a way very similar to ONNX-1.2.
 
+import numpy as np
 from ...proto import onnx_proto
 
 def _create_name_or_use_existing_one(scope, op_type, name):
@@ -438,3 +439,46 @@ def apply_softmax(scope, input_name, output_name, container, operator_name=None,
 def apply_normalization(scope, input_name, output_name, container, operator_name=None, axis=1, p=2):
     name = _create_name_or_use_existing_one(scope, 'LpNormalization', operator_name)
     container.add_node('LpNormalization', input_name, output_name, name=name, p=p, axis=axis)
+
+def apply_crop_height_width(scope, input_name, output_name, container, operator_name=None,
+        top_border=0, bottom_border=0, left_border=0, right_border=0):
+    name = scope.get_unique_operator_name('CropHeightWidth')
+    if container.target_opset < 9:
+        # If operator set < 9, we can use the experimental Crop in ONNX.
+        attrs = {'name': name, 'border': [left_border, top_border, right_border, bottom_border]}
+        container.add_node('Crop', input_name, output_name, **attrs)
+    else:
+        # The experimental Crop in ONNX is removed after operator set 9, so we
+        # switch to ONNX DynamicSlice operator.
+
+        # CoreML only crops H- and W-axes.
+        axes = [2, 3]
+        axes_name = scope.get_unique_variable_name(name + '_axes')
+        container.add_initializer(axes_name, onnx_proto.TensorProto.INT64,
+                                  [len(axes)], axes)
+
+        # Number of cropped pixels is the starting index of the remained region.
+        starts = [top_border, left_border]
+        starts_name = scope.get_unique_variable_name(name + '_starts')
+        container.add_initializer(starts_name, onnx_proto.TensorProto.INT64,
+                                  [len(starts)], starts)
+
+        # First we assume no cropping is needed at the end of those axes.
+        # We will change this right below depending on Crop's configuration.
+        ends = [np.iinfo(np.int64).max] * 2
+
+        # Crop n pixel means the end index (exclusive) is -n. Note that indexing
+        # system is zero-based.
+        if bottom_border > 0:
+            ends[0] = -bottom_border
+        if right_border > 0:
+            ends[1] = -right_border
+
+        # Add the adjusted ends.
+        ends_name = scope.get_unique_variable_name(name + '_ends')
+        container.add_initializer(ends_name, onnx_proto.TensorProto.INT64,
+                                  [len(ends)], ends)
+
+        # Collect all input names as a list because DynamicSlice has multiple inputs.
+        input_list = [input_name, starts_name, ends_name, axes_name]
+        container.add_node('DynamicSlice', input_list, output_name, op_version=9)

onnxmltools/convert/coreml/operator_converters/neural_network/Crop.py

@@ -4,18 +4,19 @@
 # license information.
 # --------------------------------------------------------------------------
 
+import numpy as np
+from .....proto import onnx_proto
+from ....common._apply_operation import apply_crop_height_width
 from ....common._registration import register_converter
 
-
 def convert_crop(scope, operator, container):
-    if len(operator.input_full_names) > 2:
-        raise RuntimeError('Unlike CoreML, ONNX only supports cropping with a single input')
-
-    op_type = 'Crop'
-    attrs = {'name': operator.full_name}
+    # Extract number of pixels cropped in CoreML operator.
     border = operator.raw_operator.crop.cropAmounts.borderAmounts
-    left_border, top_border, right_border, bottom_border = (None,) * 4
-    if len(border):
+
+    # Compute cropping amounts. left_border=1 means one pixel will be removed at the beginning
+    # of W-axis. bottom_border=1 means one pixel will be removed at the end of H-axis.
+    left_border, top_border, right_border, bottom_border = (0,) * 4
+    if len(operator.input_full_names) == 1:
         left_border = border[1].startEdgeSize
         top_border = border[0].startEdgeSize
         right_border = border[1].endEdgeSize
@@ -29,9 +30,13 @@ def convert_crop(scope, operator, container):
         right_border = in_shape[3] - left_border - out_shape[3]
         bottom_border = in_shape[2] - top_border - out_shape[2]
 
-    attrs['border'] = [left_border, top_border, right_border, bottom_border]
+    # Delegate the selection of ONNX operator to a version-dependent function.
+    apply_crop_height_width(scope,
+        operator.input_full_names[0], operator.output_full_names[0],
+        container, operator_name=operator.full_name,
+        top_border=top_border, bottom_border=bottom_border,
+        left_border=left_border, right_border=right_border)
 
-    container.add_node(op_type, operator.input_full_names[0:1], operator.output_full_names, **attrs)
 
 
 register_converter('crop', convert_crop)