Linear and cubic interpolation

interpolation.py

import numpy as np

from typing import List, Callable, Union, Optional


def cartesian(arrays, out=None):
    # type: (List[np.ndarray], np.ndarray) -> np.ndarray
    """
    From https://stackoverflow.com/a/1235363

    Generate a cartesian product of input arrays.

    Parameters
    ----------
    arrays : list of array-like
        1-D arrays to form the cartesian product of.
    out : ndarray
        Array to place the cartesian product in.

    Returns
    -------
    out : ndarray
        2-D array of shape (M, len(arrays)) containing cartesian products
        formed of input arrays.

    Examples
    --------
    >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
    array([[1, 4, 6],
           [1, 4, 7],
           [1, 5, 6],
           [1, 5, 7],
           [2, 4, 6],
           [2, 4, 7],
           [2, 5, 6],
           [2, 5, 7],
           [3, 4, 6],
           [3, 4, 7],
           [3, 5, 6],
           [3, 5, 7]])

    """

    arrays = [np.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

    n = np.prod([x.size for x in arrays])
    if out is None:
        out = np.zeros([n, len(arrays)], dtype=dtype)

    m = n // arrays[0].size
    out[:, 0] = np.repeat(arrays[0], m)
    if arrays[1:]:
        cartesian(arrays[1:], out=out[0:m, 1:])
        for j in range(1, arrays[0].size):
            out[j * m:(j + 1) * m, 1:] = out[0:m, 1:]
    return out


def interpolate_1d_with_x(data,                  # type: np.ndarray
                          scale_factor,          # type: float
                          x,                     # type: float
                          get_coeffs,            # type: Callable[[float], np.ndarray]
                          align_corners=False,   # type: bool
                          exclude_outside=False  # type: bool
                          ):                     # type: (...) -> np.ndarray
    def get_neighbor_idxes(x, n, limit):  # type: (float, int, int) -> np.ndarray
        """
        Return the n nearest indexes, prefer the indexes smaller than x to be compatible with nearest interpolation.
        As a result, the ratio must be in (0, 1]
        Examples:
        get_neighbor_idxes(4, 2, 10) == [3, 4]
        get_neighbor_idxes(4, 3, 10) == [3, 4, 5]
        get_neighbor_idxes(4.4, 3, 10) == [3, 4, 5]
        get_neighbor_idxes(4.5, 3, 10) == [3, 4, 5]
        get_neighbor_idxes(4.6, 3, 10) == [4, 5, 6]
        get_neighbor_idxes(4.4, 1, 10) == [4]
        get_neighbor_idxes(4.6, 1, 10) == [5]
        :param x:
        :param n: the number of the wanted indexes
        :param limit: the maximum value of index
        :return: An np.array containing n nearest indexes in ascending order
        """
        idxes = sorted(range(limit), key=lambda idx: (abs(x - idx), idx))[:n]
        idxes = sorted(idxes)
        return np.array(idxes)

    def get_neighbor(x, n, data):  # type: (float, int, np.ndarray) -> np.ndarray
        """
        Pad `data` in 'edge' mode, and get n nearest elements in the padded array and their indexes in the original
        array
        :param x:
        :param n:  the number of the wanted elements
        :param data: the array
        :return: A tuple containing the indexes of neighbor elements (the index can be smaller than 0 or higher than
        len(data)) and the value of these elements
        """
        pad_width = np.ceil(n / 2).astype(np.int)
        padded = np.pad(data, pad_width, mode='edge')
        x += pad_width

        idxes = get_neighbor_idxes(x, n, len(padded))
        ret = padded[idxes]
        return idxes - pad_width, ret

    input_width = len(data)
    output_width = scale_factor * input_width
    if align_corners:
        if output_width == 1:
            x_ori = 0.
        else:
            x_ori = x * (input_width - 1) / (output_width - 1)
    else:
        x_ori = (x + 0.5) / scale_factor - 0.5
    x_ori_int = np.floor(x_ori).astype(np.int).item()

    # ratio must be in (0, 1] since we prefer the pixel on the left of `x_ori` (as required by nearest interpolation)
    if x_ori.is_integer():
        ratio = 1
    else:
        ratio = x_ori - x_ori_int

    coeffs = get_coeffs(ratio)
    n = len(coeffs)

    idxes, points = get_neighbor(x_ori, n, data)

    if exclude_outside:
        for i, idx in enumerate(idxes):
            if idx < 0 or idx >= input_width:
                coeffs[i] = 0
        coeffs /= sum(coeffs)

    return np.dot(coeffs, points).item()


def interpolate_nd_with_x(data,                  # type: np.ndarray
                          n,                     # type: int
                          scale_factors,         # type: List[float]
                          x,                     # type: List[float]
                          get_coeffs,            # type: Callable[[float], np.ndarray]
                          align_corners=False,   # type: bool
                          exclude_outside=False  # type: bool
                          ):                     # type: (...) -> np.ndarray
    if n == 1:
        return interpolate_1d_with_x(data, scale_factors[0], x[0], get_coeffs, align_corners, exclude_outside)
    return interpolate_1d_with_x(
        [interpolate_nd_with_x(data[i], n - 1, scale_factors[1:], x[1:], get_coeffs, align_corners, exclude_outside)
         for i in range(data.shape[0])], scale_factors[0], x[0], get_coeffs, align_corners, exclude_outside)


def interpolate_nd(data,                  # type: np.ndarray
                   get_coeffs,            # type: Callable[[float], np.ndarray]
                   output_size=None,      # type: Optional[List[int]]
                   scale_factors=None,    # type: Optional[List[float]]
                   align_corners=False,   # type: bool
                   exclude_outside=False  # type: bool
                   ):                     # type: (...) -> np.ndarray
    def get_all_coords(data):   # type: (np.ndarray) -> np.ndarray
        return cartesian([list(range(data.shape[i])) for i in range(len(data.shape))])

    assert output_size is not None or scale_factors is not None
    if output_size is not None:
        scale_factors = np.array(output_size) / np.array(data.shape)
    else:
        output_size = (scale_factors * np.array(data.shape)).astype(np.int)
    assert scale_factors is not None

    ret = np.zeros(output_size)
    for x in get_all_coords(ret):
        ret[tuple(x)] = interpolate_nd_with_x(data, len(data.shape), scale_factors, x, get_coeffs, align_corners,
                                              exclude_outside)
    return ret


def cubic_coeffs(ratio, A=-0.75):   # type: (float, float) -> np.ndarray
    coeffs = [((A * (ratio + 1) - 5 * A) * (ratio + 1) + 8 * A) * (ratio + 1) - 4 * A,
              ((A + 2) * ratio - (A + 3)) * ratio * ratio + 1,
              ((A + 2) * (1 - ratio) - (A + 3)) * (1 - ratio) * (1 - ratio) + 1,
              ((A * ((1 - ratio) + 1) - 5 * A) * ((1 - ratio) + 1) + 8 * A) * ((1 - ratio) + 1) - 4 * A]

    return np.array(coeffs)


def linear_coeffs(ratio):           # type: (float) -> np.ndarray
    return np.array([1 - ratio, ratio])


def nearest_coeffs(ratio):          # type: (float) -> np.ndarray
    return np.array([1])



"""
Test our functions
"""

import cv2 as cv

a = np.array([[[1], [2]], [[3], [4]]]).astype(np.float)
print("OpenCV:")
print(cv.resize(a, None, fx=2, fy=2, interpolation=cv.INTER_CUBIC))

import torch
import torch.nn.functional as F

a = torch.tensor([[[[1, 2], [3, 4]]]], dtype=torch.float)
print("PyTorch, align_corners=True:")
print(F.interpolate(a, mode='bicubic', scale_factor=2, align_corners=True))
print("PyTorch, align_corners=False:")
print(F.interpolate(a, mode='bicubic', scale_factor=2, align_corners=False))

a = np.array([[1, 2], [3, 4]])
print("our implementation, align_corners=True:")
output = interpolate_nd(a, cubic_coeffs, scale_factors=[2, 2], align_corners=True)
print("our implementation, align_corners=False:")
output = interpolate_nd(a, cubic_coeffs, scale_factors=[2, 2], align_corners=False)