import torch import math import numpy import torch.nn.functional as F import torch.nn as nn from ..components import register import modules.components.upr_basic.softsplat as softsplat import modules.components.upr_basic.correlation as correlation from utils.padder import InputPadder #**************************************************************************************************# # => Feature Pyramid #**************************************************************************************************# class FeatPyramid(nn.Module): """A 3-level feature pyramid, which by default is shared by the motion estimator and synthesis network. """ def __init__(self): super(FeatPyramid, self).__init__() self.conv_stage0 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_stage1 = nn.Sequential( nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_stage2 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) def forward(self, img): C0 = self.conv_stage0(img) C1 = self.conv_stage1(C0) C2 = self.conv_stage2(C1) return [C0, C1, C2] #**************************************************************************************************# # => Motion Estimation #**************************************************************************************************# class MotionEstimator(nn.Module): """Bi-directional optical flow estimator 1) construct partial cost volume with the CNN features from the stage 2 of the feature pyramid; 2) estimate bi-directional flows, by feeding cost volume, CNN features for both warped images, CNN feature and estimated flow from previous iteration. """ def __init__(self): super(MotionEstimator, self).__init__() # (4*2 + 1) ** 2 + 128 * 2 + 128 + 4 = 469 self.conv_layer1 = nn.Sequential( nn.Conv2d(in_channels=469, out_channels=320, kernel_size=1, stride=1, padding=0), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_layer2 = nn.Sequential( nn.Conv2d(in_channels=320, out_channels=256, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_layer3 = nn.Sequential( nn.Conv2d(in_channels=256, out_channels=224, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_layer4 = nn.Sequential( nn.Conv2d(in_channels=224, out_channels=192, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_layer5 = nn.Sequential( nn.Conv2d(in_channels=192, out_channels=128, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(inplace=False, negative_slope=0.1)) self.conv_layer6 = nn.Sequential( nn.Conv2d(in_channels=128, out_channels=4, kernel_size=3, stride=1, padding=1)) def forward(self, feat0, feat1, last_feat, last_flow): corr_fn=correlation.FunctionCorrelation feat0 = softsplat.FunctionSoftsplat( tenInput=feat0, tenFlow=last_flow[:, :2]*0.25*0.5, tenMetric=None, strType='average') feat1 = softsplat.FunctionSoftsplat( tenInput=feat1, tenFlow=last_flow[:, 2:]*0.25*0.5, tenMetric=None, strType='average') volume = F.leaky_relu( input=corr_fn(tenFirst=feat0, tenSecond=feat1), negative_slope=0.1, inplace=False) input_feat = torch.cat([volume, feat0, feat1, last_feat, last_flow], 1) feat = self.conv_layer1(input_feat) feat = self.conv_layer2(feat) feat = self.conv_layer3(feat) feat = self.conv_layer4(feat) feat = self.conv_layer5(feat) flow = self.conv_layer6(feat) return flow, feat #**************************************************************************************************# # => Frame Synthesis #**************************************************************************************************# class SynthesisNetwork(nn.Module): def __init__(self): super(SynthesisNetwork, self).__init__() input_channels = 9+4+6 self.encoder_conv = nn.Sequential( nn.Conv2d(in_channels=input_channels, out_channels=64, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=64), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=64)) self.encoder_down1 = nn.Sequential( nn.Conv2d(in_channels=64 + 32 + 32, out_channels=128, kernel_size=3, stride=2, padding=1), nn.PReLU(num_parameters=128), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=128), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=128)) self.encoder_down2 = nn.Sequential( nn.Conv2d(in_channels=128 + 64 + 64, out_channels=256, kernel_size=3, stride=2, padding=1), nn.PReLU(num_parameters=256), nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=256), nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=256)) self.decoder_up1 = nn.Sequential( torch.nn.ConvTranspose2d(in_channels=256 + 128 + 128, out_channels=128, kernel_size=4, stride=2, padding=1, bias=True), nn.PReLU(num_parameters=128), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=128)) self.decoder_up2 = nn.Sequential( torch.nn.ConvTranspose2d(in_channels=128 + 128, out_channels=64, kernel_size=4, stride=2, padding=1, bias=True), nn.PReLU(num_parameters=64), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=64)) self.decoder_conv = nn.Sequential( nn.Conv2d(in_channels=64 + 64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=64), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.PReLU(num_parameters=64)) self.pred = nn.Conv2d(in_channels=64, out_channels=5, kernel_size=3, stride=1, padding=1) def get_warped_representations(self, bi_flow, c0, c1, i0=None, i1=None, time_step=0.5): flow_0t = bi_flow[:, :2] * time_step flow_1t = bi_flow[:, 2:4] * (1 - time_step) warped_c0 = softsplat.FunctionSoftsplat( tenInput=c0, tenFlow=flow_0t, tenMetric=None, strType='average') warped_c1 = softsplat.FunctionSoftsplat( tenInput=c1, tenFlow=flow_1t, tenMetric=None, strType='average') if (i0 is None) and (i1 is None): return warped_c0, warped_c1 else: warped_img0 = softsplat.FunctionSoftsplat( tenInput=i0, tenFlow=flow_0t, tenMetric=None, strType='average') warped_img1 = softsplat.FunctionSoftsplat( tenInput=i1, tenFlow=flow_1t, tenMetric=None, strType='average') flow_0t_1t = torch.cat((flow_0t, flow_1t), 1) return warped_img0, warped_img1, warped_c0, warped_c1, flow_0t_1t def forward(self, last_i, i0, i1, c0_pyr, c1_pyr, bi_flow_pyr, time_step=0.5): warped_img0, warped_img1, warped_c0, warped_c1, flow_0t_1t = \ self.get_warped_representations( bi_flow_pyr[0], c0_pyr[0], c1_pyr[0], i0, i1, time_step=time_step) input_feat = torch.cat( (last_i, warped_img0, warped_img1, i0, i1, flow_0t_1t), 1) s0 = self.encoder_conv(input_feat) s1 = self.encoder_down1(torch.cat((s0, warped_c0, warped_c1), 1)) warped_c0, warped_c1 = self.get_warped_representations( bi_flow_pyr[1], c0_pyr[1], c1_pyr[1], time_step=time_step) s2 = self.encoder_down2(torch.cat((s1, warped_c0, warped_c1), 1)) warped_c0, warped_c1 = self.get_warped_representations( bi_flow_pyr[2], c0_pyr[2], c1_pyr[2], time_step=time_step) x = self.decoder_up1(torch.cat((s2, warped_c0, warped_c1), 1)) x = self.decoder_up2(torch.cat((x, s1), 1)) x = self.decoder_conv(torch.cat((x, s0), 1)) # prediction refine = self.pred(x) refine_res = torch.sigmoid(refine[:, :3]) * 2 - 1 refine_mask0 = torch.sigmoid(refine[:, 3:4]) refine_mask1 = torch.sigmoid(refine[:, 4:5]) merged_img = (warped_img0 * refine_mask0 * (1 - time_step) + \ warped_img1 * refine_mask1 * time_step) merged_img = merged_img / (refine_mask0 * (1 - time_step) + \ refine_mask1 * time_step) interp_img = merged_img + refine_res interp_img = torch.clamp(interp_img, 0, 1) extra_dict = {} extra_dict["refine_res"] = refine_res extra_dict['refine_mask0'] = refine_mask0 extra_dict['refine_mask1'] = refine_mask1 extra_dict["warped_img0"] = warped_img0 extra_dict["warped_img1"] = warped_img1 extra_dict["merged_img"] = merged_img extra_dict['c0_pyr'] = c0_pyr extra_dict['c1_pyr'] = c1_pyr extra_dict['s0'] = s0 extra_dict['s1'] = s1 extra_dict['s2'] = s2 return interp_img, extra_dict #**************************************************************************************************# # => Unified model #**************************************************************************************************# @register('upr_basic') class Model(nn.Module): def __init__(self, pyr_level=3, nr_lvl_skipped=0, *args, **kwargs): print('@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@UPR_basic (REAL)@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@') super(Model, self).__init__() self.pyr_level = pyr_level self.nr_lvl_skipped = nr_lvl_skipped self.feat_pyramid = FeatPyramid() self.motion_estimator = MotionEstimator() self.synthesis_network = SynthesisNetwork() def forward_one_lvl(self, img0, img1, last_feat, last_flow, last_interp=None, time_step=0.5, skip_me=False): # context feature extraction feat0_pyr = self.feat_pyramid(img0) feat1_pyr = self.feat_pyramid(img1) # bi-directional flow estimation if not skip_me: flow, feat = self.motion_estimator( feat0_pyr[-1], feat1_pyr[-1], last_feat, last_flow) else: flow = last_flow feat = last_feat # frame synthesis ## optical flow is estimated at 1/4 resolution ori_resolution_flow = F.interpolate( input=flow, scale_factor=4.0, mode="bilinear", align_corners=False) ## consturct 3-level flow pyramid for synthesis network bi_flow_pyr = [] tmp_flow = ori_resolution_flow bi_flow_pyr.append(tmp_flow) for i in range(2): tmp_flow = F.interpolate( input=tmp_flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 bi_flow_pyr.append(tmp_flow) ## merge warped frames as initial interpolation for frame synthesis if last_interp is None: flow_0t = ori_resolution_flow[:, :2] * time_step flow_1t = ori_resolution_flow[:, 2:4] * (1 - time_step) warped_img0 = softsplat.FunctionSoftsplat( tenInput=img0, tenFlow=flow_0t, tenMetric=None, strType='average') warped_img1 = softsplat.FunctionSoftsplat( tenInput=img1, tenFlow=flow_1t, tenMetric=None, strType='average') last_interp = warped_img0 * (1 - time_step) \ + warped_img1 * time_step ## do synthesis interp_img, extra_dict = self.synthesis_network( last_interp, img0, img1, feat0_pyr, feat1_pyr, bi_flow_pyr, time_step=time_step) return flow, feat, interp_img, extra_dict def forward(self, img0, img1, time_step, pyr_level=None, nr_lvl_skipped=None, **kwargs): if pyr_level is None: pyr_level = self.pyr_level if nr_lvl_skipped is None: nr_lvl_skipped = self.nr_lvl_skipped N, _, H, W = img0.shape bi_flows = [] interp_imgs = [] skipped_levels = [] if nr_lvl_skipped == 0 else\ list(range(pyr_level))[::-1][-nr_lvl_skipped:] padder = InputPadder(img0.shape, divisor=int(4 * 2 ** pyr_level)) img0, img1 = padder.pad(img0, img1) N, _, H, W = img0.shape # The original input resolution corresponds to level 0. for level in list(range(pyr_level))[::-1]: if level != 0: scale_factor = 1 / 2 ** level img0_this_lvl = F.interpolate( input=img0, scale_factor=scale_factor, mode="bilinear", align_corners=False) img1_this_lvl = F.interpolate( input=img1, scale_factor=scale_factor, mode="bilinear", align_corners=False) else: img0_this_lvl = img0 img1_this_lvl = img1 # skip motion estimation, directly use up-sampled optical flow skip_me = False # the lowest-resolution pyramid level if level == pyr_level - 1: last_flow = torch.zeros( (N, 4, H // (2 ** (level+2)), W //(2 ** (level+2))) ).to(img0.device) last_feat = torch.zeros( (N, 128, H // (2 ** (level+2)), W // (2 ** (level+2))) ).to(img0.device) last_interp = None # skip some levels for both motion estimation and frame synthesis elif level in skipped_levels[:-1]: continue # last level (original input resolution), only skip motion estimation elif (level == 0) and len(skipped_levels) > 0: if len(skipped_levels) == pyr_level: last_flow = torch.zeros( (N, 4, H // 4, W // 4)).to(img0.device) last_interp = None else: resize_factor = 2 ** len(skipped_levels) last_flow = F.interpolate( input=flow, scale_factor=resize_factor, mode="bilinear", align_corners=False) * resize_factor last_interp = F.interpolate( input=interp_img, scale_factor=resize_factor, mode="bilinear", align_corners=False) skip_me = True # last level (original input resolution), motion estimation + frame # synthesis else: last_flow = F.interpolate(input=flow, scale_factor=2.0, mode="bilinear", align_corners=False) * 2 last_feat = F.interpolate(input=feat, scale_factor=2.0, mode="bilinear", align_corners=False) * 2 last_interp = F.interpolate( input=interp_img, scale_factor=2.0, mode="bilinear", align_corners=False) flow, feat, interp_img, extra_dict = self.forward_one_lvl( img0_this_lvl, img1_this_lvl, last_feat, last_flow, last_interp, time_step, skip_me=skip_me) bi_flows.append( padder.unpad(F.interpolate(input=flow, scale_factor=4.0, mode="bilinear", align_corners=False))) interp_imgs.append(padder.unpad(interp_img)) # directly up-sample estimated flow to full resolution with bi-linear # interpolation bi_flow = F.interpolate( input=flow, scale_factor=4.0, mode="bilinear", align_corners=False) result_dict = { "imgt_preds": interp_imgs, 'imgt_pred': interp_imgs[-1].contiguous(),"bi_flows": bi_flows, "flowfwd": bi_flows[-1][:,:2], "flowbwd": bi_flows[-1][:,2:] } return result_dict, extra_dict if __name__ == "__main__": pass