wehub-resource-sync
308 lines
10 KiB
Python
308 lines
10 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, line-too-long, unused-variable, too-many-locals
|
|
# ruff: noqa: E741, F841, RUF005
|
|
"""Upsampling in python"""
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
from tvm.topi.utils import nchw_pack_layout
|
|
|
|
|
|
def get_inx(x, image_width, target_width, coordinate_transformation_mode):
|
|
"""Infer input x from output x with various coordinate transformation methods"""
|
|
scale = image_width / target_width
|
|
if coordinate_transformation_mode == "half_pixel":
|
|
in_x = (x + 0.5) * scale - 0.5
|
|
elif coordinate_transformation_mode == "align_corners":
|
|
in_x = (image_width - 1) / (target_width - 1) * x if target_width > 1 else 0
|
|
elif coordinate_transformation_mode == "asymmetric":
|
|
in_x = scale * x
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported coordinate_transformation_mode: {coordinate_transformation_mode}"
|
|
)
|
|
return in_x
|
|
|
|
|
|
def get_index(x, image_width, target_width, coordinate_transformation_mode, rounding_method=""):
|
|
"""get and round the nearest index for nearest_neighbor"""
|
|
in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode)
|
|
|
|
effective_rounding_method = rounding_method
|
|
if not effective_rounding_method:
|
|
if coordinate_transformation_mode == "align_corners":
|
|
effective_rounding_method = "round"
|
|
else:
|
|
effective_rounding_method = "floor"
|
|
|
|
if effective_rounding_method == "floor":
|
|
out = math.floor(in_x)
|
|
elif effective_rounding_method == "round":
|
|
out = round(in_x)
|
|
elif effective_rounding_method == "round_prefer_floor":
|
|
out = math.ceil(in_x - 0.5)
|
|
elif effective_rounding_method == "round_prefer_ceil":
|
|
out = math.floor(in_x + 0.5)
|
|
elif effective_rounding_method == "ceil":
|
|
out = math.ceil(in_x)
|
|
else:
|
|
raise ValueError(f"Unknown rounding method: {rounding_method!r}")
|
|
|
|
out = max(min(out, image_width - 1), 0)
|
|
return int(out)
|
|
|
|
|
|
def resize3d_nearest(arr, scale, coordinate_transformation_mode, rounding_method=""):
|
|
"""Populate the array by scale factor"""
|
|
d, h, w = arr.shape
|
|
out_d, out_h, out_w = [round(i * s) for i, s in zip(arr.shape, scale)]
|
|
out = np.empty((out_d, out_h, out_w))
|
|
for z in range(out_d):
|
|
for y in range(out_h):
|
|
for x in range(out_w):
|
|
in_z = get_index(z, d, out_d, coordinate_transformation_mode, rounding_method)
|
|
in_y = get_index(y, h, out_h, coordinate_transformation_mode, rounding_method)
|
|
in_x = get_index(x, w, out_w, coordinate_transformation_mode, rounding_method)
|
|
out[z, y, x] = arr[in_z, in_y, in_x]
|
|
return out
|
|
|
|
|
|
def resize3d_linear(data_in, scale, coordinate_transformation_mode):
|
|
"""Trilinear 3d scaling using python"""
|
|
dtype = data_in.dtype
|
|
d, h, w = data_in.shape
|
|
new_d, new_h, new_w = [round(i * s) for i, s in zip(data_in.shape, scale)]
|
|
data_out = np.ones((new_d, new_h, new_w))
|
|
|
|
indexes = np.mgrid[0:2, 0:2, 0:2]
|
|
|
|
def _get_patch(zint, yint, xint):
|
|
# Get the surrounding values
|
|
indices = indexes.copy()
|
|
indices[0] = np.maximum(np.minimum(indexes[0] + zint, d - 1), 0)
|
|
indices[1] = np.maximum(np.minimum(indexes[1] + yint, h - 1), 0)
|
|
indices[2] = np.maximum(np.minimum(indexes[2] + xint, w - 1), 0)
|
|
p = data_in[indices[0], indices[1], indices[2]]
|
|
return p
|
|
|
|
for m in range(new_d):
|
|
for j in range(new_h):
|
|
for k in range(new_w):
|
|
in_z = get_inx(m, d, new_d, coordinate_transformation_mode)
|
|
in_y = get_inx(j, h, new_h, coordinate_transformation_mode)
|
|
in_x = get_inx(k, w, new_w, coordinate_transformation_mode)
|
|
zint = math.floor(in_z)
|
|
zfract = in_z - math.floor(in_z)
|
|
|
|
yint = math.floor(in_y)
|
|
yfract = in_y - math.floor(in_y)
|
|
|
|
xint = math.floor(in_x)
|
|
xfract = in_x - math.floor(in_x)
|
|
|
|
wz = np.array([1.0 - zfract, zfract], dtype=dtype)
|
|
wy = np.array([1.0 - yfract, yfract], dtype=dtype)
|
|
wx = np.array([1.0 - xfract, xfract], dtype=dtype)
|
|
|
|
p = _get_patch(zint, yint, xint)
|
|
l = np.sum(p * wx, axis=-1)
|
|
col = np.sum(l * wy, axis=-1)
|
|
data_out[m, j, k] = np.sum(col * wz)
|
|
|
|
return data_out
|
|
|
|
|
|
def resize3d_cubic(data_in, scale, coordinate_transformation_mode):
|
|
"""Tricubic 3d scaling using python"""
|
|
dtype = data_in.dtype
|
|
d, h, w = data_in.shape
|
|
new_d, new_h, new_w = [round(i * s) for i, s in zip(data_in.shape, scale)]
|
|
data_out = np.ones((new_d, new_h, new_w))
|
|
|
|
def _cubic_spline_weights(t, alpha=-0.5):
|
|
"""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 np.array([w1, w2, w3, w4])
|
|
|
|
indexes = np.mgrid[-1:3, -1:3, -1:3]
|
|
|
|
def _get_patch(zint, yint, xint):
|
|
# Get the surrounding values
|
|
indices = indexes.copy()
|
|
indices[0] = np.maximum(np.minimum(indexes[0] + zint, d - 1), 0)
|
|
indices[1] = np.maximum(np.minimum(indexes[1] + yint, h - 1), 0)
|
|
indices[2] = np.maximum(np.minimum(indexes[2] + xint, w - 1), 0)
|
|
p = data_in[indices[0], indices[1], indices[2]]
|
|
return p
|
|
|
|
for m in range(new_d):
|
|
for j in range(new_h):
|
|
for k in range(new_w):
|
|
in_z = get_inx(m, d, new_d, coordinate_transformation_mode)
|
|
in_y = get_inx(j, h, new_h, coordinate_transformation_mode)
|
|
in_x = get_inx(k, w, new_w, coordinate_transformation_mode)
|
|
zint = math.floor(in_z)
|
|
zfract = in_z - math.floor(in_z)
|
|
|
|
yint = math.floor(in_y)
|
|
yfract = in_y - math.floor(in_y)
|
|
|
|
xint = math.floor(in_x)
|
|
xfract = in_x - math.floor(in_x)
|
|
|
|
wz = _cubic_spline_weights(zfract)
|
|
wy = _cubic_spline_weights(yfract)
|
|
wx = _cubic_spline_weights(xfract)
|
|
|
|
p = _get_patch(zint, yint, xint)
|
|
|
|
l = np.sum(p * wx, axis=-1)
|
|
col = np.sum(l * wy, axis=-1)
|
|
data_out[m, j, k] = np.sum(col * wz)
|
|
|
|
return data_out
|
|
|
|
|
|
def resize3d_ncdhw(
|
|
data,
|
|
scale,
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="align_corners",
|
|
rounding_method="",
|
|
):
|
|
"""reference kernel for 3D image resizing"""
|
|
ishape = data.shape
|
|
|
|
oshape = (
|
|
ishape[0],
|
|
ishape[1],
|
|
round(ishape[2] * scale[0]),
|
|
round(ishape[3] * scale[1]),
|
|
round(ishape[4] * scale[2]),
|
|
)
|
|
|
|
output_np = np.zeros(oshape, dtype=data.dtype)
|
|
|
|
for b in range(oshape[0]):
|
|
for c in range(oshape[1]):
|
|
if method == "nearest_neighbor":
|
|
output_np[b, c, :, :, :] = resize3d_nearest(
|
|
data[b, c, :, :, :], scale, coordinate_transformation_mode, rounding_method
|
|
)
|
|
elif method == "linear":
|
|
output_np[b, c, :, :, :] = resize3d_linear(
|
|
data[b, c, :, :, :], scale, coordinate_transformation_mode
|
|
)
|
|
elif method == "cubic":
|
|
output_np[b, c, :, :, :] = resize3d_cubic(
|
|
data[b, c, :, :, :], scale, coordinate_transformation_mode
|
|
)
|
|
else:
|
|
raise ValueError("Unknown resize method", method)
|
|
|
|
return output_np
|
|
|
|
|
|
def resize1d_python(
|
|
data,
|
|
scale,
|
|
layout="NCW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="align_corners",
|
|
rounding_method="",
|
|
):
|
|
"""Python version of 3D scaling using nearest neighbour"""
|
|
|
|
if layout == "NWC":
|
|
data = data.transpose([0, 2, 1])
|
|
|
|
data = np.expand_dims(data, axis=[2, 3])
|
|
output_np = resize3d_ncdhw(
|
|
data, (1, 1) + scale, method, coordinate_transformation_mode, rounding_method
|
|
)
|
|
output_np = np.squeeze(output_np, axis=2)
|
|
output_np = np.squeeze(output_np, axis=2)
|
|
|
|
if layout == "NWC":
|
|
output_np = output_np.transpose([0, 2, 1])
|
|
|
|
return output_np
|
|
|
|
|
|
def resize2d_python(
|
|
data,
|
|
scale,
|
|
layout="NCHW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="align_corners",
|
|
rounding_method="",
|
|
):
|
|
"""Python version of scaling using nearest neighbour"""
|
|
|
|
if layout == "NHWC":
|
|
data = data.transpose([0, 3, 1, 2])
|
|
elif nchw_pack_layout(layout):
|
|
ishape = data.shape
|
|
transposed = data.transpose([0, 4, 1, 5, 2, 3])
|
|
tshape = transposed.shape
|
|
data = transposed.reshape(
|
|
tshape[0] * tshape[1], tshape[2] * tshape[3], tshape[4], tshape[5]
|
|
)
|
|
|
|
data = np.expand_dims(data, axis=2)
|
|
output_np = resize3d_ncdhw(
|
|
data, (1,) + scale, method, coordinate_transformation_mode, rounding_method
|
|
)
|
|
output_np = np.squeeze(output_np, axis=2)
|
|
|
|
if layout == "NHWC":
|
|
output_np = output_np.transpose([0, 2, 3, 1])
|
|
elif nchw_pack_layout(layout):
|
|
output_np = output_np.reshape(tshape[0:4] + output_np.shape[2:])
|
|
output_np = output_np.transpose([0, 2, 4, 5, 1, 3])
|
|
|
|
return output_np
|
|
|
|
|
|
def resize3d_python(
|
|
data,
|
|
scale,
|
|
layout="NCDHW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="align_corners",
|
|
rounding_method="",
|
|
):
|
|
"""Python version of 3D scaling using nearest neighbour"""
|
|
|
|
if layout == "NDHWC":
|
|
data = data.transpose([0, 4, 1, 2, 3])
|
|
|
|
output_np = resize3d_ncdhw(data, scale, method, coordinate_transformation_mode, rounding_method)
|
|
|
|
if layout == "NDHWC":
|
|
output_np = output_np.transpose([0, 2, 3, 4, 1])
|
|
|
|
return output_np
|