| import math |
| import torch |
| from torch import nn as nn |
| from torch.nn import functional as F |
|
|
| from basicsr.utils.registry import ARCH_REGISTRY |
| from .arch_util import flow_warp |
|
|
|
|
| class BasicModule(nn.Module): |
| """Basic Module for SpyNet. |
| """ |
|
|
| def __init__(self): |
| super(BasicModule, self).__init__() |
|
|
| self.basic_module = nn.Sequential( |
| nn.Conv2d(in_channels=8, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False), |
| nn.Conv2d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False), |
| nn.Conv2d(in_channels=64, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False), |
| nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False), |
| nn.Conv2d(in_channels=16, out_channels=2, kernel_size=7, stride=1, padding=3)) |
|
|
| def forward(self, tensor_input): |
| return self.basic_module(tensor_input) |
|
|
|
|
| @ARCH_REGISTRY.register() |
| class SpyNet(nn.Module): |
| """SpyNet architecture. |
| |
| Args: |
| load_path (str): path for pretrained SpyNet. Default: None. |
| """ |
|
|
| def __init__(self, load_path=None): |
| super(SpyNet, self).__init__() |
| self.basic_module = nn.ModuleList([BasicModule() for _ in range(6)]) |
| if load_path: |
| self.load_state_dict(torch.load(load_path, map_location=lambda storage, loc: storage)['params']) |
|
|
| self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) |
| self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) |
|
|
| def preprocess(self, tensor_input): |
| tensor_output = (tensor_input - self.mean) / self.std |
| return tensor_output |
|
|
| def process(self, ref, supp): |
| flow = [] |
|
|
| ref = [self.preprocess(ref)] |
| supp = [self.preprocess(supp)] |
|
|
| for level in range(5): |
| ref.insert(0, F.avg_pool2d(input=ref[0], kernel_size=2, stride=2, count_include_pad=False)) |
| supp.insert(0, F.avg_pool2d(input=supp[0], kernel_size=2, stride=2, count_include_pad=False)) |
|
|
| flow = ref[0].new_zeros( |
| [ref[0].size(0), 2, |
| int(math.floor(ref[0].size(2) / 2.0)), |
| int(math.floor(ref[0].size(3) / 2.0))]) |
|
|
| for level in range(len(ref)): |
| upsampled_flow = F.interpolate(input=flow, scale_factor=2, mode='bilinear', align_corners=True) * 2.0 |
|
|
| if upsampled_flow.size(2) != ref[level].size(2): |
| upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 0, 0, 1], mode='replicate') |
| if upsampled_flow.size(3) != ref[level].size(3): |
| upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 1, 0, 0], mode='replicate') |
|
|
| flow = self.basic_module[level](torch.cat([ |
| ref[level], |
| flow_warp( |
| supp[level], upsampled_flow.permute(0, 2, 3, 1), interp_mode='bilinear', padding_mode='border'), |
| upsampled_flow |
| ], 1)) + upsampled_flow |
|
|
| return flow |
|
|
| def forward(self, ref, supp): |
| assert ref.size() == supp.size() |
|
|
| h, w = ref.size(2), ref.size(3) |
| w_floor = math.floor(math.ceil(w / 32.0) * 32.0) |
| h_floor = math.floor(math.ceil(h / 32.0) * 32.0) |
|
|
| ref = F.interpolate(input=ref, size=(h_floor, w_floor), mode='bilinear', align_corners=False) |
| supp = F.interpolate(input=supp, size=(h_floor, w_floor), mode='bilinear', align_corners=False) |
|
|
| flow = F.interpolate(input=self.process(ref, supp), size=(h, w), mode='bilinear', align_corners=False) |
|
|
| flow[:, 0, :, :] *= float(w) / float(w_floor) |
| flow[:, 1, :, :] *= float(h) / float(h_floor) |
|
|
| return flow |
|
|
