import warnings from math import ceil from . import interp_methods class NoneClass: pass try: import torch from torch import nn nnModuleWrapped = nn.Module except ImportError: warnings.warn('No PyTorch found, will work only with Numpy') torch = None nnModuleWrapped = NoneClass try: import numpy except ImportError: warnings.warn('No Numpy found, will work only with PyTorch') numpy = None if numpy is None and torch is None: raise ImportError("Must have either Numpy or PyTorch but both not found") def resize(input, scale_factors=None, out_shape=None, interp_method=interp_methods.cubic, support_sz=None, antialiasing=True): # get properties of the input tensor in_shape, n_dims = input.shape, input.ndim # fw stands for framework that can be either numpy or torch, # determined by the input type fw = numpy if type(input) is numpy.ndarray else torch eps = fw.finfo(fw.float32).eps # set missing scale factors or output shapem one according to another, # scream if both missing scale_factors, out_shape = set_scale_and_out_sz(in_shape, out_shape, scale_factors, fw) # sort indices of dimensions according to scale of each dimension. # since we are going dim by dim this is efficient sorted_filtered_dims_and_scales = [(dim, scale_factors[dim]) for dim in sorted(range(n_dims), key=lambda ind: scale_factors[ind]) if scale_factors[dim] != 1.] # unless support size is specified by the user, it is an attribute # of the interpolation method if support_sz is None: support_sz = interp_method.support_sz # when using pytorch, we need to know what is the input tensor device device = input.device if fw is torch else None # output begins identical to input and changes with each iteration output = input # iterate over dims for dim, scale_factor in sorted_filtered_dims_and_scales: # get 1d set of weights and fields of view for each output location # along this dim field_of_view, weights = prepare_weights_and_field_of_view_1d( dim, scale_factor, in_shape[dim], out_shape[dim], interp_method, support_sz, antialiasing, fw, eps, device) # multiply the weights by the values in the field of view and # aggreagate output = apply_weights(output, field_of_view, weights, dim, n_dims, fw) return output class ResizeLayer(nnModuleWrapped): def __init__(self, in_shape, scale_factors=None, out_shape=None, interp_method=interp_methods.cubic, support_sz=None, antialiasing=True): super(ResizeLayer, self).__init__() # fw stands for framework, that can be either numpy or torch. since # this is a torch layer, only one option in this case. fw = torch eps = fw.finfo(fw.float32).eps # set missing scale factors or output shapem one according to another, # scream if both missing scale_factors, out_shape = set_scale_and_out_sz(in_shape, out_shape, scale_factors, fw) # unless support size is specified by the user, it is an attribute # of the interpolation method if support_sz is None: support_sz = interp_method.support_sz self.n_dims = len(in_shape) # sort indices of dimensions according to scale of each dimension. # since we are going dim by dim this is efficient self.sorted_filtered_dims_and_scales = [(dim, scale_factors[dim]) for dim in sorted(range(self.n_dims), key=lambda ind: scale_factors[ind]) if scale_factors[dim] != 1.] # iterate over dims field_of_view_list = [] weights_list = [] for dim, scale_factor in self.sorted_filtered_dims_and_scales: # get 1d set of weights and fields of view for each output # location along this dim field_of_view, weights = prepare_weights_and_field_of_view_1d( dim, scale_factor, in_shape[dim], out_shape[dim], interp_method, support_sz, antialiasing, fw, eps, input.device) # keep weights and fields of views for all dims weights_list.append(nn.Parameter(weights, requires_grad=False)) field_of_view_list.append(nn.Parameter(field_of_view, requires_grad=False)) self.field_of_view = nn.ParameterList(field_of_view_list) self.weights = nn.ParameterList(weights_list) self.in_shape = in_shape def forward(self, input): # output begins identical to input and changes with each iteration output = input for (dim, scale_factor), field_of_view, weights in zip( self.sorted_filtered_dims_and_scales, self.field_of_view, self.weights): # multiply the weights by the values in the field of view and # aggreagate output = apply_weights(output, field_of_view, weights, dim, self.n_dims, torch) return output def prepare_weights_and_field_of_view_1d(dim, scale_factor, in_sz, out_sz, interp_method, support_sz, antialiasing, fw, eps, device=None): # If antialiasing is taking place, we modify the window size and the # interpolation method (see inside function) interp_method, cur_support_sz = apply_antialiasing_if_needed( interp_method, support_sz, scale_factor, antialiasing) # STEP 1- PROJECTED GRID: The non-integer locations of the projection of # output pixel locations to the input tensor projected_grid = get_projected_grid(in_sz, out_sz, scale_factor, fw, device) # STEP 2- FIELDS OF VIEW: for each output pixels, map the input pixels # that influence it field_of_view = get_field_of_view(projected_grid, cur_support_sz, in_sz, fw, eps, device) # STEP 3- CALCULATE WEIGHTS: Match a set of weights to the pixels in the # field of view for each output pixel weights = get_weights(interp_method, projected_grid, field_of_view) return field_of_view, weights def apply_weights(input, field_of_view, weights, dim, n_dims, fw): # STEP 4- APPLY WEIGHTS: Each output pixel is calculated by multiplying # its set of weights with the pixel values in its field of view. # We now multiply the fields of view with their matching weights. # We do this by tensor multiplication and broadcasting. # this step is separated to a different function, so that it can be # repeated with the same calculated weights and fields. # for this operations we assume the resized dim is the first one. # so we transpose and will transpose back after multiplying tmp_input = fw_swapaxes(input, dim, 0, fw) # field_of_view is a tensor of order 2: for each output (1d location # along cur dim)- a list of 1d neighbors locations. # note that this whole operations is applied to each dim separately, # this is why it is all in 1d. # neighbors = tmp_input[field_of_view] is a tensor of order image_dims+1: # for each output pixel (this time indicated in all dims), these are the # values of the neighbors in the 1d field of view. note that we only # consider neighbors along the current dim, but such set exists for every # multi-dim location, hence the final tensor order is image_dims+1. neighbors = tmp_input[field_of_view] # weights is an order 2 tensor: for each output location along 1d- a list # of weighs matching the field of view. we augment it with ones, for # broadcasting, so that when multiplies some tensor the weights affect # only its first dim. tmp_weights = fw.reshape(weights, (*weights.shape, * [1] * (n_dims - 1))) # now we simply multiply the weights with the neighbors, and then sum # along the field of view, to get a single value per out pixel tmp_output = (neighbors * tmp_weights).sum(1) # we transpose back the resized dim to its original position return fw_swapaxes(tmp_output, 0, dim, fw) def set_scale_and_out_sz(in_shape, out_shape, scale_factors, fw): # eventually we must have both scale-factors and out-sizes for all in/out # dims. however, we support many possible partial arguments if scale_factors is None and out_shape is None: raise ValueError("either scale_factors or out_shape should be " "provided") if out_shape is not None: # if out_shape has less dims than in_shape, we defaultly resize the # first dims for numpy and last dims for torch # out_shape = (list(out_shape) + list(in_shape[:-len(out_shape)]) # if fw is numpy # else list(in_shape[:-len(out_shape)]) + list(out_shape)) out_shape = (list(out_shape) + list(in_shape[-len(out_shape):]) if fw is numpy else list(in_shape[:-len(out_shape)]) + list(out_shape)) if scale_factors is None: # if no scale given, we calculate it as the out to in ratio # (not recomended) scale_factors = [out_sz / in_sz for out_sz, in_sz in zip(out_shape, in_shape)] if scale_factors is not None: # by default, if a single number is given as scale, we assume resizing # two dims (most common are images with 2 spatial dims) scale_factors = (scale_factors if isinstance(scale_factors, (list, tuple)) else [scale_factors, scale_factors]) # if less scale_factors than in_shape dims, we defaultly resize the # first dims for numpy and last dims for torch scale_factors = (list(scale_factors) + [1] * (len(in_shape) - len(scale_factors)) if fw is numpy else [1] * (len(in_shape) - len(scale_factors)) + list(scale_factors)) if out_shape is None: # when no out_shape given, it is calculated by multiplying the # scale by the in_shape (not recomended) out_shape = [ceil(scale_factor * in_sz) for scale_factor, in_sz in zip(scale_factors, in_shape)] # next line intentionally after out_shape determined for stability scale_factors = [float(sf) for sf in scale_factors] return scale_factors, out_shape def get_projected_grid(in_sz, out_sz, scale_factor, fw, device=None): # we start by having the ouput coordinates which are just integer locations out_coordinates = fw.arange(out_sz) # if using torch we need to match the grid tensor device to the input device out_coordinates = fw_set_device(out_coordinates, device, fw) # This is projecting the ouput pixel locations in 1d to the input tensor, # as non-integer locations. # the following fomrula is derived in the paper # "From Discrete to Continuous Convolutions" by Shocher et al. return (out_coordinates / scale_factor + (in_sz - 1) / 2 - (out_sz - 1) / (2 * scale_factor)) def get_field_of_view(projected_grid, cur_support_sz, in_sz, fw, eps, device): # for each output pixel, map which input pixels influence it, in 1d. # we start by calculating the leftmost neighbor, using half of the window # size (eps is for when boundary is exact int) left_boundaries = fw_ceil(projected_grid - cur_support_sz / 2 - eps, fw) # then we simply take all the pixel centers in the field by counting # window size pixels from the left boundary ordinal_numbers = fw.arange(ceil(cur_support_sz - eps)) # in case using torch we need to match the device ordinal_numbers = fw_set_device(ordinal_numbers, device, fw) field_of_view = left_boundaries[:, None] + ordinal_numbers # next we do a trick instead of padding, we map the field of view so that # it would be like mirror padding, without actually padding # (which would require enlarging the input tensor) mirror = fw_cat((fw.arange(in_sz), fw.arange(in_sz - 1, -1, step=-1)), fw) field_of_view = mirror[fw.remainder(field_of_view, mirror.shape[0])] field_of_view = fw_set_device(field_of_view, device, fw) return field_of_view def get_weights(interp_method, projected_grid, field_of_view): # the set of weights per each output pixels is the result of the chosen # interpolation method applied to the distances between projected grid # locations and the pixel-centers in the field of view (distances are # directed, can be positive or negative) weights = interp_method(projected_grid[:, None] - field_of_view) # we now carefully normalize the weights to sum to 1 per each output pixel sum_weights = weights.sum(1, keepdims=True) sum_weights[sum_weights == 0] = 1 return weights / sum_weights def apply_antialiasing_if_needed(interp_method, support_sz, scale_factor, antialiasing): # antialiasing is "stretching" the field of view according to the scale # factor (only for downscaling). this is low-pass filtering. this # requires modifying both the interpolation (stretching the 1d # function and multiplying by the scale-factor) and the window size. if scale_factor >= 1.0 or not antialiasing: return interp_method, support_sz cur_interp_method = (lambda arg: scale_factor * interp_method(scale_factor * arg)) cur_support_sz = support_sz / scale_factor return cur_interp_method, cur_support_sz def fw_ceil(x, fw): if fw is numpy: return fw.int_(fw.ceil(x)) else: return x.ceil().long() def fw_cat(x, fw): if fw is numpy: return fw.concatenate(x) else: return fw.cat(x) def fw_swapaxes(x, ax_1, ax_2, fw): if fw is numpy: return fw.swapaxes(x, ax_1, ax_2) else: return x.transpose(ax_1, ax_2) def fw_set_device(x, device, fw): if fw is numpy: return x else: return x.to(device)