Issue
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!
Solution
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)
Answered By - Shai
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.