python - 如何在赋值期间以向量化方式移动输出变量(在 numpy 中)的索引

Question

动机:假设我有一个 RGB 图像 J，并且我想对 J 的像素应用变换 T(如旋转)。我将创建一个新的黑色图像 K，其像素为通过 K[x,y]=J[T[x,y]] 与 J 相关。现在的问题是 T[x,y] 必须在 J 内部，如果我想完全捕获 J 的变换图像，我可能必须处理 x 或 y 的一些负值或大于大小的值因此，首先我必须确定 K 的大小，然后将 K 的像素移动适当的向量以避免负值。

现在，假设我已经确定了适当的平移向量。我想做一个坐标转换，将 (x,y) 发送到 (x+a, y+k)。

目标:使用for循环，我想做的是:

for i in range(0,J.shape[0]):
    for j in range(0, J.shape[1]):
        K[i+a,j+b] = J[T[i,j]]

我怎样才能以矢量化的方式做到这一点？如有任何帮助，我们将不胜感激。

<小时/>

编辑:

img = face() # dummy RGB data

i,j = np.mgrid[:img.shape[0], :img.shape[1]] # 2d arrays each
i_min, i_max, j_min, j_max = func(*) # assume that these values have been found
i = i + i_min
j = j + j_min
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)
inew,jnew = np.linalg.inv(T) @ [i.ravel(), j.ravel()] # 1d arrays each

inew = np.floor(inew).astype(int)
jnew = np.floor(jnew).astype(int)

out = np.zeros((i_max - i_min, j_max - j_min, 3), dtype=img.dtype)

for i in inew:
    for j in jnew:
        out[i-i_min,j-j_min, :] = img[i,j,:]

现在我想取消数组中 i_min 和 j_min 移位的效果，就像我使用 for 循环编写的代码一样。

Answer 1

朴素版本

据我了解你的问题:你有一个输入图像，你转换它的像素位置，并希望将结果放入一个可以容纳它的更大的数组中。我的做法如下:

import numpy as np
import matplotlib.pyplot as plt # for plotting the result
from scipy.misc import face # for dummy data
img = face() # dummy RGB data

# transform pixels by 45 degrees
i,j = np.mgrid[:img.shape[0], :img.shape[1]] # 2d arrays each
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)
inew,jnew = T @ [i.ravel(), j.ravel()] # 1d arrays each

# new coordinates now range into negatives, shift back into positives
# and the non-integer pixel indices will be normalized with floor
inew = np.floor(inew - inew.min()).astype(int)
jnew = np.floor(jnew - jnew.min()).astype(int)

# now the new coordinates are all non-negative, this defines the size of the output
out = np.zeros((inew.max() + 1, jnew.max() + 1, 3), dtype=img.dtype)

# fill the necessary indices of out with pixels from img
# reshape the indices to 2d for matching broadcast
inew = inew.reshape(img.shape[:-1])
jnew = jnew.reshape(img.shape[:-1])
out[inew, jnew, :] = img
# OR, alternative with 1d index arrays:
#out[inew, jnew, :] = img.reshape(-1, 3)

# check what we've done
plt.imshow(out)
plt.show()

代码的要点是旋转的像素坐标移回正值(这对应于您的[i+a, j+b] 移位)，分配一个新的零数组这将适合所有新索引，并且索引仅应用于右侧!这与您的代码不匹配，但我相信这是您真正想要做的:对于原始(未索引)图像中的每个像素，我们将其 RGB 值设置在 new 位置结果数组。

正如您所看到的，由于非整数转换后的坐标使用floor进行舍入，因此图像中存在大量黑色像素。这不太好，所以如果我们追求这条路径，我们应该执行二维插值以消除这些伪影。请注意，这需要相当多的内存和 CPU 时间:

import numpy as np
import scipy.interpolate as interp
import matplotlib.pyplot as plt # for plotting the result
from scipy.misc import face # for dummy data
img = face() # dummy RGB data

# transform pixels by 45 degrees
i,j = np.mgrid[:img.shape[0], :img.shape[1]] # 2d arrays each
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)
inew,jnew = T @ [i.ravel(), j.ravel()] # 1d arrays each

# new coordinates now range into negatives, shift back into positives
# keep them non-integer for interpolation later
inew -= inew.min()
jnew -= jnew.min()
# (inew, jnew, img) contain the data from which the output should be interpolated


# now the new coordinates are all non-negative, this defines the size of the output
out = np.zeros((int(round(inew.max())) + 1, int(round(jnew.max())) + 1, 3), dtype=img.dtype)
i_interp,j_interp = np.mgrid[:out.shape[0], :out.shape[1]]

# interpolate for each channel
for channel in range(3):
    out[..., channel] = interp.griddata(np.array([inew.ravel(), jnew.ravel()]).T, img[..., channel].ravel(), (i_interp, j_interp), fill_value=0)

# check what we've done
plt.imshow(out)
plt.show()

至少结果看起来好多了:

scipy.ndimage: map 坐标

直接符合您想法的方法可以使用 scipy.ndimage.map_coordinates使用逆变换执行插值。这应该比之前使用 griddata 的尝试具有更好的性能，因为 map_coordinates 可以利用输入数据在网格上定义的事实。事实证明，它确实使用了更少的内存和更少的 CPU:

import numpy as np
import scipy.ndimage as ndi
import matplotlib.pyplot as plt # for plotting the result
from scipy.misc import face # for dummy data

img = face() # dummy RGB data
n,m = img.shape[:-1]

# transform pixels by 45 degrees
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)

# find out the extent of the transformed pixels from the four corners
inew_tmp,jnew_tmp = T @ [[0, 0, n-1, n-1], [0, m-1, 0, m-1]] # 1d arrays each
imin,imax,jmin,jmax = inew_tmp.min(),inew_tmp.max(),jnew_tmp.min(),jnew_tmp.max()
imin,imax,jmin,jmax = (int(round(val)) for val in (imin,imax,jmin,jmax))

# so the pixels of the original map inside [imin, imax] x [jmin, jmax]
# we need an image of size (imax - imin + 1, jmax - jmin + 1) to house this
out = np.zeros((imax - imin + 1, jmax - jmin + 1, 3), dtype=img.dtype)
# indices have to be shifted by [imin, imax]

# compute the corresponding (non-integer) coordinates on the domain for interpolation
inew,jnew = np.mgrid[:out.shape[0], :out.shape[1]]
i_back,j_back = np.linalg.inv(T) @ [inew.ravel() + imin, jnew.ravel() + jmin]

# perform 2d interpolation for each colour channel separately
for channel in range(3):
    out[inew, jnew, channel] = ndi.map_coordinates(img[..., channel], [i_back, j_back]).reshape(inew.shape)

# check what we've done
plt.imshow(out)
plt.show()

结果还是不错的:

scipy.ndimage:几何变换

最后，我意识到我们可以更高一级并使用 scipy.ndimage.geometric_transform直接地。对于旋转的浣熊情况，这似乎比使用 map_coordinates 的手动版本慢，但会导致更清晰的代码:

import numpy as np
import scipy.ndimage as ndi
import matplotlib.pyplot as plt # for plotting the result
from scipy.misc import face # for dummy data

img = face() # dummy RGB data
n,m = img.shape[:-1]

# transform pixels by 45 degrees
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)
Tinv = np.linalg.inv(T)

# find out the extent of the transformed pixels from the four corners
inew_tmp,jnew_tmp = T @ [[0, 0, n-1, n-1], [0, m-1, 0, m-1]] # 1d arrays each
imin,imax,jmin,jmax = inew_tmp.min(),inew_tmp.max(),jnew_tmp.min(),jnew_tmp.max()
imin,imax,jmin,jmax = (int(round(val)) for val in (imin,imax,jmin,jmax))

# so the pixels of the original map inside [imin, imax] x [jmin, jmax]
# we need an image of size (imax - imin + 1, jmax - jmin + 1) to house this

def transform_func(output_coords):
    """Inverse transform output coordinates back into input coordinates"""
    inew,jnew,channel = output_coords
    i,j = Tinv @ [inew + imin, jnew + jmin]
    return i,j,channel

out = ndi.geometric_transform(img, transform_func, output_shape = (imax - imin + 1, jmax - jmin + 1, 3))

# check what we've done
plt.imshow(out)
plt.show()

结果:

最终修复:仅限 numpy

我主要关心图像质量，因此上述所有解决方案都以一种或另一种方式使用插值。正如您在评论中所解释的，这不是您主要关心的问题。如果是这种情况，我们可以使用 map_coordinates 修改版本并自己计算近似(舍入整数)索引并执行矢量化分配:

import numpy as np
import matplotlib.pyplot as plt # for plotting the result
from scipy.misc import face # for dummy data

img = face() # dummy RGB data
n,m = img.shape[:-1]

# transform pixels by 45 degrees
T = np.array([[1, -1],[1, 1]])/np.sqrt(2)

# find out the extent of the transformed pixels from the four corners
inew_tmp,jnew_tmp = T @ [[0, 0, n-1, n-1], [0, m-1, 0, m-1]] # 1d arrays each
imin,imax,jmin,jmax = inew_tmp.min(),inew_tmp.max(),jnew_tmp.min(),jnew_tmp.max()
imin,imax,jmin,jmax = (int(round(val)) for val in (imin,imax,jmin,jmax))

# so the pixels of the original map inside [imin, imax] x [jmin, jmax]
# we need an image of size (imax - imin + 1, jmax - jmin + 1) to house this
out = np.zeros((imax - imin + 1, jmax - jmin + 1, 3), dtype=img.dtype)

# compute the corresponding coordinates on the domain for matching
inew,jnew = np.mgrid[:out.shape[0], :out.shape[1]]
inew = inew.ravel() # 1d array, indices of output array
jnew = jnew.ravel() # 1d array, indices of output array
i_back,j_back = np.linalg.inv(T) @ [inew + imin, jnew + jmin]

# create a mask to grab only those rounded (i_back,j_back) indices which make sense
i_back = i_back.round().astype(int)
j_back = j_back.round().astype(int)
inds = (0 <= i_back) & (i_back < n) & (0 <= j_back) & (j_back < m)
# (i_back[inds], j_back[inds]) maps to (inew[inds], jnew[inds])
# the rest stays black

out[inew[inds], jnew[inds], :] = img[i_back[inds], j_back[inds], :]

# check what we've done
plt.imshow(out)
plt.show()

结果虽然充满了单像素的不准确性，但看起来已经足够好了: