| | 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 |
| |
|
