Pixel operations in batches

Question

I have a batch of depth images, shape -> [B, 1, H, W]. For each pixel in each image of the batch I need to perform:

X = d * Kinverse @ [u, v, 1] #therefore X is in R^3 where d is float tensor[0;1] representing depth at pixel u,v; Kinverse is a constant 3X3 matrix and u, v refer to the pixel column and row respectively.

Is there some way I can vectorize the operation to obtain X(u+1,v), X(u,v) and X(u,v+1) for all the images in the batch. I eventually need to take this cross product: {X(u+1,v) - X(u,v)} x {X(u, v+1) - X(u,v)}

Thanks for the help!

so `d` is of shape `[B, 1, H, W]` ? – Shai Mar 16 '21 at 06:57 — Shai, Mar 16 '21 at 06:57
@Shai yes! Thanks for the help! – SouravBhunia Mar 16 '21 at 07:01 — SouravBhunia, Mar 16 '21 at 07:01

score 1 · Accepted Answer · answered Mar 16 '21 at 07:23

You can use torch.meshgrid to produce the u and v tensors. Once you have them, you can use torch.einsum to do the batched matrix multiplication with Kinverse. Finally, you can use torch.cross to compute the cross product:

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)

Pixel operations in batches

1 Answers1