[TOPI][RELAY][OP] add op crop_and_resize (apache#4417)

* [TOPI][RELAY][OP] add op crop_and_resize

* fix pylint

* incorporate comments

* fix ci

docs/api/python/topi.rst

@@ -104,6 +104,7 @@ List of operators
    topi.ndarray_size
    topi.layout_transform
    topi.image.resize
+   topi.image.crop_and_resize
    topi.argsort
    topi.topk
    topi.sequence_mask
@@ -207,6 +208,7 @@ topi.nn
 topi.image
 ~~~~~~~~~~
 .. autofunction:: topi.image.resize
+.. autofunction:: topi.image.crop_and_resize
 
 topi.sparse
 ~~~~~~~~~~~

docs/langref/relay_op.rst

@@ -169,6 +169,7 @@ This level enables additional math and transform operators.
    :nosignatures:
 
    tvm.relay.image.resize
+   tvm.relay.image.crop_and_resize
    tvm.relay.vision.multibox_prior
    tvm.relay.vision.multibox_transform_loc
    tvm.relay.vision.nms
@@ -335,6 +336,7 @@ Level 4 Definitions
 Level 5 Definitions
 -------------------
 .. autofunction:: tvm.relay.image.resize
+.. autofunction:: tvm.relay.image.crop_and_resize
 .. autofunction:: tvm.relay.vision.multibox_prior
 .. autofunction:: tvm.relay.vision.multibox_transform_loc
 .. autofunction:: tvm.relay.vision.nms

include/tvm/relay/attrs/image.h

@@ -63,6 +63,34 @@ struct ResizeAttrs : public tvm::AttrsNode<ResizeAttrs> {
   }
 };
 
+/*! \brief Attributes used in image crop_and_resize operator */
+struct CropAndResizeAttrs : public tvm::AttrsNode<CropAndResizeAttrs> {
+  Array<IndexExpr> crop_size;
+  std::string layout;
+  std::string method;
+  double extrapolation_value;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(CropAndResizeAttrs, "relay.attrs.CropAndResizeAttrs") {
+    TVM_ATTR_FIELD(crop_size).set_default(NullValue<Array<IndexExpr> >())
+        .describe("Target Size.");
+    TVM_ATTR_FIELD(layout).set_default("NCHW")
+        .describe("Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
+                  "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                  "dimensions respectively. Resize is applied on the 'H' and"
+                  "'W' dimensions.");
+    TVM_ATTR_FIELD(method).set_default("bilinear")
+        .describe("Specify the mode to use for scaling."
+                  "nearest_neighbor -  Nearest Neighbor"
+                  "bilinear - Bilinear Interpolation");
+    TVM_ATTR_FIELD(extrapolation_value).set_default(0.0)
+        .describe("Specify value for extrapolation.");
+    TVM_ATTR_FIELD(out_dtype)
+        .set_default(NullValue<DataType>())
+        .describe("Output data type.");
+  }
+};
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_IMAGE_H_

python/tvm/relay/frontend/tensorflow.py

@@ -546,47 +546,20 @@ def _impl(inputs, attr, params):
         # input image is a 4-D tensor of shape [batch, image_height, image_width, depth]
         # boxes is a 2-D tensor of shape [num_boxes, 4], 4 is for [y1, x1, y2, x2]
         try:
-            boxes = _get_list_param(params, inputs[1])
-            box_ind = _get_list_param(params, inputs[2])
             crop_size = _get_list_param(params, inputs[3])
         except (IndexError, KeyError):
-            boxes = _infer_value(inputs[1], params).asnumpy().tolist()
-            box_ind = _infer_value(inputs[2], params).asnumpy().tolist()
             crop_size = _infer_value(inputs[3], params).asnumpy().tolist()
 
-        data_shape = attr['_input_shapes'][inputs[0]]
-        data_dim = len(data_shape)
         method = attr['method'].decode()
-
-        attrs = {}
-        attrs['size'] = crop_size
-        attrs['layout'] = 'NHWC'
-        if method.lower() == 'nearest':
+        method = 'nearest_neighbor' if method == 'nearest' else method
+        if method not in ['bilinear', 'nearest_neighbor']:
             raise tvm.error.OpAttributeUnImplemented(
-                'Attribute method=nearest is not supported')
-        else:
-            attrs['coordinate_transformation_mode'] = 'align_corners'
-            attrs['method'] = 'bilinear'
-
-        out = None
-        begin = [0] * data_dim
-        size = data_shape[:]
-        for idx in box_ind:
-            # 1) Crop
-            # y is mapped to the image coordinate at y * (image_height - 1)
-            # x is mapped to the image coordinate at x * (image_width - 1)
-            begin[0] = idx
-            begin[1] = int(round(boxes[idx][0] * (data_shape[1] - 1)))
-            begin[2] = int(round(boxes[idx][1] * (data_shape[2] - 1)))
-            size[0] = idx + 1
-            size[1] = int(round((data_shape[1] - 1) * boxes[idx][2])) + 1
-            size[2] = int(round((data_shape[2] - 1) * boxes[idx][3])) + 1
-            res_crop = _op.strided_slice(inputs[0], begin=begin, end=size)
-
-            # 2) Resize
-            res_resize = get_relay_op('resize')(res_crop, **attrs)
-            out = _op.concatenate([out, res_resize], axis=0) if out else res_resize
-        return out
+                'Method {} is not supported'.format(method))
+        layout = attr['layout'] if 'layout' in attr else 'NHWC'
+        extrapolation_value = attr['extrapolation_value']
+
+        return get_relay_op("crop_and_resize")(inputs[0], inputs[1], inputs[2], crop_size,
+                                               layout, method, extrapolation_value)
     return _impl
 
 def _cast():

python/tvm/relay/op/image/_image.py

@@ -25,7 +25,6 @@
 # resize
 reg.register_schedule("image.resize", schedule_injective)
 
-
 @reg.register_compute("image.resize")
 def compute_resize(attrs, inputs, out_type, target):
     size = attrs.size
@@ -34,3 +33,18 @@ def compute_resize(attrs, inputs, out_type, target):
     coord_trans = attrs.coordinate_transformation_mode
     out_dtype = attrs.out_dtype
     return [topi.image.resize(inputs[0], size, layout, method, coord_trans, out_dtype)]
+
+
+# crop and resize
+reg.register_schedule("image.crop_and_resize", schedule_injective)
+
+@reg.register_compute("image.crop_and_resize")
+def compute_crop_and_resize(attrs, inputs, out_type, target):
+    crop_size = attrs.crop_size
+    layout = attrs.layout
+    method = attrs.method
+    extrapolation_value = attrs.extrapolation_value
+    out_dtype = attrs.out_dtype
+    return [topi.image.crop_and_resize(inputs[0], inputs[1], inputs[2],
+                                       crop_size, layout, method,
+                                       extrapolation_value, out_dtype)]

python/tvm/relay/op/image/image.py

@@ -31,7 +31,7 @@ def resize(data,
     with data of shape (n, c, h, w)
     out will have a shape (n, c, size[0], size[1])
 
-    method indicates the algorithm to be used while calculating ghe out value
+    method indicates the algorithm to be used while calculating the out value
     and method can be one of ("bilinear", "nearest_neighbor", "bicubic")
 
     Parameters
@@ -63,3 +63,53 @@ def resize(data,
         The resized result.
     """
     return _make.resize(data, size, layout, method, coordinate_transformation_mode, out_dtype)
+
+
+def crop_and_resize(data,
+                    boxes,
+                    box_indices,
+                    crop_size,
+                    layout,
+                    method="bilinear",
+                    extrapolation_value=0,
+                    out_dtype=None):
+    """Crop input images and resize them.
+
+    method indicates the algorithm to be used while calculating the out value
+    and method can be either "bilinear" or "nearest_neighbor".
+
+    Parameters
+    ----------
+    data : relay.Expr
+        The input data to the operator.
+
+    boxes : relay.Expr
+        A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies
+        the coordinates of a box.
+
+    box_indices : relay.Expr
+        A 1-D tensor of shape [num_boxes], box_ind[i] specifies the data that
+        the i-th box refers to.
+
+    crop_size : Tuple of Expr
+        The target size to which each box will be resized.
+
+    layout : str, optional
+        Layout of the input.
+
+    method : str, optional
+        Scale method, it can be either "nearest_neighbor" or "bilinear".
+
+    extrapolation_value : float, optional
+        Value used for extrapolation, when applicable.
+
+    out_dtype : str, optional
+        Type to return. If left None returns the same type as input.
+
+    Returns
+    -------
+    result: relay.Expr
+        The computed result.
+    """
+    return _make.crop_and_resize(data, boxes, box_indices, crop_size,
+                                 layout, method, extrapolation_value, out_dtype)

python/tvm/relay/op/op_attrs.py

@@ -114,6 +114,9 @@ class DeformableConv2DAttrs(Attrs):
 class ResizeAttrs(Attrs):
     """Attributes for image.resize"""
 
+@register_relay_attr_node
+class CropAndResizeAttrs(Attrs):
+    """Attributes for image.crop_and_resize"""
 
 @register_relay_attr_node
 class ArgsortAttrs(Attrs):

src/lang/expr_operator.cc

@@ -246,8 +246,8 @@ PrimExpr div(PrimExpr a, PrimExpr b) {
 }
 
 PrimExpr truncdiv(PrimExpr a, PrimExpr b) {
-  CHECK(a.dtype().is_int() || a.dtype().is_uint());
-  CHECK(b.dtype().is_int() || b.dtype().is_uint());
+  CHECK(a.dtype().is_int() || a.dtype().is_uint()) << a;
+  CHECK(b.dtype().is_int() || b.dtype().is_uint()) << b;
   return div(a, b);
 }
 
@@ -276,17 +276,17 @@ PrimExpr indexmod(PrimExpr a, PrimExpr b) {
 }
 
 PrimExpr floordiv(PrimExpr a, PrimExpr b) {
-  CHECK(a.dtype().is_int() || a.dtype().is_uint());
-  CHECK(b.dtype().is_int() || b.dtype().is_uint());
+  CHECK(a.dtype().is_int() || a.dtype().is_uint()) << a;
+  CHECK(b.dtype().is_int() || b.dtype().is_uint()) << b;
   BinaryOpMatchTypes(a, b);
   PrimExpr ret = arith::TryConstFold<ir::FloorDivNode>(a, b);
   if (ret.defined()) return ret;
   return ir::FloorDivNode::make(a, b);
 }
 
 PrimExpr floormod(PrimExpr a, PrimExpr b) {
-  CHECK(a.dtype().is_int() || a.dtype().is_uint());
-  CHECK(b.dtype().is_int() || b.dtype().is_uint());
+  CHECK(a.dtype().is_int() || a.dtype().is_uint()) << a;
+  CHECK(b.dtype().is_int() || b.dtype().is_uint()) << b;
   BinaryOpMatchTypes(a, b);
   PrimExpr ret = arith::TryConstFold<ir::FloorModNode>(a, b);
   if (ret.defined()) return ret;

src/relay/op/image/resize.cc

@@ -109,5 +109,89 @@ RELAY_REGISTER_OP("image.resize")
 .add_type_rel("Resize", ResizeRel)
 .set_attr<TOpPattern>("TOpPattern", kInjective);
 
+
+TVM_REGISTER_NODE_TYPE(CropAndResizeAttrs);
+
+bool CropAndResizeRel(const Array<Type>& types,
+                      int num_inputs,
+                      const Attrs& attrs,
+                      const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 4);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* boxes = types[1].as<TensorTypeNode>();
+  const auto* box_indices = types[2].as<TensorTypeNode>();
+  if (data == nullptr || boxes == nullptr ||
+      box_indices == nullptr) return false;
+
+  const CropAndResizeAttrs* param = attrs.as<CropAndResizeAttrs>();
+  CHECK(param != nullptr);
+  auto crop_size = param->crop_size;
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+
+  // 4-D tensor of shape [num_boxes, crop_height, crop_width, depth]
+  static const Layout kNCHW("NCHW");
+  const Layout in_layout(param->layout);
+  auto layout_converter = BijectiveLayoutNode::make(in_layout, kNCHW);
+  auto oshape = layout_converter.ForwardShape(data->shape);
+  oshape.Set(0, box_indices->shape[0]);
+  oshape.Set(2, crop_size[0]);
+  oshape.Set(3, crop_size[1]);
+  auto bshape = layout_converter.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[3],
+                   TensorTypeNode::make(layout_converter.BackwardShape(oshape),
+                                        out_dtype));
+  return true;
+}
+
+Expr MakeCropAndResize(Expr data,
+                       Expr boxes,
+                       Expr box_indices,
+                       Array<IndexExpr> crop_size,
+                       std::string layout,
+                       std::string method,
+                       double extrapolation_value,
+                       DataType out_dtype) {
+  auto attrs = make_object<CropAndResizeAttrs>();
+  attrs->crop_size = std::move(crop_size);
+  attrs->layout = std::move(layout);
+  attrs->method = std::move(method);
+  attrs->extrapolation_value = std::move(extrapolation_value);
+  attrs->out_dtype = out_dtype;
+  static const Op& op = Op::Get("image.crop_and_resize");
+  return CallNode::make(op, {data, boxes, box_indices}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.image._make.crop_and_resize")
+.set_body_typed(MakeCropAndResize);
+
+
+RELAY_REGISTER_OP("image.crop_and_resize")
+    .describe(R"code(Perform crop and resize to input array with nearest neighbour or bilinear interpolation.
+
+- **data**: data is 4D array of shape
+            (batch_size, channels, in_height, in_width) for NCHW
+            (batch_size, in_height, in_width, channels) for NHWC
+
+- **out**: Output is 4D array of shape
+           for layout NCHW
+           (batch_size, channels, crop_size[0], crop_size[1])
+
+           for layout NHWC
+           (batch_size, crop_size[0], crop_size[1], channels)
+)code" TVM_ADD_FILELINE)
+.set_num_inputs(3)
+.add_argument("data", "Tensor", "The input tensor.")
+.add_argument("boxes", "Tensor", "The boxes tensor.")
+.add_argument("box_indices", "Tensor", "The box indices tensor.")
+.set_attrs_type<CropAndResizeAttrs>()
+.set_support_level(5)
+.add_type_rel("CropAndResize", CropAndResizeRel)
+.set_attr<TOpPattern>("TOpPattern", kInjective);
+
 }  // namespace relay
 }  // namespace tvm