1327 lines
42 KiB
Python
1327 lines
42 KiB
Python
# 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)
|