import numpy as np import torch.nn as nn import torch.nn.functional as F import torch from mmcv.cnn import ConvModule from mmengine.runner import load_checkpoint class FlowCompletionLoss(nn.Module): """Flow completion loss""" def __init__(self): super().__init__() self.fix_spynet = SPyNet() for p in self.fix_spynet.parameters(): p.requires_grad = False self.l1_criterion = nn.L1Loss() def forward(self, pred_flows, gt_local_frames): b, l_t, c, h, w = gt_local_frames.size() with torch.no_grad(): # compute gt forward and backward flows gt_local_frames = F.interpolate(gt_local_frames.view(-1, c, h, w), scale_factor=1 / 4, mode='bilinear', align_corners=True, recompute_scale_factor=True) gt_local_frames = gt_local_frames.view(b, l_t, c, h // 4, w // 4) gtlf_1 = gt_local_frames[:, :-1, :, :, :].reshape( -1, c, h // 4, w // 4) gtlf_2 = gt_local_frames[:, 1:, :, :, :].reshape( -1, c, h // 4, w // 4) gt_flows_forward = self.fix_spynet(gtlf_1, gtlf_2) gt_flows_backward = self.fix_spynet(gtlf_2, gtlf_1) # calculate loss for flow completion forward_flow_loss = self.l1_criterion( pred_flows[0].view(-1, 2, h // 4, w // 4), gt_flows_forward) backward_flow_loss = self.l1_criterion( pred_flows[1].view(-1, 2, h // 4, w // 4), gt_flows_backward) flow_loss = forward_flow_loss + backward_flow_loss return flow_loss class SPyNet(nn.Module): """SPyNet network structure. The difference to the SPyNet in [tof.py] is that 1. more SPyNetBasicModule is used in this version, and 2. no batch normalization is used in this version. Paper: Optical Flow Estimation using a Spatial Pyramid Network, CVPR, 2017 Args: pretrained (str): path for pre-trained SPyNet. Default: None. """ def __init__( self, use_pretrain=True, pretrained='https://download.openmmlab.com/mmediting/restorers/basicvsr/spynet_20210409-c6c1bd09.pth' ): super().__init__() self.basic_module = nn.ModuleList( [SPyNetBasicModule() for _ in range(6)]) if use_pretrain: if isinstance(pretrained, str): print("load pretrained SPyNet...") load_checkpoint(self, pretrained, strict=True) elif pretrained is not None: raise TypeError('[pretrained] should be str or None, ' f'but got {type(pretrained)}.') 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 compute_flow(self, ref, supp): """Compute flow from ref to supp. Note that in this function, the images are already resized to a multiple of 32. Args: ref (Tensor): Reference image with shape of (n, 3, h, w). supp (Tensor): Supporting image with shape of (n, 3, h, w). Returns: Tensor: Estimated optical flow: (n, 2, h, w). """ n, _, h, w = ref.size() # normalize the input images ref = [(ref - self.mean) / self.std] supp = [(supp - self.mean) / self.std] # generate downsampled frames for level in range(5): ref.append( F.avg_pool2d(input=ref[-1], kernel_size=2, stride=2, count_include_pad=False)) supp.append( F.avg_pool2d(input=supp[-1], kernel_size=2, stride=2, count_include_pad=False)) ref = ref[::-1] supp = supp[::-1] # flow computation flow = ref[0].new_zeros(n, 2, h // 32, w // 32) for level in range(len(ref)): if level == 0: flow_up = flow else: flow_up = F.interpolate(input=flow, scale_factor=2, mode='bilinear', align_corners=True) * 2.0 # add the residue to the upsampled flow flow = flow_up + self.basic_module[level](torch.cat([ ref[level], flow_warp(supp[level], flow_up.permute(0, 2, 3, 1).contiguous(), padding_mode='border'), flow_up ], 1)) return flow def forward(self, ref, supp): """Forward function of SPyNet. This function computes the optical flow from ref to supp. Args: ref (Tensor): Reference image with shape of (n, 3, h, w). supp (Tensor): Supporting image with shape of (n, 3, h, w). Returns: Tensor: Estimated optical flow: (n, 2, h, w). """ # upsize to a multiple of 32 h, w = ref.shape[2:4] w_up = w if (w % 32) == 0 else 32 * (w // 32 + 1) h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1) ref = F.interpolate(input=ref, size=(h_up, w_up), mode='bilinear', align_corners=False) supp = F.interpolate(input=supp, size=(h_up, w_up), mode='bilinear', align_corners=False) # compute flow, and resize back to the original resolution flow = F.interpolate(input=self.compute_flow(ref, supp), size=(h, w), mode='bilinear', align_corners=False) # adjust the flow values flow[:, 0, :, :] *= float(w) / float(w_up) flow[:, 1, :, :] *= float(h) / float(h_up) return flow class SPyNetBasicModule(nn.Module): """Basic Module for SPyNet. Paper: Optical Flow Estimation using a Spatial Pyramid Network, CVPR, 2017 """ def __init__(self): super().__init__() self.basic_module = nn.Sequential( ConvModule(in_channels=8, out_channels=32, kernel_size=7, stride=1, padding=3, norm_cfg=None, act_cfg=dict(type='ReLU')), ConvModule(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=3, norm_cfg=None, act_cfg=dict(type='ReLU')), ConvModule(in_channels=64, out_channels=32, kernel_size=7, stride=1, padding=3, norm_cfg=None, act_cfg=dict(type='ReLU')), ConvModule(in_channels=32, out_channels=16, kernel_size=7, stride=1, padding=3, norm_cfg=None, act_cfg=dict(type='ReLU')), ConvModule(in_channels=16, out_channels=2, kernel_size=7, stride=1, padding=3, norm_cfg=None, act_cfg=None)) def forward(self, tensor_input): """ Args: tensor_input (Tensor): Input tensor with shape (b, 8, h, w). 8 channels contain: [reference image (3), neighbor image (3), initial flow (2)]. Returns: Tensor: Refined flow with shape (b, 2, h, w) """ return self.basic_module(tensor_input) # Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization def make_colorwheel(): """ Generates a color wheel for optical flow visualization as presented in: Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf Code follows the original C++ source code of Daniel Scharstein. Code follows the the Matlab source code of Deqing Sun. Returns: np.ndarray: Color wheel """ RY = 15 YG = 6 GC = 4 CB = 11 BM = 13 MR = 6 ncols = RY + YG + GC + CB + BM + MR colorwheel = np.zeros((ncols, 3)) col = 0 # RY colorwheel[0:RY, 0] = 255 colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY) col = col + RY # YG colorwheel[col:col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG) colorwheel[col:col + YG, 1] = 255 col = col + YG # GC colorwheel[col:col + GC, 1] = 255 colorwheel[col:col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC) col = col + GC # CB colorwheel[col:col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB) colorwheel[col:col + CB, 2] = 255 col = col + CB # BM colorwheel[col:col + BM, 2] = 255 colorwheel[col:col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM) col = col + BM # MR colorwheel[col:col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR) colorwheel[col:col + MR, 0] = 255 return colorwheel def flow_uv_to_colors(u, v, convert_to_bgr=False): """ Applies the flow color wheel to (possibly clipped) flow components u and v. According to the C++ source code of Daniel Scharstein According to the Matlab source code of Deqing Sun Args: u (np.ndarray): Input horizontal flow of shape [H,W] v (np.ndarray): Input vertical flow of shape [H,W] convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. Returns: np.ndarray: Flow visualization image of shape [H,W,3] """ flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8) colorwheel = make_colorwheel() # shape [55x3] ncols = colorwheel.shape[0] rad = np.sqrt(np.square(u) + np.square(v)) a = np.arctan2(-v, -u) / np.pi fk = (a + 1) / 2 * (ncols - 1) k0 = np.floor(fk).astype(np.int32) k1 = k0 + 1 k1[k1 == ncols] = 0 f = fk - k0 for i in range(colorwheel.shape[1]): tmp = colorwheel[:, i] col0 = tmp[k0] / 255.0 col1 = tmp[k1] / 255.0 col = (1 - f) * col0 + f * col1 idx = (rad <= 1) col[idx] = 1 - rad[idx] * (1 - col[idx]) col[~idx] = col[~idx] * 0.75 # out of range # Note the 2-i => BGR instead of RGB ch_idx = 2 - i if convert_to_bgr else i flow_image[:, :, ch_idx] = np.floor(255 * col) return flow_image def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False): """ Expects a two dimensional flow image of shape. Args: flow_uv (np.ndarray): Flow UV image of shape [H,W,2] clip_flow (float, optional): Clip maximum of flow values. Defaults to None. convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. Returns: np.ndarray: Flow visualization image of shape [H,W,3] """ assert flow_uv.ndim == 3, 'input flow must have three dimensions' assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]' if clip_flow is not None: flow_uv = np.clip(flow_uv, 0, clip_flow) u = flow_uv[:, :, 0] v = flow_uv[:, :, 1] rad = np.sqrt(np.square(u) + np.square(v)) rad_max = np.max(rad) epsilon = 1e-5 u = u / (rad_max + epsilon) v = v / (rad_max + epsilon) return flow_uv_to_colors(u, v, convert_to_bgr) def flow_warp(x, flow, interpolation='bilinear', padding_mode='zeros', align_corners=True): """Warp an image or a feature map with optical flow. Args: x (Tensor): Tensor with size (n, c, h, w). flow (Tensor): Tensor with size (n, h, w, 2). The last dimension is a two-channel, denoting the width and height relative offsets. Note that the values are not normalized to [-1, 1]. interpolation (str): Interpolation mode: 'nearest' or 'bilinear'. Default: 'bilinear'. padding_mode (str): Padding mode: 'zeros' or 'border' or 'reflection'. Default: 'zeros'. align_corners (bool): Whether align corners. Default: True. Returns: Tensor: Warped image or feature map. """ if x.size()[-2:] != flow.size()[1:3]: raise ValueError(f'The spatial sizes of input ({x.size()[-2:]}) and ' f'flow ({flow.size()[1:3]}) are not the same.') _, _, h, w = x.size() # create mesh grid grid_y, grid_x = torch.meshgrid(torch.arange(0, h), torch.arange(0, w)) grid = torch.stack((grid_x, grid_y), 2).type_as(x) # (w, h, 2) grid.requires_grad = False grid_flow = grid + flow # scale grid_flow to [-1,1] grid_flow_x = 2.0 * grid_flow[:, :, :, 0] / max(w - 1, 1) - 1.0 grid_flow_y = 2.0 * grid_flow[:, :, :, 1] / max(h - 1, 1) - 1.0 grid_flow = torch.stack((grid_flow_x, grid_flow_y), dim=3) output = F.grid_sample(x, grid_flow, mode=interpolation, padding_mode=padding_mode, align_corners=align_corners) return output def initial_mask_flow(mask): """ mask 1 indicates valid pixel 0 indicates unknown pixel """ B, T, C, H, W = mask.shape # calculate relative position grid_y, grid_x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W)) grid_y, grid_x = grid_y.type_as(mask), grid_x.type_as(mask) abs_relative_pos_y = H - torch.abs(grid_y[None, :, :] - grid_y[:, None, :]) relative_pos_y = H - (grid_y[None, :, :] - grid_y[:, None, :]) abs_relative_pos_x = W - torch.abs(grid_x[:, None, :] - grid_x[:, :, None]) relative_pos_x = W - (grid_x[:, None, :] - grid_x[:, :, None]) # calculate the nearest indices pos_up = mask.unsqueeze(3).repeat( 1, 1, 1, H, 1, 1).flip(4) * abs_relative_pos_y[None, None, None] * ( relative_pos_y <= H)[None, None, None] nearest_indice_up = pos_up.max(dim=4)[1] pos_down = mask.unsqueeze(3).repeat(1, 1, 1, H, 1, 1) * abs_relative_pos_y[ None, None, None] * (relative_pos_y <= H)[None, None, None] nearest_indice_down = (pos_down).max(dim=4)[1] pos_left = mask.unsqueeze(4).repeat( 1, 1, 1, 1, W, 1).flip(5) * abs_relative_pos_x[None, None, None] * ( relative_pos_x <= W)[None, None, None] nearest_indice_left = (pos_left).max(dim=5)[1] pos_right = mask.unsqueeze(4).repeat( 1, 1, 1, 1, W, 1) * abs_relative_pos_x[None, None, None] * ( relative_pos_x <= W)[None, None, None] nearest_indice_right = (pos_right).max(dim=5)[1] # NOTE: IMPORTANT !!! depending on how to use this offset initial_offset_up = -(nearest_indice_up - grid_y[None, None, None]).flip(3) initial_offset_down = nearest_indice_down - grid_y[None, None, None] initial_offset_left = -(nearest_indice_left - grid_x[None, None, None]).flip(4) initial_offset_right = nearest_indice_right - grid_x[None, None, None] # nearest_indice_x = (mask.unsqueeze(1).repeat(1, img_width, 1) * relative_pos_x).max(dim=2)[1] # initial_offset_x = nearest_indice_x - grid_x # handle the boundary cases final_offset_down = (initial_offset_down < 0) * initial_offset_up + ( initial_offset_down > 0) * initial_offset_down final_offset_up = (initial_offset_up > 0) * initial_offset_down + ( initial_offset_up < 0) * initial_offset_up final_offset_right = (initial_offset_right < 0) * initial_offset_left + ( initial_offset_right > 0) * initial_offset_right final_offset_left = (initial_offset_left > 0) * initial_offset_right + ( initial_offset_left < 0) * initial_offset_left zero_offset = torch.zeros_like(final_offset_down) # out = torch.cat([final_offset_left, zero_offset, final_offset_right, zero_offset, zero_offset, final_offset_up, zero_offset, final_offset_down], dim=2) out = torch.cat([ zero_offset, final_offset_left, zero_offset, final_offset_right, final_offset_up, zero_offset, final_offset_down, zero_offset ], dim=2) return out