import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from collections import OrderedDict from networks.modules.extractor import BasicEncoder from networks.modules.corr import CorrBlock from networks.modules.gru import ConvGRU from networks.modules.clipping import GradientClip from lietorch import SE3 from networks.geom.ba import BA import networks.geom.projective_ops as pops from networks.geom.graph_utils import graph_to_edge_list, keyframe_indicies from torch_scatter import scatter_mean def cvx_upsample(data, mask): """ upsample pixel-wise transformation field """ batch, ht, wd, dim = data.shape data = data.permute(0, 3, 1, 2) mask = mask.view(batch, 1, 9, 8, 8, ht, wd) mask = torch.softmax(mask, dim=2) up_data = F.unfold(data, [3,3], padding=1) up_data = up_data.view(batch, dim, 9, 1, 1, ht, wd) up_data = torch.sum(mask * up_data, dim=2) up_data = up_data.permute(0, 4, 2, 5, 3, 1) up_data = up_data.reshape(batch, 8*ht, 8*wd, dim) return up_data def upsample_disp(disp, mask): batch, num, ht, wd = disp.shape disp = disp.view(batch*num, ht, wd, 1) mask = mask.view(batch*num, -1, ht, wd) return cvx_upsample(disp, mask).view(batch, num, 8*ht, 8*wd) class GraphAgg(nn.Module): def __init__(self): super(GraphAgg, self).__init__() self.conv1 = nn.Conv2d(128, 128, 3, padding=1) self.conv2 = nn.Conv2d(128, 128, 3, padding=1) self.relu = nn.ReLU(inplace=True) self.eta = nn.Sequential( nn.Conv2d(128, 1, 3, padding=1), GradientClip(), nn.Softplus()) self.upmask = nn.Sequential( nn.Conv2d(128, 8*8*9, 1, padding=0)) def forward(self, net, ii): batch, num, ch, ht, wd = net.shape net = net.view(batch*num, ch, ht, wd) _, ix = torch.unique(ii, return_inverse=True) net = self.relu(self.conv1(net)) net = net.view(batch, num, 128, ht, wd) net = scatter_mean(net, ix, dim=1) net = net.view(-1, 128, ht, wd) net = self.relu(self.conv2(net)) eta = self.eta(net).view(batch, -1, ht, wd) upmask = self.upmask(net).view(batch, -1, 8*8*9, ht, wd) return .01 * eta, upmask class UpdateModule(nn.Module): def __init__(self): super(UpdateModule, self).__init__() cor_planes = 4 * (2*3 + 1)**2 self.corr_encoder = nn.Sequential( nn.Conv2d(cor_planes, 128, 1, padding=0), nn.ReLU(inplace=True), nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(inplace=True)) self.flow_encoder = nn.Sequential( nn.Conv2d(4, 128, 7, padding=3), nn.ReLU(inplace=True), nn.Conv2d(128, 64, 3, padding=1), nn.ReLU(inplace=True)) self.weight = nn.Sequential( nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(128, 2, 3, padding=1), GradientClip(), nn.Sigmoid()) self.delta = nn.Sequential( nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(128, 2, 3, padding=1), GradientClip()) # Inputs to gru are: # i) hidden state (128) # ii) input state (128+128+64) # a) context_encoder_output(128), # b) corr_encoder_output (128), # c) flow_encoder_output (64), # Input planes: 128=constant_context, 128=correlation, 64=motion, 64=sigma self.gru = ConvGRU(128, 128+128+64) self.agg = GraphAgg() def forward(self, net, inp, corr, flow=None, ii=None, jj=None): """ RaftSLAM update operator """ batch, num, ch, ht, wd = net.shape if flow is None: # Initialize flow to zero... (TODO: better initialization using IMU) flow = torch.zeros(batch, num, 4, ht, wd, device=net.device) output_dim = (batch, num, -1, ht, wd) net = net.view(batch*num, -1, ht, wd) inp = inp.view(batch*num, -1, ht, wd) corr = corr.view(batch*num, -1, ht, wd) flow = flow.view(batch*num, -1, ht, wd) corr = self.corr_encoder(corr) flow = self.flow_encoder(flow) net = self.gru(net, inp, corr, flow) ### update variables ### delta = self.delta(net).view(*output_dim) weight = self.weight(net).view(*output_dim) delta = delta.permute(0,1,3,4,2)[...,:2].contiguous() weight = weight.permute(0,1,3,4,2)[...,:2].contiguous() net = net.view(*output_dim) if ii is not None: eta, upmask = self.agg(net, ii.to(net.device)) return net, delta, weight, eta, upmask else: return net, delta, weight class DroidNet(nn.Module): def __init__(self): super(DroidNet, self).__init__() self.feature_net = BasicEncoder(output_dim=128, norm_fn='instance') self.context_net = BasicEncoder(output_dim=256, norm_fn='none') self.update_net = UpdateModule() #### EVERYTHING BELOW IS ONLY USED FOR TRAINING NOT INFERENCE #### def extract_features(self, images): """ run feature extraction networks """ # normalize images b, n, c1, h1, w1 = x.shape images = images[:, :, [2,1,0]] / 255.0 mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device) std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device) images = images.sub_(mean[:, None, None]).div_(std[:, None, None]) feature_maps = self.feature_net(images) context_maps = self.context_net(images) context_maps, gru_input_maps = context_maps.split([128,128], dim=2) context_maps = torch.tanh(context_maps) gru_input_maps = torch.relu(gru_input_maps) return feature_maps, context_maps, gru_input_maps def forward(self, Gs, images, disps, intrinsics, graph=None, num_steps=12, fixedp=2): """ Estimates SE3 or Sim3 between pair of frames """ u = keyframe_indicies(graph) ii, jj, kk = graph_to_edge_list(graph) ii = ii.to(device=images.device, dtype=torch.long) jj = jj.to(device=images.device, dtype=torch.long) feature_maps, context_maps, gru_input_maps = self.extract_features(images) context_maps, gru_input_maps = context_maps[:,ii], gru_input_maps[:,ii] corr_fn = CorrBlock(feature_maps[:,ii], feature_maps[:,jj], num_levels=4, radius=3) ht, wd = images.shape[-2:] coords0 = pops.coords_grid(ht//8, wd//8, device=images.device) coords1, _ = pops.projective_transform(Gs, disps, intrinsics, ii, jj) target = coords1.clone() Gs_list, disp_list, residual_list = [], [], [] for _ in range(num_steps): Gs = Gs.detach() disps = disps.detach() coords1 = coords1.detach() target = target.detach() # extract motion features corr = corr_fn(coords1) residual = target - coords1 flow = coords1 - coords0 motion = torch.cat([flow, residual], dim=-1) motion = motion.permute(0,1,4,2,3).clamp(-64.0, 64.0) context_maps, delta, weight, eta, upmask = \ self.update_net(context_maps, gru_input_maps, corr, motion, ii, jj) target = coords1 + delta for _ in range(2): Gs, disps = BA(target, weight, eta, Gs, disps, intrinsics, ii, jj, fixedp=2) coords1, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj) residual = (target - coords1) Gs_list.append(Gs) disp_list.append(upsample_disp(disps, upmask)) residual_list.append(valid_mask * residual) return Gs_list, disp_list, residual_list