33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import torch
|
|
|
|
def warp(x, flo):
|
|
'''
|
|
warp an image/tensor (im2) back to im1, according to the optical flow
|
|
x: [B, C, H, W] (im2)
|
|
flo: [B, 2, H, W] (flow)
|
|
'''
|
|
B, C, H, W = x.size()
|
|
# mesh grid
|
|
xx = torch.arange(0, W, device=x.device).view(1, -1).repeat(H, 1)
|
|
yy = torch.arange(0, H, device=x.device).view(-1, 1).repeat(1, W)
|
|
xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
|
|
yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
|
|
grid = torch.cat((xx, yy), 1).float()
|
|
|
|
if x.is_cuda:
|
|
grid = grid.to(x.device)
|
|
vgrid = torch.autograd.Variable(grid) + flo
|
|
|
|
# scale grid to [-1, 1]
|
|
vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0
|
|
vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0
|
|
|
|
vgrid = vgrid.permute(0, 2, 3, 1)
|
|
output = torch.nn.functional.grid_sample(x, vgrid, align_corners=True)
|
|
mask = torch.autograd.Variable(torch.ones((B, C, H, W), device=x.device))
|
|
mask = torch.nn.functional.grid_sample(mask, vgrid, align_corners=True)
|
|
|
|
mask[mask < 0.9999] = 0
|
|
mask[mask > 0] = 1
|
|
|
|
return output * mask, mask |