import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from .update import BasicUpdateBlock, SmallUpdateBlock from .extractor import BasicEncoder, SmallEncoder from .corr import CorrBlock, AlternateCorrBlock from .utils.utils import bilinear_sampler, coords_grid, upflow8 import argparse from pathlib import Path try: autocast = torch.cuda.amp.autocast except: # dummy autocast for PyTorch < 1.6 class autocast: def __init__(self, enabled): pass def __enter__(self): pass def __exit__(self, *args): pass class Dummy: def __init__(self, enabled): pass def __enter__(self): pass def __exit__(self, *args): pass def get_args(cmd=None): parser = argparse.ArgumentParser() parser.add_argument('--corr_levels', type=int, default=4) parser.add_argument('--corr_radius', type=int, default=4) parser.add_argument('--dropout', type=float, default=0.0) parser.add_argument('--mixed_precision', action='store_true') parser.add_argument('--small', action='store_true') parser.add_argument('--gpus', type=int, nargs='+', default=[0]) if cmd is None: args = parser.parse_args() else: args = parser.parse_args(cmd) return args def load_raft_model(load_path, ignore_prefix=None, multiframe=False, scale_inputs=False, **kwargs): path = Path(load_path) if load_path else None args = get_args("") for k,v in kwargs.items(): args.__setattr__(k,v) args.multiframe = multiframe args.scale_inputs = scale_inputs model = RAFT(args) if load_path is not None: weight_dict = torch.load(load_path, map_location=torch.device("cpu")) new_dict = dict() for k in weight_dict.keys(): if 'module' in k: new_dict[k.replace('module.', '')] = weight_dict[k] else: new_dict[k] = weight_dict[k] if ignore_prefix is not None: new_dict_1 = dict() for k, v in new_dict.items(): new_dict_1[k.replace(ignore_prefix, '')] = v new_dict = new_dict_1 did_load = model.load_state_dict(new_dict, strict=False) print(did_load, type(model).__name__, load_path) else: print("created a new %s with %d parameters" % ( type(model).__name__, sum([v.numel() for v in model.parameters()]))) return model def get_raft_flow(x, raft_model, iters=24, backward=False, t_dim=1): assert len(x.shape) == 5, x.shape assert x.shape[t_dim] >= 2, x.shape x = x * 255.0 inds = torch.tensor([0,1]).to(x.device) x1, x2 = torch.index_select(x, t_dim, inds).unbind(t_dim) if backward: flow = raft_model(x2, x1, test_mode=True, iters=iters)[-1] else: flow = raft_model(x1, x2, test_mode=True, iters=iters)[-1] return flow class RAFT(nn.Module): def __init__(self, args): super(RAFT, self).__init__() self.args = args self.multiframe = self.args.multiframe self.scale_inputs = self.args.scale_inputs if args.small: self.hidden_dim = hdim = 96 self.context_dim = cdim = 64 args.corr_levels = 4 args.corr_radius = 3 else: self.hidden_dim = hdim = 128 self.context_dim = cdim = 128 args.corr_levels = 4 args.corr_radius = 4 if 'dropout' not in self.args: self.args.dropout = 0 if 'alternate_corr' not in self.args: self.args.alternate_corr = False # feature network, context network, and update block if args.small: self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout) self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout) self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim) else: self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout) self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout) self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim) def freeze_bn(self): for m in self.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() def initialize_flow(self, img): """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0""" N, C, H, W = img.shape coords0 = coords_grid(N, H//8, W//8, device=img.device, dtype=img.dtype) coords1 = coords_grid(N, H//8, W//8, device=img.device, dtype=img.dtype) # optical flow computed as difference: flow = coords1 - coords0 return coords0, coords1 def upsample_flow(self, flow, mask): """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ N, _, H, W = flow.shape mask = mask.view(N, 1, 9, 8, 8, H, W) mask = torch.softmax(mask, dim=2) up_flow = F.unfold(8 * flow, [3,3], padding=1) up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) up_flow = torch.sum(mask * up_flow, dim=2) up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) return up_flow.reshape(N, 2, 8*H, 8*W) @property def iters(self): if getattr(self, '_iters', None) is None: return None return self._iters @iters.setter def iters(self, value=None): self._iters = value def _forward_two_images( self, image1, image2, iters=24, flow_init=None, upsample=True, test_mode=True, **kwargs): """ Estimate optical flow between pair of frames """ if self.iters is not None: iters = self.iters image1 = 2 * (image1 / 255.0) - 1.0 image2 = 2 * (image2 / 255.0) - 1.0 image1 = image1.contiguous() image2 = image2.contiguous() hdim = self.hidden_dim cdim = self.context_dim # run the feature network decorator = autocast(enabled=True) if \ (self.args.mixed_precision or (image1.dtype in [torch.float16, torch.bfloat16])) \ else Dummy(enabled=False) with decorator: fmap1, fmap2 = self.fnet([image1, image2]) fmap1 = fmap1.float() fmap2 = fmap2.float() if self.args.alternate_corr: corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius) else: corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius) # run the context network # with autocast(enabled=self.args.mixed_precision): with decorator: cnet = self.cnet(image1) net, inp = torch.split(cnet, [hdim, cdim], dim=1) net = torch.tanh(net) inp = torch.relu(inp) coords0, coords1 = self.initialize_flow(image1) if flow_init is not None: coords1 = coords1 + flow_init flow_predictions = [] for itr in range(iters): coords1 = coords1.detach() corr = corr_fn(coords1) # index correlation volume flow = coords1 - coords0 # with autocast(enabled=self.args.mixed_precision): with decorator: net, up_mask, delta_flow, motion_features = self.update_block(net, inp, corr, flow) # F(t+1) = F(t) + \Delta(t) coords1 = coords1 + delta_flow # upsample predictions if up_mask is None: flow_up = upflow8(coords1 - coords0) else: flow_up = self.upsample_flow(coords1 - coords0, up_mask) flow_predictions.append(flow_up) if test_mode: return coords1 - coords0, flow_up, motion_features return flow_predictions, motion_features def forward(self, *args, **kwargs): if not self.multiframe: return self._forward_two_images(*args, **kwargs) x = (args[0] * 255.0) if self.scale_inputs else args[0] assert len(x.shape) == 5, x.shape assert x.shape[1] >= 2, x.shape num_frames = x.size(1) flows = [] motion_features = [] backward = kwargs.get('backward', False) for t in range(num_frames-1): x1, x2 = torch.index_select( x, 1, torch.tensor([t,t+1]).to(x.device)).unbind(1) _args = (x2, x1) if backward else (x1, x2) _, flow, features = self._forward_two_images(*_args, *args[1:], **kwargs) flows.insert(0, flow) if backward else flows.append(flow) motion_features.append(features) return torch.stack(flows, 1), torch.stack(motion_features, 1)