# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name # ruff: noqa: E741, F821 """TVM operator input resize compute.""" import tvm from tvm import te from tvm.topi.utils import nchw_pack_layout, nchw_xc_layout from .. import tag def can_convert_multiply_to_intdiv(origin_size, scaled_size): """Check whether can convert multiplication to division""" # Only support IntImm type if not isinstance(scaled_size, tvm.tirx.expr.IntImm): return False div = scaled_size / origin_size.astype("float") if div.value % 1 != 0: return False epsilon = 1e-5 check = 1 / (epsilon * origin_size + epsilon) if div > check: return False return True def get_1d_indices(indices, layout="NCW"): """Get 1d indices""" (cc, inum, ic) = (0, 0, 0) if layout == "NWC": n, x, c = indices cc = None elif layout == "NCW": n, c, x = indices cc = None elif ncw_pack_layout(layout): n, c, x, inum, ic = indices else: # else must be NCHWxc assert ncw_xc_layout(layout) n, c, x, cc = indices return n, c, x, cc, inum, ic def get_2d_indices(indices, layout="NCHW"): """Get 2d indices""" (cc, inum, ic) = (0, 0, 0) if layout == "NHWC": n, y, x, c = indices cc = None elif layout == "NCHW": n, c, y, x = indices cc = None elif nchw_pack_layout(layout): n, c, y, x, inum, ic = indices else: # else must be NCHWxc assert nchw_xc_layout(layout) n, c, y, x, cc = indices return n, c, y, x, cc, inum, ic def get_3d_indices(indices, layout="NCDHW"): """Get 3d indices""" if layout == "NDHWC": n, z, y, x, c = indices cc = None elif layout == "NCDHW": n, c, z, y, x = indices cc = None else: n, c, z, y, x, cc = indices return n, c, z, y, x, cc def get_1d_pixel(data, layout, image_width, n, c, x, cc, ib, ic): """Get 1d pixel""" x = tvm.te.max(tvm.te.min(x, image_width - 1), 0) if layout == "NWC": return data(n, x, c).astype("float") if layout == "NCW": return data(n, c, x).astype("float") if ncw_pack_layout(layout): return data(n, c, x, ib, ic).astype("float") # else must be NCHWxc assert ncw_xc_layout(layout) return data(n, c, x, cc).astype("float") def get_2d_pixel(data, layout, image_height, image_width, n, c, y, x, cc, ib, ic): """Get 2d pixel""" y = tvm.te.max(tvm.te.min(y, image_height - 1), 0) x = tvm.te.max(tvm.te.min(x, image_width - 1), 0) if layout == "NHWC": return data(n, y, x, c).astype("float") if layout == "NCHW": return data(n, c, y, x).astype("float") if nchw_pack_layout(layout): return data(n, c, y, x, ib, ic).astype("float") # else must be NCHWxc assert nchw_xc_layout(layout) return data(n, c, y, x, cc).astype("float") def get_3d_pixel(data, layout, image_depth, image_height, image_width, n, c, z, y, x, cc): """Get 3d pixel""" z = tvm.te.max(tvm.te.min(z, image_depth - 1), 0) y = tvm.te.max(tvm.te.min(y, image_height - 1), 0) x = tvm.te.max(tvm.te.min(x, image_width - 1), 0) if layout == "NDHWC": return data(n, z, y, x, c).astype("float") if layout == "NCDHW": return data(n, c, z, y, x).astype("float") # else must be NCDHWxc return data(n, c, z, y, x, cc).astype("float") def get_inx( x, image_width, target_width, coordinate_transformation_mode, start_x=0, end_x=-1, use_int_div=False, ): """Infer input x from output x with various coordinate transformation methods""" scale_x = te.div(image_width.astype("float"), target_width.astype("float")) if coordinate_transformation_mode == "half_pixel": in_x = (x + 0.5) * scale_x - 0.5 elif coordinate_transformation_mode == "align_corners": in_x = (image_width - 1).astype("float") / (target_width - 1) * x elif coordinate_transformation_mode == "asymmetric": if use_int_div: in_x = te.div(x, te.div(target_width, image_width)) else: in_x = scale_x * x elif coordinate_transformation_mode == "pytorch_half_pixel": in_x = te.if_then_else(target_width > 1, (x + 0.5) * scale_x - 0.5, 0.0) elif coordinate_transformation_mode == "tf_half_pixel_for_nn": in_x = (x + 0.5) * scale_x elif coordinate_transformation_mode == "tf_crop_and_resize": in_x = te.if_then_else( target_width > 1, start_x * (image_width - 1) + x * (end_x - start_x) * (image_width - 1).astype("float") / (target_width - 1), 0.5 * (start_x + end_x) * (image_width - 1), ) else: raise ValueError( f"Unsupported coordinate_transformation_mode: {coordinate_transformation_mode}" ) return in_x def get_closest_index(in_x, rounding_method, boxes, use_int_div=False): """get the closest index to a value based on a certain rounding method""" if use_int_div: closest_x_index = in_x.astype("int32") return closest_x_index if rounding_method == "round" or boxes is not None: closest_x_index = te.round(in_x).astype("int32") elif rounding_method == "round_prefer_floor": closest_x_index = te.ceil(in_x - 0.5).astype("int32") elif rounding_method == "round_prefer_ceil": closest_x_index = te.floor(in_x + 0.5).astype("int32") elif rounding_method == "floor": # Add epsilon to floor to prevent gpu rounding errors. epsilon = 1e-5 closest_x_index = te.floor(in_x + epsilon).astype("int32") elif rounding_method == "ceil": # Subract epsilon from ceil to prevent gpu rounding errors. epsilon = 1e-5 closest_x_index = te.ceil(in_x - epsilon).astype("int32") else: raise ValueError(f"Unknown rounding method: {rounding_method}") return closest_x_index def _lerp(A, B, t): """Perform Linear interpolation in 1D""" return A * (1.0 - t) + B * t def _cubic_spline_weights(t, alpha): """create cubic spline weights in 1D""" t2 = t * t t3 = t * t * t w1 = alpha * (t3 - 2 * t2 + t) w2 = (alpha + 2) * t3 - (3 + alpha) * t2 + 1 w3 = -(alpha + 2) * t3 + (3 + 2 * alpha) * t2 - alpha * t w4 = -alpha * t3 + alpha * t2 return [w1, w2, w3, w4] def _cubic_kernel(inputs, w): """perform cubic interpolation in 1D""" return sum([a_i * w_i for a_i, w_i in zip(inputs, w)]) def _resize_1d( indices, data, roi, image_width, target_width, boxes=None, box_indices=None, method=None, extrapolation_value=0.0, layout="NCW", coordinate_transformation_mode="align_corners", rounding_method="", alpha=-0.5, exclude_outside=0, out_dtype=None, ): """Perform resize operation on the data with selected method and options. Parameters ---------- indices : tuple The indices of input data data : tvm.te.Tensor inputs is a 3-D tensor with shape [batch, channel, in_width] or [batch, in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 2, and format [start_w, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. image_width : integer Input image width target_width : integer The target resized image width boxes : tvm.te.Tensor, optional A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies the coordinates of a box. box_indices : tvm.te.Tensor, optional A 1-D tensor of shape [num_boxes], box_indices[i] specifies the data that the i-th box refers to. extrapolation_value: float, optional Value used for extrapolation, when applicable. layout: string, optional "NCW", "NWC", or "NCWc". method: string, optional method of interpolation ("nearest", "linear", "bicubic") coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: string, optional indicates how to find the "nearest" pixel in nearest_neighbor method [round, floor, ceil] alpha: float, optional Bicubic spline coefficient exclude_outside: bool, optional: Exclude values outside the image fdor bicubic interpolation out_dtype: string, optional Type to return. If left None will be same as input type. Returns ------- output : out_dtype The computed result with type out_dtype """ def _cast_output(value, data_dtype="float32", out_dtype=None): if out_dtype: dtype = out_dtype else: dtype = data_dtype return value.astype(dtype) n, c, x, cc, inum, ic = get_1d_indices(indices, layout) box_idx = box_indices(n) if box_indices is not None else n if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[1]) if method == "nearest_neighbor": if rounding_method == "": if coordinate_transformation_mode == "align_corners": rounding_method = "round" else: rounding_method = "floor" closest_x_index = get_closest_index(in_x, rounding_method, boxes) value = get_1d_pixel(data, layout, image_width, box_idx, c, closest_x_index, cc, inum, ic) elif method == "linear": x_int = te.floor(in_x).astype("int32") x_lerp = in_x - x_int p = [0 for i in range(2)] for i in range(2): p[i] = get_1d_pixel(data, layout, image_width, box_idx, c, x_int + i, cc, inum, ic) value = _lerp(*p, x_lerp) elif method == "cubic": xint = te.floor(in_x).astype("int32") xfract = in_x - te.floor(in_x) # Get the surrounding values p = [0 for i in range(4)] for i in range(4): p[i] = get_1d_pixel(data, layout, image_width, box_idx, c, xint + i - 1, cc, inum, ic) wx = _cubic_spline_weights(xfract, alpha) if exclude_outside: for i in range(4): wx[i] = te.if_then_else( te.any(xint - 1 + i < 0, xint + i > image_width), 0.0, wx[i] ) sum_wx = sum(wx) wx = [w / sum_wx for w in wx] value = _cubic_kernel(p, wx) else: raise ValueError("Unknown resize method:", method) if coordinate_transformation_mode == "tf_crop_and_resize": # use extrapolation_value if in_x is out of boundary value = tvm.tirx.if_then_else( in_x < 0, extrapolation_value, tvm.tirx.if_then_else(in_x > image_width - 1, extrapolation_value, value), ) return _cast_output(value, data.dtype, out_dtype=out_dtype) def resize1d( data, roi, size, layout="NCW", method="linear", coordinate_transformation_mode="half_pixel", rounding_method="", bicubic_alpha=-0.75, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None, ): """Perform resize operation on the data. Parameters ---------- data : tvm.te.Tensor inputs is a 3-D tensor with shape [batch, channel in_width] or [batch in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 2, and format [start_w, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. size: Tuple Output resolution scale to layout: string, optional "NCW", "NWC", or "NCWc". coordinate_transformation_mode: string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. Refer to the ONNX Resize operator specification for details. Available options are "half_pixel", "align_corners" and "asymmetric". method: string, optional method of interpolation ("nearest", "linear", "bicubic") coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: Method for rounding coordinate locations bicubic_alpha: float, optional Bicubic spline coefficient bicubic_exclude: bool, optional: Exclude values outside the image fdor bicubic interpolation 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. output_shape: tvm_ffi.Array, optional Shape to return. If left None will be inferred (If shape is determined dynamically, pass out_dtype.shape as output_shape) Returns ------- output : tvm.te.Tensor 4-D with shape [batch, chananel, in_width*scale] or [batch, in_width*scale, channel] or 5-D with shape [batch, channel-major, in_width*scale, channel-minor] """ method = method.lower() if layout == "NWC": in_n, in_w, in_c = data.shape if output_shape is None: output_shape = [in_n, size[0], in_c] elif layout == "NCW": in_n, in_c, in_w = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0]] elif ncw_pack_layout(layout): # for NCWinic in_n, in_c, in_w, in_inum, in_ic = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0], in_inum, in_ic] elif ncw_xc_layout(layout): # for NCWxc in_n, in_c, in_w, in_cc = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0], in_cc] else: raise ValueError(f"{layout} layout is not supported.") if isinstance(size, tuple): size = list(size) for i in range(1): if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) def compute_func(*indices): return _resize_1d( indices, data, roi, in_w, size[0], method=method, layout=layout, coordinate_transformation_mode=coordinate_transformation_mode, rounding_method=rounding_method, alpha=bicubic_alpha, exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) def _resize_2d( indices, data, roi, image_height, image_width, target_height, target_width, boxes=None, box_indices=None, method=None, extrapolation_value=0.0, layout="NCHW", coordinate_transformation_mode="align_corners", rounding_method="", alpha=-0.5, exclude_outside=0, out_dtype=None, ): """Perform resize operation on the data with selected method and options. Parameters ---------- indices : tuple The indices of input data data : tvm.te.Tensor inputs is a 4-D tensor with shape [batch, channel, in_height, in_width] or [batch, in_height, in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 4, and format [start_h, start_w, end_h, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. image_height : integer Input image height image_width : integer Input image width target_height : integer The target resized image height target_width : integer The target resized image width boxes : tvm.te.Tensor, optional A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies the coordinates of a box. method: string, optional method of interpolation ("nearest", "linear", "bicubic") box_indices : tvm.te.Tensor, optional A 1-D tensor of shape [num_boxes], box_indices[i] specifies the data that the i-th box refers to. extrapolation_value: float, optional Value used for extrapolation, when applicable. layout: string, optional "NCHW", "NHWC", or "NCHWc". coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: string, optional indicates how to find the "nearest" pixel in nearest_neighbor method [round, floor, ceil] alpha: float, optional Bicubic spline coefficient exclude_outside: bool, optional: Exclude values outside the image fdor bicubic interpolation out_dtype: string, optional Type to return. If left None will be same as input type. Returns ------- output : out_dtype The computed result with type out_dtype """ def _cast_output(value, data_dtype="float32", out_dtype=None): if out_dtype: dtype = out_dtype else: dtype = data_dtype return value.astype(dtype) height_use_int_div = False width_use_int_div = False if method == "nearest_neighbor" and coordinate_transformation_mode == "asymmetric": if rounding_method == "floor" or rounding_method == "": height_use_int_div = can_convert_multiply_to_intdiv(image_height, target_height) width_use_int_div = can_convert_multiply_to_intdiv(image_width, target_width) n, c, y, x, cc, inum, ic = get_2d_indices(indices, layout) box_idx = box_indices(n) if box_indices is not None else n if boxes is not None: 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 = in_h.astype("float") / (target_height - 1) w_scale = in_w.astype("float") / (target_width - 1) in_y = y1 * (image_height - 1) + h_scale * y in_x = x1 * (image_width - 1) + w_scale * x else: in_x = get_inx( x, image_width, target_width, coordinate_transformation_mode, roi[1], roi[3], width_use_int_div, ) in_y = get_inx( y, image_height, target_height, coordinate_transformation_mode, roi[0], roi[2], height_use_int_div, ) if method == "nearest_neighbor": if rounding_method == "": if coordinate_transformation_mode == "align_corners": rounding_method = "round" else: rounding_method = "floor" closest_x_index = get_closest_index(in_x, rounding_method, boxes, width_use_int_div) closest_y_index = get_closest_index(in_y, rounding_method, boxes, height_use_int_div) value = get_2d_pixel( data, layout, image_height, image_width, box_idx, c, closest_y_index, closest_x_index, cc, inum, ic, ) elif method == "linear": y_int = te.floor(in_y).astype("int32") x_int = te.floor(in_x).astype("int32") y_lerp = in_y - y_int x_lerp = in_x - x_int p = [[0 for i in range(2)] for j in range(2)] for j in range(2): for i in range(2): p[j][i] = get_2d_pixel( data, layout, image_height, image_width, box_idx, c, y_int + j, x_int + i, cc, inum, ic, ) top = _lerp(*p[0], x_lerp) bottom = _lerp(*p[1], x_lerp) value = _lerp(top, bottom, y_lerp) elif method == "cubic": xint = te.floor(in_x).astype("int32") xfract = in_x - te.floor(in_x) yint = te.floor(in_y).astype("int32") yfract = in_y - te.floor(in_y) # Get the surrounding values p = [[0 for i in range(4)] for j in range(4)] for j in range(4): for i in range(4): p[j][i] = get_2d_pixel( data, layout, image_height, image_width, box_idx, c, yint + j - 1, xint + i - 1, cc, inum, ic, ) wx = _cubic_spline_weights(xfract, alpha) wy = _cubic_spline_weights(yfract, alpha) if exclude_outside: for i in range(4): wx[i] = te.if_then_else( te.any(xint - 1 + i < 0, xint + i > image_width), 0.0, wx[i] ) wy[i] = te.if_then_else( te.any(yint - 1 + i < 0, yint + i > image_height), 0.0, wy[i] ) sum_wx = sum(wx) sum_wy = sum(wy) wx = [w / sum_wx for w in wx] wy = [w / sum_wy for w in wy] col0 = _cubic_kernel(p[0], wx) col1 = _cubic_kernel(p[1], wx) col2 = _cubic_kernel(p[2], wx) col3 = _cubic_kernel(p[3], wx) value = _cubic_kernel([col0, col1, col2, col3], wy) else: raise ValueError("Unknown resize method:", method) if coordinate_transformation_mode == "tf_crop_and_resize": out = tvm.tirx.if_then_else( in_y < 0, extrapolation_value, tvm.tirx.if_then_else(in_y > image_height - 1, extrapolation_value, value), ) # use extrapolation_value if in_x is out of boundary value = tvm.tirx.if_then_else( in_x < 0, extrapolation_value, tvm.tirx.if_then_else(in_x > image_width - 1, extrapolation_value, out), ) return _cast_output(value, data.dtype, out_dtype=out_dtype) def resize2d( data, roi, size, layout="NCHW", method="linear", coordinate_transformation_mode="half_pixel", rounding_method="", bicubic_alpha=-0.75, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None, ): """Perform resize operation on the data. Parameters ---------- data : tvm.te.Tensor inputs is a 4-D tensor with shape [batch, channel, in_height, in_width] or [batch, in_height, in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 4, and format [start_h, start_w, end_h, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. size: Tuple Output resolution scale to layout: string, optional "NCHW", "NHWC", or "NCHWc". method: string, optional method of interpolation ("nearest", "linear", "bicubic") coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: Method for rounding coordinate locations bicubic_alpha: float, optional Bicubic spline coefficient bicubic_exclude: bool, optional: Exclude values outside the image fdor bicubic interpolation 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. output_shape: tvm_ffi.Array, optional Shape to return. If left None will be inferred (If shape is determined dynamically, pass out_dtype.shape as output_shape) Returns ------- output : tvm.te.Tensor 4-D with shape [batch, channel, in_height*scale, in_width*scale] or [batch, in_height*scale, in_width*scale, channel] or 5-D with shape [batch, channel-major, in_height*scale, in_width*scale, channel-minor] """ method = method.lower() if layout == "NHWC": in_n, in_h, in_w, in_c = data.shape if output_shape is None: output_shape = [in_n, size[0], size[1], in_c] elif layout == "NCHW": in_n, in_c, in_h, in_w = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0], size[1]] elif nchw_pack_layout(layout): # for NCHWinic in_n, in_c, in_h, in_w, in_inum, in_ic = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0], size[1], in_inum, in_ic] elif nchw_xc_layout(layout): # for NCHWxc in_n, in_c, in_h, in_w, in_cc = data.shape if output_shape is None: output_shape = [in_n, in_c, size[0], size[1], in_cc] else: raise ValueError(f"{layout} layout is not supported.") if isinstance(size, tuple): size = list(size) for i in range(2): if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) def compute_func(*indices): return _resize_2d( indices, data, roi, in_h, in_w, size[0], size[1], method=method, layout=layout, coordinate_transformation_mode=coordinate_transformation_mode, rounding_method=rounding_method, alpha=bicubic_alpha, exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, ) return te.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=None, out_dtype=None, ): """Perform crop and resize operation on the data. Parameters ---------- data : tvm.te.Tensor inputs is a 4-D tensor with shape [batch, channel, in_height, in_width] or [batch, in_height, in_width, channel] boxes : tvm.te.Tensor A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies the coordinates of a box. box_indices : tvm.te.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.te.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_h = data.shape[1].astype("int32") image_w = data.shape[2].astype("int32") elif layout == "NCHW": output_shape = [box_indices.shape[0], data.shape[1], crop_size[0], crop_size[1]] image_h = data.shape[2].astype("int32") image_w = data.shape[3].astype("int32") elif layout.startswith("NCHW"): # for NCHWxc output_shape = [ box_indices.shape[0], data.shape[1], crop_size[0], crop_size[1], data.shape[4], ] image_h = data.shape[2].astype("int32") image_w = data.shape[3].astype("int32") else: raise ValueError(f"{layout} layout is not supported.") if method == "bilinear": method = "linear" def compute_func(*indices): return _resize_2d( indices, data, [0.0] * 4, image_h, image_w, target_h, target_w, boxes, box_indices, method=method, extrapolation_value=extrapolation_value, layout=layout, coordinate_transformation_mode="tf_crop_and_resize", out_dtype=out_dtype, ) return te.compute(output_shape, compute_func, name="crop_and_resize", tag=tag.INJECTIVE) def _resize_3d( indices, data, roi, image_depth, image_height, image_width, target_depth, target_height, target_width, boxes=None, box_indices=None, method=None, extrapolation_value=0.0, layout="NCHW", coordinate_transformation_mode="align_corners", rounding_method="", alpha=-0.5, exclude_outside=0, out_dtype=None, ): """Perform resize operation on the data with selected method and options. Parameters ---------- indices : tuple The indices of input data data : tvm.te.Tensor inputs is a 4-D tensor with shape [batch, channel, in_height, in_width] or [batch, in_height, in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 6, and format [start_d, start_h, start_w, end_d, end_h, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. image_depth : integer Input image depth image_height : integer Input image height image_width : integer Input image width target_depth : integer The target resized image depth target_height : integer The target resized image height target_width : integer The target resized image width boxes : tvm.te.Tensor, optional A 2-D tensor of shape [num_boxes, 4]. Each row of the tensor specifies the coordinates of a box. box_indices : tvm.te.Tensor, optional A 1-D tensor of shape [num_boxes], box_indices[i] specifies the data that the i-th box refers to. method: string, optional method of interpolation ("nearest", "linear", "bicubic") extrapolation_value: float, optional Value used for extrapolation, when applicable. layout: string, optional "NCHW", "NHWC", or "NCHWc". coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: string, optional indicates how to find the "nearest" pixel in nearest_neighbor method [round, floor, ceil] alpha: float, optional Bicubic spline coefficient exclude_oiutside: bool, optional: Exclude values outside the image fdor bicubic interpolation out_dtype: string, optional Type to return. If left None will be same as input type. Returns ------- output : out_dtype The computed result with type out_dtype """ def _cast_output(value, data_dtype="float32", out_dtype=None): if out_dtype: dtype = out_dtype else: dtype = data_dtype return value.astype(dtype) n, c, z, y, x, cc = get_3d_indices(indices, layout) box_idx = box_indices(n) if box_indices is not None else n if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") in_z = get_inx(z, image_depth, target_depth, coordinate_transformation_mode, roi[2], roi[5]) in_y = get_inx(y, image_height, target_height, coordinate_transformation_mode, roi[1], roi[4]) in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[3]) if method == "nearest_neighbor": if rounding_method == "": if coordinate_transformation_mode == "align_corners": rounding_method = "round" else: rounding_method = "floor" closest_z_index = get_closest_index(in_z, rounding_method, boxes) closest_y_index = get_closest_index(in_y, rounding_method, boxes) closest_x_index = get_closest_index(in_x, rounding_method, boxes) value = get_3d_pixel( data, layout, image_depth, image_height, image_width, box_idx, c, closest_z_index, closest_y_index, closest_x_index, cc, ) elif method == "linear": z_int = te.floor(in_z).astype("int32") y_int = te.floor(in_y).astype("int32") x_int = te.floor(in_x).astype("int32") z_lerp = in_z - z_int y_lerp = in_y - y_int x_lerp = in_x - x_int p = [[[0 for i in range(2)] for j in range(2)] for k in range(2)] for k in range(2): for j in range(2): for i in range(2): p[k][j][i] = get_3d_pixel( data, layout, image_depth, image_height, image_width, box_idx, c, z_int + k, y_int + j, x_int + i, cc, ) l = [[0 for i in range(2)] for j in range(2)] for j in range(2): for i in range(2): l[j][i] = _lerp(*p[j][i], x_lerp) top = _lerp(*l[0], y_lerp) bottom = _lerp(*l[1], y_lerp) value = _lerp(top, bottom, z_lerp) elif method == "cubic": zint = te.floor(in_z).astype("int32") zfract = in_z - te.floor(in_z) yint = te.floor(in_y).astype("int32") yfract = in_y - te.floor(in_y) xint = te.floor(in_x).astype("int32") xfract = in_x - te.floor(in_x) # Get the surrounding values p = [[[0 for i in range(4)] for j in range(4)] for k in range(4)] for k in range(4): for j in range(4): for i in range(4): p[k][j][i] = get_3d_pixel( data, layout, image_depth, image_height, image_width, box_idx, c, zint + k - 1, yint + j - 1, xint + i - 1, cc, ) wz = _cubic_spline_weights(zfract, alpha) wy = _cubic_spline_weights(yfract, alpha) wx = _cubic_spline_weights(xfract, alpha) if exclude_outside: for i in range(4): wz[i] = te.if_then_else( te.any(xint - 1 + i < 0, xint + i > image_height), 0.0, wx[i] ) wy[i] = te.if_then_else( te.any(yint - 1 + i < 0, yint + i > image_height), 0.0, wy[i] ) wx[i] = te.if_then_else( te.any(xint - 1 + i < 0, xint + i > image_width), 0.0, wx[i] ) sum_wz = sum(wz) sum_wy = sum(wy) sum_wx = sum(wx) wz = [w / sum_wz for w in wz] wy = [w / sum_wy for w in wy] wx = [w / sum_wx for w in wx] l = [[0 for i in range(4)] for j in range(4)] for j in range(4): for i in range(4): l[j][i] = _cubic_kernel(p[j][i], wx) col0 = _cubic_kernel(l[0], wy) col1 = _cubic_kernel(l[1], wy) col2 = _cubic_kernel(l[2], wy) col3 = _cubic_kernel(l[3], wy) value = _cubic_kernel([col0, col1, col2, col3], wz) else: raise ValueError("Unknown resize method:", method) if coordinate_transformation_mode == "tf_crop_and_resize": out = tvm.tirx.if_then_else( in_z < 0, extrapolation_value, tvm.tirx.if_then_else(in_z > image_depth - 1, extrapolation_value, value), ) out = tvm.tirx.if_then_else( in_y < 0, extrapolation_value, tvm.tirx.if_then_else(in_y > image_height - 1, extrapolation_value, value), ) # use extrapolation_value if in_x is out of boundary value = tvm.tirx.if_then_else( in_x < 0, extrapolation_value, tvm.tirx.if_then_else(in_x > image_width - 1, extrapolation_value, out), ) return _cast_output(value, data.dtype, out_dtype=out_dtype) def resize3d( data, roi, size, layout="NCDHW", method="linear", coordinate_transformation_mode="half_pixel", rounding_method="", bicubic_alpha=-0.75, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None, ): """Perform resize operation on the data. Parameters ---------- data : tvm.te.Tensor inputs is a 5-D tensor with shape [batch, channel, in_depth, in_height, in_width] or [batch, in_depth, in_height, in_width, channel] roi: Tuple of Float or Expr The region of interest for cropping the input image. Expected to be of size 6, and format [start_d, start_h, start_w, end_d, end_h, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize. size: Tuple Output resolution scale to layout: string, optional "NCDHW", "NDHWC", or "NCDHWc". method: string, optional method of interpolation ("nearest", "linear", "bicubic") coordinate_transformation_mode : string, optional Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize]. rounding_method: Method for rounding coordinate locations bicubic_alpha: float, optional Bicubic spline coefficient bicubic_exclude: bool, optional: Exclude values outside the image fdor bicubic interpolation 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. output_shape: tvm_ffi.Array, optional Shape to return. If left None will be inferred (If shape is determined dynamically, pass out_dtype.shape as output_shape) Returns ------- output : tvm.te.Tensor 4-D with shape [batch, channel, in_depth*scale, in_height*scale, in_width*scale] or [batch, in_depth*scale, in_height*scale, in_width*scale, channel] or 5-D with shape [batch, channel-major, in_depth*scale, in_height*scale, in_width*scale, channel-minor] """ method = method.lower() if layout == "NDHWC": in_n, in_d, in_h, in_w, in_c = data.shape output_shape = [in_n, size[0], size[1], size[2], in_c] elif layout == "NCDHW": in_n, in_c, in_d, in_h, in_w = data.shape output_shape = [in_n, in_c, size[0], size[1], size[2]] # Otherwise layout must be NCHWxc else: in_n, in_c, in_d, in_h, in_w, in_cc = data.shape output_shape = [in_n, in_c, size[0], size[1], size[2], in_cc] if isinstance(size, tuple): size = list(size) for i in range(3): if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) def compute_func(*indices): return _resize_3d( indices, data, roi, in_d, in_h, in_w, size[0], size[1], size[2], method=method, layout=layout, coordinate_transformation_mode=coordinate_transformation_mode, rounding_method=rounding_method, alpha=bicubic_alpha, exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE)