|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from raft import RAFT |
|
from nnutils import make_conv_2d, make_upscale_2d, make_downscale_2d, ResBlock2d, Identity |
|
|
|
|
|
class ImportanceWeights(torch.nn.Module): |
|
def __init__(self, opt): |
|
super().__init__() |
|
|
|
if opt.small: |
|
in_dim = 128 |
|
else: |
|
in_dim = 256 |
|
fn_0 = 16 |
|
self.input_fn = fn_0 + 3 * 2 |
|
fn_1 = 16 |
|
self.conv1 = torch.nn.Conv2d(in_channels=in_dim, out_channels=fn_0, kernel_size=3, stride=1, padding=1) |
|
|
|
if opt.use_batch_norm: |
|
custom_batch_norm = torch.nn.BatchNorm2d |
|
else: |
|
custom_batch_norm = Identity |
|
|
|
self.model = nn.Sequential( |
|
make_conv_2d(self.input_fn, fn_1, n_blocks=1, normalization=custom_batch_norm), |
|
ResBlock2d(fn_1, normalization=custom_batch_norm), |
|
ResBlock2d(fn_1, normalization=custom_batch_norm), |
|
ResBlock2d(fn_1, normalization=custom_batch_norm), |
|
nn.Conv2d(fn_1, 1, kernel_size=3, padding=1) |
|
|
|
) |
|
|
|
def forward(self, x, features): |
|
|
|
features = self.conv1(features) |
|
x = torch.cat([features, x], 1) |
|
assert x.shape[1] == self.input_fn |
|
x = self.model(x) |
|
print(x) |
|
print(x.max(), x.min(), x.mean()) |
|
|
|
return torch.nn.Sigmoid()(x) |
|
|
|
class NeuralNRT(nn.Module): |
|
def __init__(self, opt, path=None, device="cuda:0"): |
|
super(NeuralNRT, self).__init__() |
|
self.opt = opt |
|
self.CorresPred = RAFT(opt) |
|
self.ImportanceW = ImportanceWeights(opt) |
|
if path is not None: |
|
data = torch.load(path,map_location='cpu') |
|
if 'state_dict' in data.keys(): |
|
self.CorresPred.load_state_dict(data['state_dict']) |
|
print("load done") |
|
else: |
|
self.CorresPred.load_state_dict({k.replace('module.', ''):v for k,v in data.items()}) |
|
print("load done") |
|
def forward(self, src_im,tar_im, src_im_raw, tar_im_raw, Crop_param): |
|
N=src_im.shape[0] |
|
src_im = src_im*255.0 |
|
tar_im = tar_im*255.0 |
|
flow_fw_crop, feature_fw_crop = self.CorresPred(src_im, tar_im, iters=self.opt.iters) |
|
|
|
xx = torch.arange(0, self.opt.width).view(1,-1).repeat(self.opt.height,1) |
|
yy = torch.arange(0, self.opt.height).view(-1,1).repeat(1,self.opt.width) |
|
xx = xx.view(1,1,self.opt.height,self.opt.width).repeat(N,1,1,1) |
|
yy = yy.view(1,1,self.opt.height,self.opt.width).repeat(N,1,1,1) |
|
grid = torch.cat((xx,yy),1).float() |
|
grid = grid.to(src_im.device) |
|
|
|
grid_crop = grid[:, :, :self.opt.crop_height, :self.opt.crop_width] |
|
|
|
flow_fw = torch.zeros((N, 2, self.opt.height, self.opt.width), device=src_im.device) |
|
|
|
leftup1 = torch.cat((Crop_param[:, 0:1, 0], Crop_param[:, 2:3, 0]), 1)[:, :, None, None] |
|
leftup2 = torch.cat((Crop_param[:, 4:5, 0], Crop_param[:, 6:7, 0]), 1)[:, :, None, None] |
|
|
|
scale1 = torch.cat(((Crop_param[:, 1:2, 0]-Crop_param[:, 0:1, 0]).float() / self.opt.crop_width, (Crop_param[:, 3:4, 0]-Crop_param[:, 2:3, 0]).float() / self.opt.crop_height), 1)[:, :, None, None] |
|
scale2 = torch.cat(((Crop_param[:, 5:6, 0]-Crop_param[:, 4:5, 0]).float() / self.opt.crop_width, (Crop_param[:, 7:8, 0]-Crop_param[:, 6:7, 0]).float() / self.opt.crop_height), 1)[:, :, None, None] |
|
|
|
flow_fw_crop = (scale2 - scale1) * grid_crop + scale2 * flow_fw_crop |
|
for i in range(N): |
|
flow_fw_cropi = F.interpolate(flow_fw_crop[i:(i+1)], ((Crop_param[i, 3, 0]-Crop_param[i, 2, 0]).item(), (Crop_param[i, 1, 0]-Crop_param[i, 0, 0]).item()), mode='bilinear', align_corners=True) |
|
flow_fw_cropi =flow_fw_cropi + (leftup2 - leftup1)[i:(i+1), :, :, :] |
|
flow_fw[i, :, Crop_param[i, 2, 0]:Crop_param[i, 3, 0], Crop_param[i, 0, 0]:Crop_param[i, 1, 0]] = flow_fw_cropi[0] |
|
return flow_fw |
|
|
