python - 是否有将位移图/矩阵应用于图像的功能？

我生成一个位移图/大小为(H，W，2)的矩阵，每个通道对应于x和y中值的位移。目标是使用位移图的值“移动”形状为(H，W，3)的给定图像的像素值。

一个只有一个通道(例如x)的简单示例为:

disp_map = np.array([[ 0, 1, 1, 0 ],
                      [0, 0, 0, 0],
                      [1, 1, 0, 0]])

im = np.array([[ 0, 1, 1, 0 ],
              [0, 1, 1, 0],
              [0, 0, 0, 0]])

output = apply_disp_function(im, disp_map)
output
np.array([[ 1, 1, 0, 0 ],
          [0, 1, 1, 0],
          [0, 0, 0, 0]])

我不知道是否已经有一个函数可以执行此操作，或者我必须用 slice 在numpy中编写一个函数。

最佳答案

这是一种可能的解决方案。在此，disp_map数组在其最后一个维度中为每个轴都具有一个映射。超出图像边界的位移将环绕。

import numpy as np

def displace(im, disp_map):
    im = np.asarray(im)
    disp_map = np.asarray(disp_map)
    grid = np.ogrid[list(map(slice, disp_map.shape[:-1]))]
    result = np.zeros_like(im)
    np.add.at(result, tuple((g + disp_map[..., i]) % im.shape[i]
                            for i, g in enumerate(grid)), im)
    return result

# Each pair of values represents the number of rows and columns
# that each element will be displaced
disp_map = np.array([[[ 0,  0], [ 0, -1], [ 0,  1], [ 0,  0]],
                     [[ 0,  0], [ 1,  0], [ 1,  1], [ 0,  0]],
                     [[ 0,  1], [ 0,  1], [ 0,  0], [ 0,  0]]])
im = np.array([[ 0,  1,  1,  0],
               [ 0,  1,  1,  0],
               [ 0,  0,  0,  0]])
output = displace(im, disp_map)
print(output)
# [[1 0 0 1]
#  [0 0 0 0]
#  [0 1 0 1]]