TVL1计算光流,并进行warp
import cv2
def compute_TVL1(prev, curr):
"""Compute the TV-L1 optical flow."""
TVL1 = cv2.optflow.DualTVL1OpticalFlow_create()
flow = TVL1.calc(prev, curr,None)
return flow
img1 = cv2.imread("18_98.jpg",cv2.IMREAD_GRAYSCALE)
img2 = cv2.imread("18_99.jpg", cv2.IMREAD_GRAYSCALE)
print(img1.shape)
flow = compute_TVL1(img1,img2)
import torch
from torch.autograd import Variable
from torch import nn
import numpy as np
import torchvision.transforms as transforms
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()
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()
x = x
grid = grid
vgrid = Variable(grid) + flo
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 = nn.functional.grid_sample(x, vgrid,mode="bilinear",padding_mode="zeros")
mask = torch.autograd.Variable(torch.ones(x.size()))
mask = nn.functional.grid_sample(mask, vgrid,align_corners=True)
return output
tvl1 = transforms.ToTensor()(flow).unsqueeze(0)
img=transforms.ToTensor((cv2.imread("18_98.jpg")).unsqueeze(0)
w = warp(img,tvl1)
w = w.squeeze(0)
w = np.transpose(w, (1,2,0)).numpy()
w = (w*255).astype("int")