Files
apache--tvm/python/tvm/topi/image/resize.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Has been cancelled
CI / MacOS (push) Has been cancelled
CI / Windows (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

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)