yongwww commented on a change in pull request #4417: [TOPI][RELAY][OP] add op
crop_and_resize
URL: https://github.com/apache/incubator-tvm/pull/4417#discussion_r351021052
##########
File path: topi/python/topi/image/resize.py
##########
@@ -210,3 +210,185 @@ def _bicubic(*indices):
raise ValueError('%s method is not supported.' % method)
return tvm.compute(output_shape, compute_func, name='resize',
tag=tag.INJECTIVE)
+
+
+def crop_and_resize(data, boxes, box_indices, crop_size, layout="NCHW",
+ method="bilinear", extrapolation_value=0, out_dtype=None):
+ """Perform crop and resize operation on the data.
+
+ Parameters
+ ----------
+ data : tvm.Tensor
+ inputs is a 4-D tensor with shape
+ [batch, channel, in_height, in_width]
+ or [batch, in_height, in_width, channel]
+
+ boxes : tvm.Tensor
+ A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies
+ the coordinates of a box.
+
+ box_indices : tvm.Tensor
+ A 1-D tensor of shape [num_boxes], box_indices[i] specifies the data
that
+ the i-th box refers to.
+
+ crop_size : Tuple
+ The target size of each box.
+
+ layout : string, optional
+ "NCHW", "NHWC"
+
+ method : {"bilinear", "nearest_neighbor"}
+ Method to be used for resizing.
+
+ extrapolation_value: float, optional
+ Value used for extrapolation, when applicable.
+
+ out_dtype : string, optional
+ Type to return. If left None will be same as input type.
+
+ Returns
+ -------
+ output : tvm.Tensor
+ 4-D with shape [num_boxes, channel, crop_height, crop_width]
+ or [num_boxes, crop_height, crop_width, channel]
+ """
+ method = method.lower()
+ target_h = crop_size[0]
+ target_w = crop_size[1]
+
+ if layout == 'NHWC':
+ output_shape = [box_indices.shape[0], crop_size[0], crop_size[1],
data.shape[3]]
+ image_height = data.shape[1]
+ image_width = data.shape[2]
+ elif layout == 'NCHW':
+ output_shape = [box_indices.shape[0], data.shape[1], crop_size[0],
crop_size[1]]
+ image_height = data.shape[2]
+ image_width = data.shape[3]
+ # Otherwise layout must be NCHWxc
+ else:
+ output_shape = [box_indices.shape[0], data.shape[1],
+ crop_size[0], crop_size[1], data.shape[4]]
+ image_height = data.shape[2]
+ image_width = data.shape[3]
+
+ def _get_pixel(n, c, y, x, cc):
+ if layout.lower() == 'nhwc':
+ return data(n, y.astype("int32"), x.astype("int32"),
c).astype('float')
+ if layout.lower() == 'nchw':
+ return data(n, c, y.astype("int32"),
x.astype("int32")).astype('float')
+ # else must be NCHWxc
+ return data(n, c, y.astype("int32"), x.astype("int32"),
cc).astype('float')
+
+ def _get_indices(*indices):
+ if layout == 'NHWC':
+ n, y, x, c = indices
+ cc = None
+ elif layout == 'NCHW':
+ n, c, y, x = indices
+ cc = None
+ else:
+ n, c, y, x, cc = indices
+
+ return n, c, y, x, cc
+
+ def _cast_output(value):
+ if out_dtype:
+ dtype = out_dtype
+ else:
+ dtype = data.dtype
+ return value.astype(dtype)
+
+ # Nearest neighbor computation
+ def _nearest_neighbor(*indices):
+ n, c, y, x, cc = _get_indices(*indices)
+ box_idx = box_indices(n)
+
+ y1, x1 = boxes(n, 0), boxes(n, 1)
+ y2, x2 = boxes(n, 2), boxes(n, 3)
+
+ in_h = (image_height - 1) * (y2 - y1)
+ in_w = (image_width - 1) * (x2 - x1)
+ h_scale = tvm.div(in_h, target_h - 1)
+ w_scale = tvm.div(in_w, target_w - 1)
+
+ in_y = y1 * (image_height - 1) + h_scale * y
+ in_x = x1 * (image_width - 1) + w_scale * x
+ closest_x_index = tvm.round(in_x)
+ closest_y_index = tvm.round(in_y)
+
+ value = _get_pixel(box_idx, c, closest_y_index, closest_x_index, cc)
+ out_y = tvm.if_then_else(in_y < 0,
+ extrapolation_value,
+ tvm.if_then_else(in_y > image_height - 1,
+ extrapolation_value,
+ value))
+
+ # use extrapolation_value if in_x is out of boundary
+ out = tvm.if_then_else(in_x < 0,
+ extrapolation_value,
+ tvm.if_then_else(in_x > image_width - 1,
+ extrapolation_value,
+ out_y))
+ return _cast_output(out)
+
+
+ # Bilinear helper functions and computation.
+ def _lerp(A, B, t):
+ return A * (1.0 - t) + B * t
+
+ def _bilinear(*indices):
+ n, c, y, x, cc = _get_indices(*indices)
+ box_idx = box_indices(n)
+
+ y1, x1 = boxes(n, 0), boxes(n, 1)
+ y2, x2 = boxes(n, 2), boxes(n, 3)
+
+ in_h = (image_height - 1) * (y2 - y1)
+ in_w = (image_width - 1) * (x2 - x1)
+ h_scale = tvm.div(in_h, target_h - 1)
+ w_scale = tvm.div(in_w, target_w - 1)
Review comment:
It is sligthtly different from crop, for example, h_scale is calculated
based on boxes while crop has no boxes, etc. I have moved these kernels out of
the op and generalized them to work for both resize and crop_and_resize.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
