问如何使用光流和grid_sample对图像进行扭曲？

EN

import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np
from imageio import imread

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).view(1 ,-1).repeat(H ,1)
yy = torch.arange(0, H).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.cuda()
vgrid = 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)
flo = flo.permute(0 ,2 ,3 ,1)
output = F.grid_sample(x, vgrid)
mask = torch.autograd.Variable(torch.ones(x.size())).cuda()
mask = F.grid_sample(mask, vgrid)

mask[mask <0.9999] = 0
mask[mask >0] = 1

return output*mask
test = cv2.readOpticalFlow(‘reference_frame_0010.flo’)
test = torch.from_numpy(test)
test = test.view(1, 2, H, W)
test = Variable(test).cuda()

test_img2 = cv2.imread(‘frame_0011.png’)
test_img2 = torch.from_numpy(test_img2)
test_img2 = test_img2.view(1, 3, H, W)
test_img2 = Variable(test_img2).cuda()

out = warp(test_img2, test)
out = out.view(H, W, 3)
out = out.cpu()
out = np.float32(out)
cv2.imwrite(‘test3.png’, out*255)