像素分批操作

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)