wehub-resource-sync
275 lines
8.5 KiB
Python
275 lines
8.5 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.
|
|
"""Image operators."""
|
|
|
|
from typing import cast
|
|
|
|
from tvm import DataType
|
|
from tvm.ir import is_prim_expr
|
|
|
|
from ...expr import Expr, ShapeExpr
|
|
from . import _ffi_api
|
|
|
|
PrimExprLike = int | Expr
|
|
SizeLike = PrimExprLike | tuple[PrimExprLike, ...]
|
|
|
|
|
|
def resize2d(
|
|
data: Expr,
|
|
size: SizeLike,
|
|
roi: float | tuple[float] | None = None,
|
|
layout: str = "NCHW",
|
|
method: str = "linear",
|
|
coordinate_transformation_mode: str = "half_pixel",
|
|
rounding_method: str = "round",
|
|
cubic_alpha: float = -0.75,
|
|
cubic_exclude: int = 0,
|
|
extrapolation_value: float = 0.0,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
"""Image resize2d operator.
|
|
|
|
This operator takes data as input and does 2D scaling to the given scale factor.
|
|
In the default case, where the data_layout is `NCHW`
|
|
with data of shape (n, c, h, w)
|
|
out will have a shape (n, c, size[0], size[1])
|
|
|
|
method indicates the algorithm to be used while calculating the out value
|
|
and method can be one of ("linear", "nearest_neighbor", "cubic")
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
size: SizeLike
|
|
The out size to which the image will be resized.
|
|
If specified as a list, it is required to have length either 1 or 2.
|
|
If specified as an Expr, it is required to have ndim 2.
|
|
|
|
roi: Optional[Union[float, Tuple[float]]]
|
|
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.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
method : str
|
|
Scale method to used [nearest_neighbor, linear, cubic].
|
|
|
|
coordinate_transformation_mode : str
|
|
Describes how to transform the coordinate in the resized tensor
|
|
to the coordinate in the original tensor. Definitions can be found
|
|
in topi/image/resize.py.
|
|
[half_pixel, align_corners, asymmetric, pytorch_half_pixel,
|
|
tf_half_pixel_for_nn, and tf_crop_and_resize].
|
|
|
|
rounding_method: str
|
|
indicates how to find the "nearest" pixel in nearest_neighbor method
|
|
[round, floor, ceil]
|
|
|
|
cubic_alpha: float
|
|
Spline Coefficient for bicubic interpolation
|
|
|
|
cubic_exclude: int
|
|
Flag to exclude exterior of the image during bicubic interpolation
|
|
|
|
extrapolation_value: float
|
|
Fill value to use when roi is outside of the image
|
|
|
|
out_dtype : Optional[str | DataType]
|
|
The dtype of the output tensor.
|
|
It it is not specified, the output will have the same dtype as input if not specified.
|
|
|
|
Returns
|
|
-------
|
|
result: relax.Expr
|
|
The resized result.
|
|
"""
|
|
if roi is None:
|
|
roi = (0.0, 0.0, 0.0, 0.0) # type: ignore
|
|
elif isinstance(roi, float):
|
|
roi = (roi, roi, roi, roi) # type: ignore
|
|
elif isinstance(roi, tuple | list):
|
|
roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
|
|
else:
|
|
raise NotImplementedError(f"Unsupported roi type {type(roi)}")
|
|
|
|
if isinstance(size, int) or is_prim_expr(size):
|
|
size = (size, size)
|
|
if isinstance(size, tuple | list):
|
|
if len(size) == 1:
|
|
size = ShapeExpr([size[0], size[0]])
|
|
else:
|
|
size = ShapeExpr(size)
|
|
|
|
return _ffi_api.resize2d( # type: ignore
|
|
data,
|
|
size,
|
|
roi,
|
|
layout,
|
|
method,
|
|
coordinate_transformation_mode,
|
|
rounding_method,
|
|
cubic_alpha,
|
|
cubic_exclude,
|
|
extrapolation_value,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def resize3d(
|
|
data: Expr,
|
|
size: SizeLike,
|
|
roi: float | tuple[float] | None = None,
|
|
layout: str = "NCDHW",
|
|
method: str = "linear",
|
|
coordinate_transformation_mode: str = "half_pixel",
|
|
rounding_method: str = "",
|
|
cubic_alpha: float = -0.75,
|
|
cubic_exclude: int = 0,
|
|
extrapolation_value: float = 0.0,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
"""Image resize3d operator.
|
|
|
|
This operator takes data as input and does 3D scaling to the given output size.
|
|
In the default case, where data layout is `NCDHW`
|
|
with data of shape (n, c, d, h, w),
|
|
the output has shape (n, c, size[0], size[1], size[2]).
|
|
"""
|
|
if roi is None:
|
|
roi = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0) # type: ignore
|
|
elif isinstance(roi, float):
|
|
roi = (roi, roi, roi, roi, roi, roi) # type: ignore
|
|
elif isinstance(roi, tuple | list):
|
|
roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
|
|
else:
|
|
raise NotImplementedError(f"Unsupported roi type {type(roi)}")
|
|
|
|
if isinstance(size, int) or is_prim_expr(size):
|
|
size = (size, size, size)
|
|
if isinstance(size, tuple | list):
|
|
if len(size) == 1:
|
|
size = ShapeExpr([size[0], size[0], size[0]])
|
|
else:
|
|
size = ShapeExpr(size)
|
|
|
|
return _ffi_api.resize3d( # type: ignore
|
|
data,
|
|
size,
|
|
roi,
|
|
layout,
|
|
method,
|
|
coordinate_transformation_mode,
|
|
rounding_method,
|
|
cubic_alpha,
|
|
cubic_exclude,
|
|
extrapolation_value,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def grid_sample(
|
|
data: Expr,
|
|
grid: Expr,
|
|
method: str = "bilinear",
|
|
layout: str = "NCHW",
|
|
padding_mode: str = "zeros",
|
|
align_corners: bool = False,
|
|
) -> Expr:
|
|
"""Applies grid sampling to input feature map.
|
|
|
|
Given data and grid, the output is computed by sampling from data using
|
|
the grid coordinates.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data tensor with shape [N, C, H, W] for NCHW layout.
|
|
|
|
grid : relax.Expr
|
|
The grid tensor with shape [N, H_out, W_out, 2]. The values are normalized
|
|
to [-1, 1], where (-1, -1) is the top-left corner and (1, 1) is the bottom-right.
|
|
|
|
method : str
|
|
Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.
|
|
|
|
layout : str
|
|
Layout of the input data. Default is 'NCHW'.
|
|
|
|
padding_mode : str
|
|
Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.
|
|
|
|
align_corners : bool
|
|
If True, the corner pixels of the input and output tensors are aligned.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The sampled output tensor with shape [N, C, H_out, W_out].
|
|
"""
|
|
return _ffi_api.grid_sample( # type: ignore
|
|
data,
|
|
grid,
|
|
method,
|
|
layout,
|
|
padding_mode,
|
|
align_corners,
|
|
)
|
|
|
|
|
|
def affine_grid(
|
|
data: Expr,
|
|
size: SizeLike,
|
|
align_corners: bool = True,
|
|
) -> Expr:
|
|
"""Generate a 2D or 3D sampling grid using an affine transformation matrix.
|
|
|
|
This operation is described in https://arxiv.org/pdf/1506.02025.pdf.
|
|
It generates a uniform sampling grid within the target shape, normalizes it
|
|
to [-1, 1], and applies the provided affine transformation.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input affine matrix tensor with shape [batch, 2, 3] for 2D or
|
|
[batch, 3, 4] for 3D.
|
|
|
|
size : SizeLike
|
|
The target output spatial shape, (H, W) for 2D or (D, H, W) for 3D. If a
|
|
single integer or PrimExpr is provided, it is interpreted as a square 2D
|
|
output shape (size, size).
|
|
|
|
align_corners : bool
|
|
If True, normalized grid coordinates map to corner pixels; if False, to
|
|
pixel centers (the PyTorch / ONNX default).
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The output grid tensor with shape [batch, 2, H, W] for 2D or
|
|
[batch, 3, D, H, W] for 3D.
|
|
"""
|
|
if isinstance(size, int) or is_prim_expr(size):
|
|
size = (size, size)
|
|
if isinstance(size, tuple | list):
|
|
size = ShapeExpr(size)
|
|
|
|
return cast(Expr, _ffi_api.affine_grid(data, size, align_corners))
|