| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .warp import warp | |
| def resize(x, scale_factor): | |
| return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False) | |
| def convrelu(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True): | |
| return nn.Sequential( | |
| nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=bias), | |
| nn.PReLU(out_channels) | |
| ) | |
| class ResBlock(nn.Module): | |
| def __init__(self, in_channels, side_channels, bias=True): | |
| super(ResBlock, self).__init__() | |
| self.side_channels = side_channels | |
| self.conv1 = nn.Sequential( | |
| nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), | |
| nn.PReLU(in_channels) | |
| ) | |
| self.conv2 = nn.Sequential( | |
| nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), | |
| nn.PReLU(side_channels) | |
| ) | |
| self.conv3 = nn.Sequential( | |
| nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), | |
| nn.PReLU(in_channels) | |
| ) | |
| self.conv4 = nn.Sequential( | |
| nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), | |
| nn.PReLU(side_channels) | |
| ) | |
| self.conv5 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias) | |
| self.prelu = nn.PReLU(in_channels) | |
| def forward(self, x): | |
| out = self.conv1(x) | |
| res_feat = out[:, :-self.side_channels, ...] | |
| side_feat = out[:, -self.side_channels:, :, :] | |
| side_feat = self.conv2(side_feat) | |
| out = self.conv3(torch.cat([res_feat, side_feat], 1)) | |
| res_feat = out[:, :-self.side_channels, ...] | |
| side_feat = out[:, -self.side_channels:, :, :] | |
| side_feat = self.conv4(side_feat) | |
| out = self.conv5(torch.cat([res_feat, side_feat], 1)) | |
| out = self.prelu(x + out) | |
| return out | |
| class Encoder(nn.Module): | |
| def __init__(self, channels, large=False): | |
| super(Encoder, self).__init__() | |
| self.channels = channels | |
| prev_ch = 3 | |
| for idx, ch in enumerate(channels, 1): | |
| k = 7 if large and idx == 1 else 3 | |
| p = 3 if k == 7 else 1 | |
| self.register_module(f'pyramid{idx}', | |
| nn.Sequential( | |
| convrelu(prev_ch, ch, k, 2, p), | |
| convrelu(ch, ch, 3, 1, 1) | |
| )) | |
| prev_ch = ch | |
| def forward(self, in_x): | |
| fs = [] | |
| for idx in range(len(self.channels)): | |
| out_x = getattr(self, f'pyramid{idx + 1}')(in_x) | |
| fs.append(out_x) | |
| in_x = out_x | |
| return fs | |
| class InitDecoder(nn.Module): | |
| def __init__(self, in_ch, out_ch, skip_ch) -> None: | |
| super().__init__() | |
| self.convblock = nn.Sequential( | |
| convrelu(in_ch * 2 + 1, in_ch * 2), | |
| ResBlock(in_ch * 2, skip_ch), | |
| nn.ConvTranspose2d(in_ch * 2, out_ch + 4, 4, 2, 1, bias=True) | |
| ) | |
| def forward(self, f0, f1, embt): | |
| h, w = f0.shape[2:] | |
| embt = embt.repeat(1, 1, h, w) | |
| out = self.convblock(torch.cat([f0, f1, embt], 1)) | |
| flow0, flow1 = torch.chunk(out[:, :4, ...], 2, 1) | |
| ft_ = out[:, 4:, ...] | |
| return flow0, flow1, ft_ | |
| class IntermediateDecoder(nn.Module): | |
| def __init__(self, in_ch, out_ch, skip_ch) -> None: | |
| super().__init__() | |
| self.convblock = nn.Sequential( | |
| convrelu(in_ch * 2 + 2, in_ch * 2), | |
| ResBlock(in_ch * 2, skip_ch), | |
| nn.ConvTranspose2d(in_ch * 2, out_ch, 4, 2, 1, bias=True) | |
| ) | |
| self.conv_flow = nn.Conv2d(out_ch, 2, 3, 1, 1) | |
| def forward(self, f0, f1, flow_fwd, flow_bwd): | |
| f0_warp = warp(f0, flow_bwd) | |
| f1_warp = warp(f1, flow_fwd) | |
| f0_in = torch.cat([f0, f1_warp, flow_fwd], 1) | |
| f1_in = torch.cat([f1, f0_warp, flow_bwd], 1) | |
| out0 = self.convblock(f0_in) | |
| out1 = self.convblock(f1_in) | |
| flow_fwd = 2.0 * resize(flow_fwd, scale_factor=2.0) + self.conv_flow(out0) | |
| flow_bwd = 2.0 * resize(flow_bwd, scale_factor=2.0) + self.conv_flow(out1) | |
| return flow_fwd, flow_bwd, out0, out1 | |