import math
import numpy as np
from enum import Enum, unique
class GetNearestPixel:
def __init__(self, mode: str):
self.func = {
'round_prefer_floor': GetNearestPixel.prefer_floor_func,
'round_prefer_ceil': GetNearestPixel.prefer_ceil_func,
'floor': GetNearestPixel.floor_func,
'ceil': GetNearestPixel.ceil_func,
'simple': GetNearestPixel.simple_func
}[mode]
def __call__(self, x_original, is_downsample):
return self.func(x_original, is_downsample)
@staticmethod
def prefer_floor_func(x_original, is_downsample):
if x_original == int(x_original) + 0.5:
return int(math.floor(x_original))
else:
return int(round(x_original))
@staticmethod
def prefer_ceil_func(x_original, is_downsample):
return int(round(x_original))
@staticmethod
def floor_func(x_original, is_downsample):
return int(math.floor(x_original))
@staticmethod
def ceil_func(x_original, is_downsample):
return int(math.ceil(x_original))
@staticmethod
def simple_func(x_original, is_downsample):
if is_downsample:
return int(math.ceil(x_original))
else:
return int(x_original)
class GetOriginalCoordinate:
def __init__(self, mode: str):
self.func = {
'half_pixel': GetOriginalCoordinate.half_pixel_func,
'pytorch_half_pixel': GetOriginalCoordinate.pytorch_half_pixel_func,
'asymmetric': GetOriginalCoordinate.asymmetric_func,
'tf_half_pixel_for_nn': GetOriginalCoordinate.tf_half_pixel_for_nn_func,
'align_corners': GetOriginalCoordinate.align_corners_func
}[mode]
def __call__(self, x_resized, x_scale, length_resized, length_original):
return self.func(x_resized, x_scale, length_resized, length_original)
@staticmethod
def half_pixel_func(x_resized, x_scale, length_resized, length_original):
return ((x_resized + 0.5) / x_scale) - 0.5
@staticmethod
def pytorch_half_pixel_func(x_resized, x_scale, length_resized, length_original):
return (x_resized + 0.5) / x_scale - 0.5 if length_resized > 1 else 0.0
@staticmethod
def asymmetric_func(x_resized, x_scale, length_resized, length_original):
return x_resized / x_scale
@staticmethod
def tf_half_pixel_for_nn_func(x_resized, x_scale, length_resized, length_original):
return (x_resized + 0.5) / x_scale
@staticmethod
def align_corners_func(x_resized, x_scale, length_resized, length_original):
return 0 if length_resized == 1 else x_resized * (length_original - 1) / (length_resized - 1)
def get_cubic_coeff(s, a):
abs_s = abs(s)
coeff = np.zeros(4)
coeff[0] = a * (abs_s - 1.0) * (abs_s - 1.0) * abs_s
coeff[1] = ((a + 2.0) * abs_s - (a + 3.0)) * abs_s * abs_s + 1.0
coeff[2] = (((-a -2.0) * abs_s+ (2.0 * a + 3.0)) * abs_s - a) * abs_s
coeff[3] = - a * abs_s * abs_s * (abs_s - 1.0)
return coeff
def triangle_coeffs(dz):
return np.maximum(0.0, 1.0 - np.abs(dz))
@unique
class ShapeCalculationMode(Enum):
SIZES = 0
SCALES = 1
class InterpolateCalculation:
def __init__(self, attrs: dict):
self.mode = attrs['mode']
self.func = {
'nearest': self.nearest_interpolation,
'linear': self.linear_interpolation,
'cubic': self.cubic_interpolation,
'linear_onnx': self.onnx_linear_interpolation
}[self.mode]
self.attrs = attrs
self.pads_begin = attrs.get('pads_begin', [0])
self.pads_end = attrs.get('pads_end', [0])
self.coordinate_transformation_mode = attrs.get('coordinate_transformation_mode', 'half_pixel')
self.nearest_mode = attrs.get('nearest_mode', 'round_prefer_floor')
self.cube_coeff = attrs.get('cube_coeff', -0.75)
self.antialias = attrs.get('antialias', False)
self.shape_calculation_mode = {
'sizes': ShapeCalculationMode.SIZES,
'scales': ShapeCalculationMode.SCALES
}[attrs['shape_calculation_mode']]
self.get_original_coordinate = self.get_coordinate_transformation_mode()
self.get_nearest_pixel = GetNearestPixel(self.nearest_mode)
def get_coordinate_transformation_mode(self):
return GetOriginalCoordinate(self.coordinate_transformation_mode)
def shape_infer(self, input_data, sizes, scales):
result = input_data.shape + self.pads_begin + self.pads_end
if self.shape_calculation_mode == ShapeCalculationMode.SIZES:
for i, axis in enumerate(self.axes):
result[axis] = sizes[i]
else:
for i, axis in enumerate(self.axes):
result[axis] = math.floor(scales[i] * result[axis])
return result
@staticmethod
def correct_pad(pad, rank):
pad_len = len(pad)
if pad_len < rank:
return np.pad(pad, (0, rank - pad_len), 'constant').astype(np.int64)
elif pad_len > rank:
return np.array(pad[: rank - 1]).astype(np.int64)
else:
return np.array(pad, dtype=np.int64)
def __call__(self, input_data, sizes, scales, axes):
rank = input_data.ndim
self.pads_begin = InterpolateCalculation.correct_pad(self.pads_begin, rank)
self.pads_end = InterpolateCalculation.correct_pad(self.pads_end, rank)
self.pads = list(zip(self.pads_begin, self.pads_end))
self.axes = np.array(axes).astype(np.int64)
self.output_shape = self.shape_infer(input_data, sizes, scales)
padded_data = np.pad(input_data, self.pads, 'constant')
if self.shape_calculation_mode == ShapeCalculationMode.SIZES:
num_of_axes = len(self.axes)
self.scales = np.zeros(num_of_axes)
for i, axis in enumerate(axes):
self.scales[i] = self.output_shape[axis] / padded_data.shape[axis]
else:
self.scales = scales
if self.mode == 'nearest':
self.all_scales = np.ones(rank).astype(np.float)
for i, axis in enumerate(self.axes):
self.all_scales[axis] = self.scales[i]
self.input_shape = padded_data.shape
return self.func(padded_data)
def clip_coord(self, coord, axis):
return max(0, min(coord, self.input_shape[axis] - 1))
def cubic_interpolation(self, input_data):
rank = len(self.input_shape)
result = np.zeros(self.output_shape)
num_of_axes = len(self.axes)
indices = [ind for ind in np.ndindex(tuple(4 for _ in range(num_of_axes)))]
for coordinates in np.ndindex(tuple(self.output_shape)):
input_coords = np.array(coordinates, dtype=np.int64)
cubic_coeffs = np.zeros((rank, 4))
for i, axis in enumerate(self.axes):
in_coord = self.get_original_coordinate(coordinates[axis], self.scales[i], self.output_shape[axis], self.input_shape[axis])
in_coord_int = math.floor(in_coord)
input_coords[axis] = in_coord_int
cubic_coeffs[axis] = get_cubic_coeff(in_coord - in_coord_int, self.cube_coeff)
summa = 0.0
for index in indices:
coords_for_sum = input_coords.copy()
coeffs_prod = 1.0
for i, axis in enumerate(self.axes):
coords_for_sum[axis] = self.clip_coord(input_coords[axis] + index[i] - 1, axis)
for i, axis in enumerate(self.axes):
coeffs_prod = coeffs_prod * cubic_coeffs[axis][index[i]]
summa += coeffs_prod * input_data[tuple(coords_for_sum)]
result[coordinates] = summa
return result
def linear_interpolation(self, input_data):
result = np.zeros(self.output_shape)
num_of_axes = len(self.axes)
is_downsample = False
for scale in self.scales:
is_downsample = is_downsample or (scale < 1)
antialias = is_downsample and self.antialias
a = np.zeros(num_of_axes)
for i, _ in enumerate(self.axes):
a[i] = self.scales[i] if antialias else 1.0
prod_of_a = np.prod(a)
r = np.zeros(num_of_axes).astype(np.int64)
for i, _ in enumerate(self.axes):
r[i] = 2 if self.scales[i] > 1.0 else int(math.ceil(2.0/a[i]))
indices = [tuple(np.array(ind).astype(np.int64) - r) for ind in np.ndindex(tuple(2 * r + 1))]
for coordinates in np.ndindex(tuple(self.output_shape)):
icoords = np.array(coordinates).astype(np.float64)
icoords_r = np.array(coordinates).astype(np.float64)
for i, axis in enumerate(self.axes):
in_coord = self.get_original_coordinate(coordinates[axis], self.scales[i], self.output_shape[axis], self.input_shape[axis])
icoords[axis] = in_coord
icoords_r[axis] = round(in_coord)
summa = 0.0
wsum = 0.0
for index in indices:
inner_coords = np.array(coordinates)
for i, axis in enumerate(self.axes):
inner_coords[axis] = index[i] + icoords_r[axis]
conditions = [inner_coords[axis] >= 0 and inner_coords[axis] < self.input_shape[axis] for axis in self.axes]
if not all(conditions):
continue
dz = np.zeros(num_of_axes)
for i, axis in enumerate(self.axes):
dz[i] = icoords[axis] - inner_coords[axis]
w = prod_of_a * np.prod(triangle_coeffs(a * dz))
wsum += w
summa += w * input_data[tuple(inner_coords)]
if wsum == 0:
result[coordinates] = 0.0
else:
result[coordinates] = summa / wsum
return result
def onnx_linear_interpolation5D(self, input_data):
rank = len(self.input_shape)
assert rank in [3, 5], "mode 'linear_onnx' supports only 3D or 5D tensors"
assert set(self.axes) == {2, 3, 4} or set(self.axes) == {0, 1, 2}, \
"mode 'linear_onnx' supports only case when axes = {2, 3, 4} or axes = {0, 1, 2}"
result = np.zeros(self.output_shape)
if rank == 3:
reshaped_data = np.reshape(input_data, (1, 1, self.input_shape[0], self.input_shape[1], self.input_shape[2]))
result = np.reshape(result, (1, 1, self.output_shape[0], self.output_shape[1], self.output_shape[2]))
else:
reshaped_data = input_data
input_shape = np.array(reshaped_data.shape).astype(np.int64)
output_shape = np.array(result.shape).astype(np.int64)
batch_size = input_shape[0];
num_channels = input_shape[1];
input_depth = input_shape[2];
input_height = input_shape[3];
input_width = input_shape[4];
output_depth = output_shape[2];
output_height = output_shape[3];
output_width = output_shape[4];
depth_scale = self.scales[0];
height_scale = self.scales[1];
width_scale = self.scales[2];
z_original = np.zeros(output_depth).astype(np.float)
y_original = np.zeros(output_height).astype(np.float)
x_original = np.zeros(output_width).astype(np.float)
in_z1 = np.zeros(output_depth).astype(np.int64)
in_z2 = np.zeros(output_depth).astype(np.int64)
in_y1 = np.zeros(output_height).astype(np.int64)
in_y2 = np.zeros(output_height).astype(np.int64)
in_x1 = np.zeros(output_width).astype(np.int64)
in_x2 = np.zeros(output_width).astype(np.int64)
dz1 = np.zeros(output_depth).astype(np.float)
dz2 = np.zeros(output_depth).astype(np.float)
dy1 = np.zeros(output_height).astype(np.float)
dy2 = np.zeros(output_height).astype(np.float)
dx1 = np.zeros(output_width).astype(np.float)
dx2 = np.zeros(output_width).astype(np.float)
for z in range(0, output_depth):
in_z = self.get_original_coordinate(z, depth_scale, output_depth, input_depth)
z_original[z] = in_z
in_z = max(0, min(in_z, input_depth - 1))
in_z1[z] = max(0, min(int(in_z), input_depth - 1))
in_z2[z] = min(in_z1[z] + 1, input_depth - 1)
dz1[z] = abs(in_z - in_z1[z])
dz2[z] = abs(in_z - in_z2[z])
if in_z1[z] == in_z2[z]:
dz1[z] = 0.5
dz2[z] = 0.5
for y in range(0, output_height):
in_y = self.get_original_coordinate(y, height_scale, output_height, input_height)
y_original[y] = in_y
in_y = max(0, min(in_y, input_height - 1))
in_y1[y] = max(0, min(int(in_y), input_height - 1))
in_y2[y] = min(in_y1[y] + 1, input_height - 1)
dy1[y] = abs(in_y - in_y1[y])
dy2[y] = abs(in_y - in_y2[y])
if in_y1[y] == in_y2[y]:
dy1[y] = 0.5
dy2[y] = 0.5
for x in range(0, output_width):
in_x = self.get_original_coordinate(x, width_scale, output_width, input_width);
x_original[x] = in_x
in_x = max(0.0, min(in_x, input_width - 1));
in_x1[x] = min(in_x, input_width - 1);
in_x2[x] = min(in_x1[x] + 1, input_width - 1);
dx1[x] = abs(in_x - in_x1[x]);
dx2[x] = abs(in_x - in_x2[x]);
if in_x1[x] == in_x2[x]:
dx1[x] = 0.5
dx2[x] = 0.5
for n in range(0, batch_size):
for c in range(0, num_channels):
for z in range(0, output_depth):
for y in range(0, output_height):
for x in range(0, output_width):
x111 = reshaped_data[n, c, in_z1[z], in_y1[y], in_x1[x]]
x211 = reshaped_data[n, c, in_z1[z], in_y1[y], in_x2[x]]
x121 = reshaped_data[n, c, in_z1[z], in_y2[y], in_x1[x]]
x221 = reshaped_data[n, c, in_z1[z], in_y2[y], in_x2[x]]
x112 = reshaped_data[n, c, in_z2[z], in_y1[y], in_x1[x]]
x212 = reshaped_data[n, c, in_z2[z], in_y1[y], in_x2[x]]
x122 = reshaped_data[n, c, in_z2[z], in_y2[y], in_x1[x]]
x222 = reshaped_data[n, c, in_z2[z], in_y2[y], in_x2[x]]
temp = dx2[x] * dy2[y] * dz2[z] * x111 + dx1[x] * dy2[y] * dz2[z] * x211
temp += dx2[x] * dy1[y] * dz2[z] * x121 + dx1[x] * dy1[y] * dz2[z] * x221
temp += dx2[x] * dy2[y] * dz1[z] * x112 + dx1[x] * dy2[y] * dz1[z] * x212
temp += dx2[x] * dy1[y] * dz1[z] * x122 + dx1[x] * dy1[y] * dz1[z] * x222
result[n, c, z, y, x] = temp
return np.reshape(result, self.output_shape)
def onnx_linear_interpolation4D(self, input_data):
rank = len(self.input_shape)
assert rank in [2, 4], "mode 'linear_onnx' supports only 2D or 4D tensors"
assert set(self.axes) == {2, 3} or set(self.axes) == {0, 1}, \
"mode 'linear_onnx' supports only case when axes = {2, 3} or axes = {0, 1}"
result = np.zeros(self.output_shape)
if rank == 2:
reshaped_data = np.reshape(input_data, (1, 1, self.input_shape[0], self.input_shape[1]))
result = np.reshape(result, (1, 1, self.output_shape[0], self.output_shape[1]))
else:
reshaped_data = input_data
input_shape = np.array(reshaped_data.shape).astype(np.int64)
output_shape = np.array(result.shape).astype(np.int64)
output_height = output_shape[2]
output_width = output_shape[3]
input_height = input_shape[2]
input_width = input_shape[3]
height_scale = self.scales[0]
width_scale = self.scales[1]
batch_size = input_shape[0]
num_channels = input_shape[1]
y_original = np.zeros(output_height).astype(np.float)
x_original = np.zeros(output_width).astype(np.float)
in_y1 = np.zeros(output_height).astype(np.int64)
in_y2 = np.zeros(output_height).astype(np.int64)
in_x1 = np.zeros(output_width).astype(np.int64)
in_x2 = np.zeros(output_width).astype(np.int64)
dy1 = np.zeros(output_height).astype(np.float)
dy2 = np.zeros(output_height).astype(np.float)
dx1 = np.zeros(output_width).astype(np.float)
dx2 = np.zeros(output_width).astype(np.float)
for y in range(0, output_height):
in_y = self.get_original_coordinate(y, height_scale, output_height, input_height)
y_original[y] = in_y
in_y = max(0, min(in_y, input_height - 1))
in_y1[y] = max(0, min(int(in_y), input_height - 1))
in_y2[y] = min(in_y1[y] + 1, input_height - 1)
dy1[y] = abs(in_y - in_y1[y])
dy2[y] = abs(in_y - in_y2[y])
if in_y1[y] == in_y2[y]:
dy1[y] = 0.5
dy2[y] = 0.5
for x in range(0, output_width):
in_x = self.get_original_coordinate(x, width_scale, output_width, input_width);
x_original[x] = in_x
in_x = max(0.0, min(in_x, input_width - 1));
in_x1[x] = min(in_x, input_width - 1);
in_x2[x] = min(in_x1[x] + 1, input_width - 1);
dx1[x] = abs(in_x - in_x1[x]);
dx2[x] = abs(in_x - in_x2[x]);
if in_x1[x] == in_x2[x]:
dx1[x] = 0.5
dx2[x] = 0.5
for n in range(0, batch_size):
for c in range(0, num_channels):
for y in range(0, output_height):
for x in range(0, output_width):
x11 = reshaped_data[n, c, in_y1[y], in_x1[x]]
x21 = reshaped_data[n, c, in_y1[y], in_x2[x]]
x12 = reshaped_data[n, c, in_y2[y], in_x1[x]]
x22 = reshaped_data[n, c, in_y2[y], in_x2[x]]
temp = dx2[x] * dy2[y] * x11 + dx1[x] * dy2[y] * x21 + dx2[x] * dy1[y] * x12 + dx1[x] * dy1[y] * x22
result[n, c, y, x] = temp
return np.reshape(result, self.output_shape)
def onnx_linear_interpolation(self, input_data):
rank = len(self.input_shape)
assert rank in [2, 3, 4, 5], "mode 'linear_onnx' supports only 2D, 3D, 4D, or 5D tensors"
if rank in [2, 4]:
self.onnx_linear_interpolation4D(input_data)
else:
self.onnx_linear_interpolation5D(input_data)
def nearest_interpolation(self, input_data):
result = np.zeros(self.output_shape)
num_of_axes = len(self.axes)
for coordinates in np.ndindex(tuple(self.output_shape)):
input_coords = np.array(coordinates, dtype=np.int64)
for axis, scale in enumerate(self.all_scales):
in_coord = self.get_original_coordinate(coordinates[axis], scale, self.output_shape[axis], self.input_shape[axis])
nearest_pixel = self.get_nearest_pixel(in_coord, scale < 1)
input_coords[axis] = max(0, min(nearest_pixel, self.input_shape[axis] - 1))
result[coordinates] = input_data[tuple(input_coords)]
return result