31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def get_grid(x):
|
|
torchHorizontal = (
|
|
torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3))
|
|
)
|
|
torchVertical = (
|
|
torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3))
|
|
)
|
|
return torch.cat([torchHorizontal, torchVertical], 1).cuda()
|
|
|
|
|
|
class WarpingLayer(nn.Module):
|
|
def __init__(self, device):
|
|
super(WarpingLayer, self).__init__()
|
|
self.device = device
|
|
|
|
def forward(self, x, flow):
|
|
# WarpingLayer uses F.grid_sample, which expects normalized grid
|
|
# we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code
|
|
# so here we need to denormalize the flow
|
|
flow_for_grip = torch.zeros_like(flow).cuda()
|
|
flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0)
|
|
flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0)
|
|
|
|
grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1)
|
|
return F.grid_sample(x, grid, align_corners=True)
|