批量像素操作

Question

我有一批深度图像，形状 -> [B, 1, H, W]。对于我需要执行的批次的每个图像中的每个像素：

X = d * Kinverse @ [u, v, 1] #therefore X is in R^3 其中 d 是浮点张量 [0;1]，表示像素 u,v 处的深度； Kinverse 是一个常数 3X3 矩阵，u, v 分别代表像素的列和行。

是否有某种方法可以对操作进行矢量化以获得批处理中所有图像的 X(u+1,v)、X(u,v) 和 X(u,v+1)。 我最终需要采取这个交叉产品： {X(u+1,v) - X(u,v)} x {X(u, v+1) - X(u,v)}

感谢您的帮助！

Answer 1

您可以使用torch.meshgrid 生成u 和v 张量。拥有它们后，您可以使用torch.einsum 与Kinverse 进行批量矩阵乘法。 最后，您可以使用torch.cross 计算叉积：

u, v = torch.meshgrid(*[torch.arange(s_, dtype=d.dtype, device=d.device) for s_ in d.shape[2:]])
# make a single 1x1xHxW for [u v 1] per pixel:
uv = torch.cat((u[None, None, ...], v[None, None, ...], torch.ones_like(u)[None, None, ...]), dim=1)
# compute X
X = d * torch.einsum('ij,bjhw->bihw',Kinverse,uv)
# the cross product
out = torch.cross(X[..., 1:, :-1] - X[..., :-1, :-1], X[..., :-1, 1:] - X[..., :-1, :-1], dim=1)