#!/usr/bin/env python3 from abc import abstractclassmethod from collections import OrderedDict from icecream import ic import cv2 import numpy as np import torch from torch import nn import torch.nn.functional as F from torch_scatter import scatter_sum from utils.flow_viz import * import networks.geom.projective_ops as pops from networks.modules.corr import CorrBlock, AltCorrBlock import lietorch from lietorch import SE3 import droid_backends import gtsam from gtsam import (HessianFactor) from gtsam import Values from gtsam import (Pose3, Rot3, Point3) from gtsam import PriorFactorPose3 from gtsam import NonlinearFactorGraph from gtsam import GaussianFactorGraph from gtsam.symbol_shorthand import X # utility functions for scattering ops def safe_scatter_add_mat(A, ii, jj, n, m): v = (ii >= 0) & (jj >= 0) & (ii < n) & (jj < m) return scatter_sum(A[:,v], ii[v]*m + jj[v], dim=1, dim_size=n*m) def lietorch_pose_to_gtsam(pose : lietorch.SE3): trans, quat = pose.vec().split([3,4], -1) trans = trans.cpu().numpy() quat = quat.cpu().numpy() return Pose3(Rot3.Quaternion(quat[3], quat[0], quat[1], quat[2]), Point3(trans)) def gtsam_pose_to_torch(pose: gtsam.Pose3, device, dtype): t = pose.translation() q = pose.rotation().quaternion() return torch.tensor([t[0], t[1], t[2], q[1], q[2], q[3], q[0]], device=device, dtype=dtype) class VisualFrontend(nn.Module): def __init__(self): super().__init__() @abstractclassmethod def forward(self, mini_batch): pass from networks.modules.extractor import BasicEncoder from networks.droid_net import UpdateModule class RaftVisualFrontend(VisualFrontend): def __init__(self, world_T_body_t0, body_T_cam0, args, device="cpu"): super().__init__() self.args = args self.kf_idx = 0 # Keyframe index self.kf_idx_to_f_idx = {} # Keyframe index to frame index self.f_idx_to_kf_idx = {} # Frame index to keyframe index self.last_kf_idx = 0 self.last_k = None self.global_ba = False self.stop = False # stop module self.compute_covariances = True self.last_state = gtsam.Values() self.initial_x0 = None self.initial_priors = None self.factors_to_remove = gtsam.KeyVector() self.buffer = args.buffer self.stereo = args.stereo self.device = device self.is_initialized = False self.keyframe_warmup = 8 self.max_age = 25 self.max_factors = 48 self.kf_init_count = 8 self.motion_filter_thresh = 2.4 # To determine if we are moving, how much mean optical flow before considering new frame [px] self.viz = False # Whether to visualize the results self.world_T_body_t0 = world_T_body_t0 self.body_t0_T_world = gtsam_pose_to_torch(self.world_T_body_t0.inverse(), self.device, torch.float) self.body_T_cam0 = body_T_cam0 self.world_T_cam0_t0 = world_T_body_t0 * body_T_cam0 self.cam0_t0_T_world = gtsam_pose_to_torch( self.world_T_cam0_t0.inverse(), self.device, torch.float) self.cam0_T_body = gtsam_pose_to_torch( body_T_cam0.inverse(), self.device, torch.float) # Frontend params self.keyframe_thresh = 4.0 # Distance to consider a keyframe, threshold to create a new keyframe [m] # why not 0.4! self.frontend_thresh = 16.0 # Add edges between frames within this distance self.frontend_window = 25 # frontend optimization window self.frontend_radius = 2 # force edges between frames within radius self.frontend_nms = 1 # non-maximal supression of edges self.beta = 0.3 # weight for translation / rotation components of flow # also used in backend # Backend params self.backend_thresh = 22.0 self.backend_radius = 2 self.backend_nms = 3 self.iters1 = 4 # number of iterations for first optimization self.iters2 = 2 # number of iterations for second optimization # DownSamplingFactor: resolution of the images with respect to the features extracted. # 8.0 means that the features are at 1/8th of the original resolution. self.dsf = 8 # perhaps the most important parameter # Type of correlation computation to use: "volume" or "alt" # "volume" takes a lot of memory (but is faster), "alt" takes less memory and should be as fast as volume but it's not self.corr_impl = "volume" # Build Networks self.feature_net = BasicEncoder(output_dim=128, norm_fn='instance') self.context_net = BasicEncoder(output_dim=256, norm_fn='none') self.update_net = UpdateModule() # Load network weights weights = self.load_weights(args.weights) missing_keys = self.load_state_dict(weights) self.to(device) self.eval() # Uncertainty sigmas, initial sigmas for initialization (but not priors?) self.translation_sigma = torch.tensor(0.01, device=self.device) # standard deviation of translation [m] self.rotation_sigma = torch.tensor(0.01, device=self.device) # standard deviation of rotation [rad] # TODO: given that the values are much larger than 1.0... we should increase this much more... self.sigma_idepth = torch.tensor(0.1, device=self.device) # standard deviation of depth [m] (or inverse depth?) [1/m], we don't know the scale anyway... self.t_cov = torch.pow(self.translation_sigma, 2) * torch.eye(3, device=self.device) self.r_cov = torch.pow(self.rotation_sigma, 2) * torch.eye(3, device=self.device) self.idepth_prior_cov = torch.pow(self.sigma_idepth, 2) self.g_prior_cov = torch.block_diag(self.r_cov, self.t_cov) # GTSAM convention, rotation first, then translation def __del__(self): print("Calling frontend dtor...") torch.cuda.empty_cache() def stop_condition(self): return self.stop # Pre-allocate all the memory in the GPU def initialize_buffers(self, image_size): self.img_height = h = image_size[0] self.img_width = w = image_size[1] ic(self.dsf) ic(h) ic(w) ic(h//self.dsf) ic(w//self.dsf) self.coords0 = pops.coords_grid(h//self.dsf, w//self.dsf, device=self.device) self.ht, self.wd = self.coords0.shape[:2] ### Input attributes ### self.cam0_timestamps = torch.zeros(self.buffer, dtype=torch.float, device=self.device).share_memory_() # TODO this should be in the euroc parser, so that we don't allocate memory without bounds self.cam0_images = torch.zeros(self.buffer, 3, h, w, dtype=torch.uint8, device=self.device).share_memory_() # TODO why not shared memory? # This is a looot of memory self.cam0_intrinsics = torch.zeros(self.buffer, 4, dtype=torch.float, device=self.device).share_memory_() self.gt_poses = torch.zeros(self.buffer, 4, 4, dtype=torch.float, device=self.device).share_memory_() self.gt_depths = torch.zeros(self.buffer, 1, h, w, dtype=torch.float, device=self.device).share_memory_() ### State attributes ### self.cam0_T_world = torch.zeros(self.buffer, 7, dtype=torch.float, device=self.device).share_memory_() self.world_T_body = torch.zeros(self.buffer, 7, dtype=torch.float, device=self.device).share_memory_() self.world_T_body_cov = torch.zeros(self.buffer, 6, 6, dtype=torch.float, device=self.device).share_memory_() self.cam0_idepths = torch.ones(self.buffer, h//self.dsf, w//self.dsf, dtype=torch.float, device=self.device).share_memory_() self.cam0_idepths_cov = torch.ones(self.buffer, h//self.dsf, w//self.dsf, dtype=torch.float, device=self.device).share_memory_() self.cam0_depths_cov = torch.ones(self.buffer, h//self.dsf, w//self.dsf, dtype=torch.float, device=self.device).share_memory_() self.cam0_idepths_sensed = torch.zeros(self.buffer, h//self.dsf, w//self.dsf, dtype=torch.float, device=self.device).share_memory_() self.cam0_idepths_up = torch.zeros(self.buffer, h, w, dtype=torch.float, device=self.device).share_memory_() # This is a looot of memory self.cam0_depths_cov_up = torch.ones(self.buffer, h, w, dtype=torch.float, device=self.device).share_memory_() # This is a looot of memory # INITIALIZE state: # - poses all to initial state transformation # - velocities all to 0 except first to initial state # - biases all to 0 except first to initial state # TODO: why not shared memory? self.cam0_T_world[:] = self.cam0_t0_T_world self.world_T_body[:] = gtsam_pose_to_torch(self.world_T_body_t0, device=self.device, dtype=torch.float) self.world_T_body_cov[:] = self.g_prior_cov * torch.eye(6, device=self.device) #* torch.as_tensor([0.001, 0.001, 0.001, 0.001, 0.001, 0.001], device=self.device)[None] self.cam0_idepths_cov *= self.idepth_prior_cov # For multi-view, we could set this to >2, but we need to know overlapping FOVs # we could do that automatically by looking at the mean flow between frames cameras = 2 if self.stereo else 1 ### Feature attributes ### Every keyframe has a feature/context/gru_input self.features_imgs = torch.zeros(self.buffer, cameras, 128, h//self.dsf, w//self.dsf, dtype=torch.half, device=self.device)#.share_memory_() self.contexts_imgs = torch.zeros(self.buffer, cameras, 128, h//self.dsf, w//self.dsf, dtype=torch.half, device=self.device)#.share_memory_() self.cst_contexts_imgs = torch.zeros(self.buffer, cameras, 128, h//self.dsf, w//self.dsf, dtype=torch.half, device=self.device)#.share_memory_() ### Correlations, Flows, and Hidden States ### Every pair of co-visible keyframes has a correlation volume, flow, and hidden state # These are created on-the-fly so we can't really pre-allocate memory self.correlation_volumes = None self.gru_hidden_states = None # initialized as context, but evolves as hidden state self.gru_contexts_input = None # initialized as context, and remains as such self.gru_estimated_flow = torch.zeros([1, 0, h//self.dsf, w//self.dsf, 2], device=self.device, dtype=torch.float) self.gru_estimated_flow_weight = torch.zeros([1, 0, h//self.dsf, w//self.dsf, 2], device=self.device, dtype=torch.float) self.damping = 1e-6 * torch.ones_like(self.cam0_idepths) # not sure what this does ### Co-visibility Graph ### self.ii = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj = torch.as_tensor([], dtype=torch.long, device=self.device) self.age = torch.as_tensor([], dtype=torch.long, device=self.device) # inactive factors self.ii_inactive = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj_inactive = torch.as_tensor([], dtype=torch.long, device=self.device) self.ii_bad = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj_bad = torch.as_tensor([], dtype=torch.long, device=self.device) self.gru_estimated_flow_inactive = torch.zeros([1, 0, h//self.dsf, w//self.dsf, 2], device=self.device, dtype=torch.float) self.gru_estimated_flow_weight_inactive = torch.zeros([1, 0, h//self.dsf, w//self.dsf, 2], device=self.device, dtype=torch.float) # For visualization, True: needs viz update, False: not changed self.viz_idx = torch.zeros(self.buffer, device=self.device, dtype=torch.bool) #@abstractclassmethod def forward(self, batch): # The output of RaftVisualFrontend is not dense optical flow # but rather a bunch of pose-to-pose factors resulting from the reduced camera matrix. print("RaftVisualFrontend.forward") k = batch["k"][0] # The output x0 = Values()# None # factors = NonlinearFactorGraph()# None # viz_out = None imgs_k = torch.as_tensor(batch["images"], device=self.device)[None].permute(0, 1, 4, 2, 3)#.shared_memory() imgs_norm_k = self._normalize_imgs(imgs_k) if self.viz: for i, img in enumerate(imgs_k[0]): cv2.imshow(f'Img{i} input', img.permute(1,2,0).cpu().numpy()) for i, img in enumerate(imgs_norm_k[0]): cv2.imshow(f'Img{i} normalized', img.permute(1,2,0).cpu().numpy()) if self.last_k is None: ic(k) assert k == 0 assert self.kf_idx == 0 assert self.last_kf_idx == 0 # Initialize network buffers self.initialize_buffers(imgs_k.shape[-2:]) # last two dims are h,w self.gt_poses[self.kf_idx] = torch.tensor(batch["poses"][0], device=self.device) if batch["depths"][0] is not None: self.gt_depths[self.kf_idx] = torch.tensor(batch["depths"][0], device=self.device).permute(2,0,1) self.cam0_timestamps[self.kf_idx] = torch.tensor(batch["t_cams"][0], device=self.device) self.cam0_images[self.kf_idx] = torch.tensor(batch["images"][0], device=self.device)[..., :3].permute(2,0,1) self.cam0_intrinsics[self.kf_idx] = (1.0 / self.dsf) * torch.tensor(batch["calibs"][0].camera_model.numpy(), device=self.device) # Initialize the state # Compute its dense features for next iteration self.features_imgs[self.kf_idx] = self.__feature_encoder(imgs_norm_k) # Compute its context features for next iteration self.contexts_imgs[self.kf_idx], self.cst_contexts_imgs[self.kf_idx] = self.__context_encoder(imgs_norm_k) # Store things for next iteration self.last_k = k self.last_kf_idx = self.kf_idx self.kf_idx_to_f_idx[self.kf_idx] = k self.f_idx_to_kf_idx[k] = self.kf_idx viz_out = self.get_viz_out(batch) self.kf_idx += 1 return x0, factors, viz_out assert k > 0 assert self.kf_idx < self.buffer # Add frame as keyframe if we have enough motion, otherwise discard: current_imgs_features = self.__feature_encoder(imgs_norm_k) if not self.has_enough_motion(current_imgs_features): if batch["is_last_frame"]: self.kf_idx -= 1 # Because in the last iter we increased it, but aren't taking any... print("Last frame reached, and no new motion: starting GLOBAL BA") self.terminate() # Send the whole viz_out to update the fact that BA has changed all poses/depths viz_out = self.get_viz_out(batch) return x0, factors, viz_out # By returning, we do not increment self.kf_idx return x0, factors, viz_out # Ok, we got enough motion, consider this frame as a keyframe # Compute dense optical flow self.gt_poses[self.kf_idx] = torch.tensor(batch["poses"][0], device=self.device) if batch["depths"][0] is not None: self.gt_depths[self.kf_idx] = torch.tensor(batch["depths"][0], device=self.device).permute(2,0,1) self.cam0_timestamps[self.kf_idx] = torch.tensor(batch["t_cams"][0], device=self.device) self.cam0_images[self.kf_idx] = torch.tensor(batch["images"][0], device=self.device)[..., :3].permute(2,0,1) self.cam0_intrinsics[self.kf_idx] = (1.0 / self.dsf) * torch.tensor(batch["calibs"][0].camera_model.numpy(), device=self.device) self.features_imgs[self.kf_idx] = current_imgs_features self.contexts_imgs[self.kf_idx], self.cst_contexts_imgs[self.kf_idx] = self.__context_encoder(imgs_norm_k) self.kf_idx_to_f_idx[self.kf_idx] = k self.f_idx_to_kf_idx[k] = self.kf_idx # Build the flow graph: ii -> jj edges # TODO: for now just do a chain # Just adds the `r' sequential frames to the graph # do initialization if not self.is_initialized: if self.kf_idx >= self.keyframe_warmup: self.__initialize() else: # We don't return here so that we increment kf_idx # ic("Warming up: pre-processing frame with enough motion") pass # do update else: if not self.__update(): # We did not accept this keyframe, reinit its properties (really needed? they will be overwritten no?) # Remove as well factors connected to it used for estimating its distance... self.rm_keyframe(self.kf_idx - 1) # TODO: the -1 here changed the whole behavior, check if it is correct # Decrease kf_idx since we are literally removing the keyframe... # But this means that we need to keep track the difference btw the kf_idx in the backend # and the keyframe_idx in the frontend. #self.video.counter.value -= 1 # self.kf_idx -= 1 # so that on the next pass we use the previous keyframe # By returning, we do not increment self.kf_idx return x0, factors, viz_out #x0.insert(X(k), pose_to_gtsam(last_pose)) self.last_k = k self.last_kf_idx = self.kf_idx self.kf_idx_to_f_idx[self.kf_idx] = k # not really necessary I think self.f_idx_to_kf_idx[k] = self.kf_idx viz_out = self.get_viz_out(batch) # build viz_out after updating self.kf_idx_to_f_idx if self.viz: cv2.waitKey(1) if self.kf_idx + 1 >= self.buffer or batch["is_last_frame"]: print("Buffer full or last frame reached: starting GLOBAL BA") self.terminate() viz_out = self.get_viz_out(batch) return x0, factors, viz_out self.kf_idx += 1 return x0, factors, viz_out # If kf0 is None, then it is init to keyframe number 1 or min(ii)+1, TODO: why not 0? # If kf1 is None, then it is init to the max(ii, jj) +1, TODO: again, why the +1? @torch.cuda.amp.autocast(enabled=True) def update(self, kf0=None, kf1=None, itrs=2, use_inactive=False, EP=1e-7, motion_only=False): """ run update operator on factor graph """ #ic("Memory usage before update: {} Mb".format(torch.cuda.memory_allocated()/1024**2)) # motion features with torch.cuda.amp.autocast(enabled=False): # try mixed precision? # Coords1 shape is: (batch_size, num_edges, ht, wd, 2) coords1, mask, (Ji, Jj, Jz) = self.reproject(self.ii, self.jj, cam_T_body=self.cam0_T_body, jacobian=True) # this is not using the cuda kernels... # mask is not used... # "coords1 - coords0": coords1 = coords0 + flow, so this is the current # flow induced by the estimated pose/depth. # "target - coords1": residual from current # - `measured` flow (target) by GRU: target = coords1 + flow_delta (from GRU), and # - `estimated' flow (coords1): coords1 = coords0 + reproject() motion = torch.cat([coords1 - self.coords0, self.gru_estimated_flow - coords1], dim=-1) motion = motion.permute(0,1,4,2,3).clamp(-64.0, 64.0) # correlation features corr = self.correlation_volumes(coords1) # We then pool the hidden state over all features which share the same source view i and predict a # pixel-wise damping factor λ. We use the softplus operator to ensure that the damping term is positive. # Additionally, we use the pooled features to predict a 8x8 mask which can be used to upsample the # inverse depth estimate self.gru_hidden_states, flow_delta, gru_estimated_flow_weight, damping, upmask = \ self.update_net(self.gru_hidden_states, self.gru_contexts_input, corr, flow=motion, ii=self.ii, jj=self.jj) if kf0 is None: # In droid: kf0 = max(1, self.ii.min().item()+1) # It is max(1, min(ii)) because the first pose is fixed. kf0 = max(0, self.ii.min().item()) else: ic(kf0) raise with torch.cuda.amp.autocast(enabled=False): # flow_delta shape is: (batch_size, num_edges, ht, wd, 2) self.gru_estimated_flow = coords1 + flow_delta.to(dtype=torch.float) self.gru_estimated_flow_weight = gru_estimated_flow_weight.to(dtype=torch.float) self.damping[torch.unique(self.ii)] = damping # TODO What is this damping? See also `damping` below if use_inactive: # It is always set to True... # TODO What is this doing? I think this is somehow setting the priors!! Or # the marginalization priors... # The thing is that, there are two considerations: # - depth_maps to optimize over (K = len(torch.unique(ii))) # - keyframe poses to optimize over (P = kf1-kf0) # where K > P, so that there are some poses we do not optimize over (but we do optimize their depth-maps) mask = (self.ii_inactive >= kf0 - 3) & (self.jj_inactive >= kf0 - 3) ii = torch.cat([self.ii_inactive[mask], self.ii], 0) jj = torch.cat([self.jj_inactive[mask], self.jj], 0) gru_estimated_flow = torch.cat([self.gru_estimated_flow_inactive[:,mask], self.gru_estimated_flow], 1) gru_estimated_flow_weight = torch.cat([self.gru_estimated_flow_weight_inactive[:,mask], self.gru_estimated_flow_weight], 1) else: ii, jj, gru_estimated_flow, gru_estimated_flow_weight = self.ii, self.jj, self.gru_estimated_flow, self.gru_estimated_flow_weight damping = .2 * self.damping[torch.unique(ii)].contiguous() + EP # TODO What is this damping? # gru_estimated_flow(_weight) shape after this line is: (num_edges, 2, ht, wd) gru_estimated_flow = gru_estimated_flow.view(-1, self.ht, self.wd, 2).permute(0,3,1,2).contiguous() gru_estimated_flow_weight = gru_estimated_flow_weight.view(-1, self.ht, self.wd, 2).permute(0,3,1,2).contiguous() # TODO We should output at this point the GRU estimated flow and weights. # Or rather the factors? # Dense bundle adjustment ic("BA!") x0, rcm_factor = self.ba(gru_estimated_flow, gru_estimated_flow_weight, damping, ii, jj, kf0, kf1, itrs=itrs, lm=1e-4, ep=0.1, motion_only=motion_only, compute_covariances=self.compute_covariances) # Stores depths_up, depths_cov_up kx = torch.unique(self.ii) self.cam0_idepths_up[kx] = cvx_upsample(self.cam0_idepths[kx].unsqueeze(-1), upmask).squeeze() self.cam0_depths_cov_up[kx] = cvx_upsample(self.cam0_depths_cov[kx].unsqueeze(-1), upmask, pow=1.0).squeeze() if self.viz: viz_idepth(self.cam0_idepths[kx], upmask) viz_idepth_sigma(self.cam0_idepths_cov[kx], upmask, fix_range=True, bg_img=self.cam0_images[kx]) viz_depth_sigma(self.cam0_depths_cov_up[kx].unsqueeze(-1).sqrt(), fix_range=True, bg_img=self.cam0_images[kx], sigma_thresh=20.0) viz_flow("gru_flow", gru_estimated_flow[-1] - self.coords0.permute(2,0,1)) reprojection_flow = (coords1 - self.coords0).squeeze().permute(0,3,1,2) viz_flow("reprojection_flow", reprojection_flow[-1]) # Viz weight if self.viz: # Visualize input image as well #for k, i in enumerate(self.ii): # self.viz_weight(gru_estimated_flow_weight[k], self.cam0_images[ii[k]]) #self.ii[-1].item(), self.jj[-1].item()) self.viz_weight(gru_estimated_flow_weight[-1], self.cam0_images[ii[-1]]) #self.ii[-1].item(), self.jj[-1].item()) # Update visualization kf1 = max(ii.max().item(), jj.max().item()) assert kf1 == self.kf_idx self.viz_idx[kf0:self.kf_idx+1] = True self.age += 1 return x0, rcm_factor @torch.cuda.amp.autocast(enabled=False) def update_lowmem(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, steps=8): """ run update operator on factor graph - reduced memory implementation """ # alternate corr implementation kfs, cameras, ch, ht, wd = self.features_imgs.shape corr_op = AltCorrBlock(self.features_imgs.view(1, kfs*cameras, ch, ht, wd)) for step in range(steps): print(f"Global BA Iteration #{step}/{steps}") with torch.cuda.amp.autocast(enabled=False): coords1, mask, _ = self.reproject(self.ii, self.jj) motion = torch.cat([coords1 - self.coords0, self.gru_estimated_flow - coords1], dim=-1) motion = motion.permute(0,1,4,2,3).clamp(-64.0, 64.0) # CONVGRU RUNS # Optimize the flow as much as possible, s = 8 for i in range(0, self.jj.max() + 1, s): # what does this do? print(f"ConvGRU Iteration #{i/s}/{(self.jj.max() + 1)/s}") v = (self.ii >= i) & (self.ii < i + s) # kind-of like a sliding optimization window iis = self.ii[v] jjs = self.jj[v] corr = corr_op(coords1[:,v], cameras * iis, cameras * jjs + (iis == jjs).long()) with torch.cuda.amp.autocast(enabled=True): # TODO: somehow the damping and upmask have weird shapes... what is going on? gru_hidden_states, flow_delta, gru_estimated_flow_weight, damping, upmask = \ self.update_net(self.gru_hidden_states[:, v], self.gru_contexts_input[:, iis], corr, motion[:, v], iis, jjs) kx = torch.unique(iis) all_kf_ids = torch.unique(torch.cat([iis, jjs], 0)) self.gru_hidden_states[:,v] = gru_hidden_states self.gru_estimated_flow[:,v] = coords1[:,v] + flow_delta.float() self.gru_estimated_flow_weight[:,v] = gru_estimated_flow_weight.float() self.damping[all_kf_ids] = damping # TODO What is this damping? See also `damping` below # Stores depths_up, depths_cov_up self.cam0_idepths_up[all_kf_ids] = cvx_upsample(self.cam0_idepths[all_kf_ids].unsqueeze(-1), upmask).squeeze() self.cam0_depths_cov_up[all_kf_ids] = cvx_upsample(self.cam0_depths_cov[all_kf_ids].unsqueeze(-1), upmask, pow=1.0).squeeze() #ii = torch.cat([self.ii_inactive[mask], self.ii], 0) damping = .2 * self.damping[torch.unique(self.ii)].contiguous() + EP gru_estimated_flow = self.gru_estimated_flow.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous() gru_estimated_flow_weight = self.gru_estimated_flow_weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous() # dense bundle adjustment ic("Global BA!") # TODO: do not compute cov for global BA until we fix the memory issue when building Eiz x0, rcm_factor = self.ba(gru_estimated_flow, gru_estimated_flow_weight, damping, self.ii, self.jj, kf0=0, kf1=None, itrs=itrs, lm=1e-5, ep=1e-2, motion_only=False, compute_covariances=False) @torch.cuda.amp.autocast(enabled=True) def rm_keyframe(self, kf_idx): """ drop nodes from factor graph """ # TODO: how does this work if we rm a keyframe that is not the one before the last one?? # As of now, kf_idx is the last one, and so kf_idx+1 is all 0s. self.gt_poses[kf_idx] = self.gt_poses[kf_idx+1] self.gt_depths[kf_idx] = self.gt_depths[kf_idx+1] self.cam0_images[kf_idx] = self.cam0_images[kf_idx+1] self.cam0_timestamps[kf_idx] = self.cam0_timestamps[kf_idx+1] self.cam0_T_world[kf_idx] = self.cam0_T_world[kf_idx+1] self.world_T_body[kf_idx] = self.world_T_body[kf_idx+1] self.world_T_body_cov[kf_idx] = self.world_T_body_cov[kf_idx+1] self.cam0_idepths[kf_idx] = self.cam0_idepths[kf_idx+1] self.cam0_idepths_cov[kf_idx] = self.cam0_idepths_cov[kf_idx+1] self.cam0_depths_cov[kf_idx] = self.cam0_depths_cov[kf_idx+1] self.cam0_idepths_sensed[kf_idx] = self.cam0_idepths_sensed[kf_idx+1] self.cam0_intrinsics[kf_idx] = self.cam0_intrinsics[kf_idx+1] self.features_imgs[kf_idx] = self.features_imgs[kf_idx+1] self.contexts_imgs[kf_idx] = self.contexts_imgs[kf_idx+1] self.cst_contexts_imgs[kf_idx] = self.cst_contexts_imgs[kf_idx+1] # Remove all inactive edges that are connected to the keyframe mask = (self.ii_inactive == kf_idx) | (self.jj_inactive == kf_idx) # Reindex the inactive edges that we are going to keep self.ii_inactive[self.ii_inactive >= kf_idx] -= 1 self.jj_inactive[self.jj_inactive >= kf_idx] -= 1 # Remove the inactive edges concerning this keyframe if torch.any(mask): self.ii_inactive = self.ii_inactive[~mask] self.jj_inactive = self.jj_inactive[~mask] self.gru_estimated_flow_inactive = self.gru_estimated_flow_inactive[:,~mask] self.gru_estimated_flow_weight_inactive = self.gru_estimated_flow_weight_inactive[:,~mask] # Remove all edges that are connected to the keyframe mask = (self.ii == kf_idx) | (self.jj == kf_idx) # Reindex the edges that we are going to keep self.ii[self.ii >= kf_idx] -= 1 self.jj[self.jj >= kf_idx] -= 1 # Remove the data concerning this keyframe (correlation volumes, etc.) self.rm_factors(mask, store=False) def __update(self): """ add edges, perform update """ # self.count += 1 # TODO #self.kf_idx += 1 # TODO I think this is our kf_idx if self.correlation_volumes is not None: # Really only drops edges... #ic("Removing factors, and storing.") self.rm_factors(self.age > self.max_age, store=True) # TODO: unclear how this works # t = self.kf_idx # ix = torch.arange(kf0, t) # jx = torch.arange(kf1, t) #ic("Adding proximity factors") self.add_proximity_factors(kf0=self.kf_idx - 4, kf1=max(self.kf_idx + 1 - self.frontend_window, 0), rad=self.frontend_radius, nms=self.frontend_nms, thresh=self.frontend_thresh, beta=self.beta, remove=True) # Initialize current cam0_depths with the sensed depths if valid, # otherwise with the previous mean depth (during init it is the mean amongst last 4 frames) # (perhaps try with the previous raw depth) self.cam0_idepths[self.kf_idx] = torch.where(self.cam0_idepths_sensed[self.kf_idx] > 0, self.cam0_idepths_sensed[self.kf_idx], self.cam0_idepths[self.kf_idx]) # TODO: should we lower the sigmas of idepth here since it is a measurement #ic("First update") for itr in range(self.iters1): x0, rcm_factor= self.update(kf0=None, kf1=None, use_inactive=True) # Get distance between the previous keyframe (kf_idx-2) and the current frame (kf_idx-1) d = self.distance([self.kf_idx-2], [self.kf_idx-1], beta=self.beta, bidirectional=True) #ic(d.item()) if d.item() < self.keyframe_thresh: #ic("Not a keyframe") return False else: #ic("Found a keyframe") #ic("Second update") for itr in range(self.iters2): x0, rcm_factor = self.update(None, None, use_inactive=True) # TODO: I believe this should be inside this conditional rather than outside, # Because in the previous case we decided not to accept the keyframe... # set pose for next iteration next_kf = self.kf_idx + 1 if next_kf < self.buffer: # if we have not reached the end of the buffer (aka end of sequence) self.cam0_T_world[next_kf] = self.cam0_T_world[self.kf_idx] self.world_T_body[next_kf] = self.world_T_body[self.kf_idx] self.world_T_body_cov[next_kf] = self.world_T_body_cov[self.kf_idx] # Why not just the previous depths as in init? #self.cam0_idepths[next_kf] = self.cam0_idepths[self.kf_idx] self.cam0_idepths[next_kf] = self.cam0_idepths[self.kf_idx].mean() self.cam0_idepths_cov[next_kf] = self.cam0_idepths_cov[self.kf_idx] self.cam0_depths_cov[next_kf] = self.cam0_depths_cov[self.kf_idx] #self.viz_idx[next_kf] = True return True def __initialize(self): """ initialize the SLAM system """ assert(self.kf_idx > 4) assert(self.kf_idx >= self.keyframe_warmup) # Just adds the `radius' sequential frames to the graph self.add_neighborhood_factors(kf0=0, kf1=self.kf_idx, radius=3) for _ in range(8): x0, rcm_factor = self.update(kf0=None, kf1=None, use_inactive=True) # Adds factors between frames, but unsure how this actually works... # if kf0 and kf1 are 0, then it tries to add prox factors btw all keyframes # t = self.kf_idx # ix = torch.arange(kf0, t) # jx = torch.arange(kf1, t) #ic("Add proximity factors") self.add_proximity_factors(kf0=0, kf1=0, rad=2, nms=2, thresh=self.frontend_thresh, remove=False) for _ in range(8): x0, rcm_factor = self.update(kf0=None, kf1=None, use_inactive=True) # TODO: next kf_idx shouldn't be kf_idx+1? # Set initial pose/depth for next iteration self.cam0_T_world[self.kf_idx + 1] = self.cam0_T_world[self.kf_idx].clone() self.world_T_body[self.kf_idx + 1] = self.world_T_body[self.kf_idx].clone() self.world_T_body_cov[self.kf_idx + 1] = self.world_T_body_cov[self.kf_idx].clone() # TODO: Next depth is just the mean of the previous 4 depths? # We just retrieve our global implicit map here.... self.cam0_idepths[self.kf_idx + 1] = self.cam0_idepths[self.kf_idx - 3:self.kf_idx+1].mean() # TODO: here we are doing something very wrong and it is to previous sigmas and worst the initial sigma which is just one # because I don't have a better initialization... but looking at the numbers one is very confident... self.cam0_idepths_cov[self.kf_idx + 1] = self.cam0_idepths_cov[self.kf_idx-3:self.kf_idx+1].mean() self.cam0_depths_cov[self.kf_idx + 1] = self.cam0_depths_cov[self.kf_idx-3:self.kf_idx+1].mean() # initialization complete self.is_initialized = True # Update visualization self.viz_idx[:self.kf_idx+1] = True # TODO: what is this 4 here? What if warmup is less than 4... # Remove edges that point to the first keyframes... #ic("Remove factors after init, and storing") self.rm_factors(self.ii < (self.keyframe_warmup - 4), store=True) return x0, rcm_factor def add_neighborhood_factors(self, kf0, kf1, radius=3): """ add edges between neighboring frames within radius r """ # Build dense adjacency matrix, +1 bcs arange stops but does not include the last index. ii, jj = torch.meshgrid(torch.arange(kf0, kf1+1), torch.arange(kf0, kf1+1)) ii = ii.reshape(-1).to(dtype=torch.long, device=self.device) jj = jj.reshape(-1).to(dtype=torch.long, device=self.device) c = 1 if self.stereo else 0 # Remove from dense adjacency matrix those that are not within `r' frames of each other # This basically keeps the `r' sub-diagonals at the top/bottom of the diagonal. distances = torch.abs(ii - jj) keep_radius = distances <= radius keep_stereo = distances > c # TODO: not sure why this is like this... keep = keep_stereo & keep_radius # Computes the correlation between these frames self.add_factors(ii[keep], jj[keep]) # TODO: not sure how this really works... def add_proximity_factors(self, kf0=0, kf1=0, rad=2, nms=2, beta=0.25, thresh=16.0, remove=False): """ add edges to the factor graph based on distance """ t = self.kf_idx + 1 ix = torch.arange(kf0, t) jx = torch.arange(kf1, t) ii, jj = torch.meshgrid(ix, jx) ii = ii.reshape(-1) jj = jj.reshape(-1) d = self.distance(ii, jj, beta=beta) d[(ii - rad) < jj] = np.inf # Set closer than rad frames distance to infinity d[d > 100] = np.inf # Set distances greater than 100 to infinity ii1 = torch.cat([self.ii, self.ii_bad, self.ii_inactive], 0) jj1 = torch.cat([self.jj, self.jj_bad, self.jj_inactive], 0) for i, j in zip(ii1.cpu().numpy(), jj1.cpu().numpy()): for di in range(-nms, nms+1): for dj in range(-nms, nms+1): if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0): i1 = i + di j1 = j + dj if (kf0 <= i1 < t) and (kf1 <= j1 < t): d[(i1-kf0)*(t-kf1) + (j1-kf1)] = np.inf es = [] for i in range(kf0, t): if self.stereo: es.append((i, i)) d[(i-kf0)*(t-kf1) + (i-kf1)] = np.inf for j in range(max(i-rad-1,0), i): es.append((i,j)) es.append((j,i)) d[(i-kf0)*(t-kf1) + (j-kf1)] = np.inf ix = torch.argsort(d) for k in ix: if d[k].item() > thresh: continue if len(es) > self.max_factors: break i = ii[k] j = jj[k] # bidirectional es.append((i, j)) es.append((j, i)) for di in range(-nms, nms+1): for dj in range(-nms, nms+1): if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0): i1 = i + di j1 = j + dj if (kf0 <= i1 < t) and (kf1 <= j1 < t): d[(i1-kf0)*(t-kf1) + (j1-kf1)] = np.inf ii, jj = torch.as_tensor(es, device=self.device).unbind(dim=-1) self.add_factors(ii, jj, remove) def distance(self, ii=None, jj=None, beta=0.3, bidirectional=True): """ frame distance metric """ return_distance_matrix = False if ii is None: return_distance_matrix = True N = self.kf_idx ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N)) ii, jj = self.format_indicies(ii, jj, device=self.device) if bidirectional: poses = self.cam0_T_world[:self.kf_idx+1].clone() # TODO: why clone? why +1? d1 = droid_backends.frame_distance(poses, self.cam0_idepths, self.cam0_intrinsics[0], ii, jj, beta) d2 = droid_backends.frame_distance(poses, self.cam0_idepths, self.cam0_intrinsics[0], jj, ii, beta) d = .5 * (d1 + d2) else: d = droid_backends.frame_distance(self.cam0_T_world, self.cam0_idepths, self.cam0_intrinsics[0], ii, jj, beta) if return_distance_matrix: return d.reshape(N, N) return d # Computes correlations for each edge in (ii, jj) # Inits gru_hidden_states for each edge from self.contexts_imgs # Computes the gru_context_input from the cst_contexts_imgs # Inits gru_estimated_flow to the reprojected flow # Inits gru_estimated_flow_weight to 0.0 @torch.cuda.amp.autocast(enabled=True) def add_factors(self, ii, jj, remove=False): """ add edges to factor graph """ if not isinstance(ii, torch.Tensor): ii = torch.as_tensor(ii, dtype=torch.long, device=self.device) if not isinstance(jj, torch.Tensor): jj = torch.as_tensor(jj, dtype=torch.long, device=self.device) # remove duplicate edges, duplication happens because neighborhood and proximity factors # may overlap. ii, jj = self.__filter_repeated_edges(ii, jj) if ii.shape[0] == 0: return # place limit on number of factors old_factors_count = self.ii.shape[0] new_factors_count = ii.shape[0] if self.max_factors > 0 and \ old_factors_count + new_factors_count > self.max_factors \ and self.correlation_volumes is not None and remove: ix = torch.arange(len(self.age))[torch.argsort(self.age).cpu()] #ic("Remove old factors, and storing.") self.rm_factors(ix >= (self.max_factors - new_factors_count), store=True) ### Add new factors self.ii = torch.cat([self.ii, ii], 0) self.jj = torch.cat([self.jj, jj], 0) self.age = torch.cat([self.age, torch.zeros_like(ii)], 0) ### Add correlation volumes for new edges, if we do not use alt implementation # that computes correlations on the fly... (a bit slower for some reason, but saves # tons of memory, since most of the corr volume will be unexplored I think) if self.corr_impl == "volume": is_stereo = (ii == jj).long() # TODO: does this still hold for multi-camera? feature_img_ii = self.features_imgs[None, ii, 0] feature_img_jj = self.features_imgs[None, jj, is_stereo] corr = CorrBlock(feature_img_ii, feature_img_jj) self.correlation_volumes = self.correlation_volumes.cat(corr) \ if self.correlation_volumes is not None else corr ### Gru hidden states are initialized to the context features, and then they evolve gru_hidden_state = self.contexts_imgs[None, ii, 0] # Only for cam0 self.gru_hidden_states = torch.cat([self.gru_hidden_states, gru_hidden_state], 1) \ if self.gru_hidden_states is not None else gru_hidden_state ### Gru input states are initialized to the context features, and they do not evolve gru_context_input = self.cst_contexts_imgs[None, ii, 0] # Only for cam0 self.gru_contexts_input = torch.cat([self.gru_contexts_input, gru_context_input], 1) \ if self.gru_contexts_input is not None else gru_context_input ### Gru estimated flow is initialized to the reprojected flow, and then it evolves with torch.cuda.amp.autocast(enabled=False): target, _, _ = self.reproject(ii, jj) weight = torch.zeros_like(target) # Initialize weights to 0 # TODO: not sure why we concat gru_estimated_flow with target, instead of directly init with target... # The first gru_estimated_flow is init to zero! self.gru_estimated_flow = torch.cat([self.gru_estimated_flow, target], 1) ## Init gru flow with the one from reprojection! self.gru_estimated_flow_weight = torch.cat([self.gru_estimated_flow_weight, weight], 1) # Frees up memory as well, # TODO it doesn't do any marginalization? # TODO should we update kf_idx to re-use the empty slots? @torch.cuda.amp.autocast(enabled=True) def rm_factors(self, mask, store=False): """ drop edges from factor graph """ # store estimated factors that are removed if store: self.ii_inactive = torch.cat([self.ii_inactive, self.ii[mask]], 0) self.jj_inactive = torch.cat([self.jj_inactive, self.jj[mask]], 0) self.gru_estimated_flow_inactive = torch.cat([self.gru_estimated_flow_inactive, self.gru_estimated_flow[:,mask]], 1) self.gru_estimated_flow_weight_inactive = torch.cat([self.gru_estimated_flow_weight_inactive, self.gru_estimated_flow_weight[:,mask]], 1) # Actually remove edges self.ii = self.ii[~mask] self.jj = self.jj[~mask] self.age = self.age[~mask] if self.corr_impl == "volume": self.correlation_volumes = self.correlation_volumes[~mask] if self.gru_hidden_states is not None: self.gru_hidden_states = self.gru_hidden_states[:,~mask] if self.gru_contexts_input is not None: self.gru_contexts_input = self.gru_contexts_input[:,~mask] self.gru_estimated_flow = self.gru_estimated_flow[:,~mask] self.gru_estimated_flow_weight = self.gru_estimated_flow_weight[:,~mask] # TODO: what about those keyframes that are isolated because no more factors touch them? def __filter_repeated_edges(self, ii, jj): """ remove duplicate edges """ keep = torch.zeros(ii.shape[0], dtype=torch.bool, device=ii.device) eset = set( [(i.item(), j.item()) for i, j in zip(self.ii, self.jj)] + [(i.item(), j.item()) for i, j in zip(self.ii_inactive, self.jj_inactive)]) for k, (i, j) in enumerate(zip(ii, jj)): keep[k] = (i.item(), j.item()) not in eset return ii[keep], jj[keep] def reproject(self, ii, jj, cam_T_body=None, jacobian=False): """ project points from ii -> jj """ ii, jj = self.format_indicies(ii, jj, device=self.device) # TODO: is this really necessary? Gs = lietorch.SE3(self.cam0_T_world[None]) # TODO: this is a bit expensive no? # TODO: It would be great to visualize both coords, valid_mask, (Ji, Jj, Jz) = \ pops.projective_transform(Gs, self.cam0_idepths[None], self.cam0_intrinsics[None], ii, jj, cam_T_body=cam_T_body, jacobian=jacobian) return coords, valid_mask, (Ji, Jj, Jz) def print_edges(self): ii = self.ii.cpu().numpy() jj = self.jj.cpu().numpy() ix = np.argsort(ii) ii = ii[ix] jj = jj[ix] w = torch.mean(self.gru_estimated_flow_weight, dim=[0,2,3,4]).cpu().numpy() w = w[ix] for e in zip(ii, jj, w): print(e) print() @staticmethod def format_indicies(ii, jj, device="cpu"): """ to device, long, {-1} """ if not isinstance(ii, torch.Tensor): ii = torch.as_tensor(ii) if not isinstance(jj, torch.Tensor): jj = torch.as_tensor(jj) ii = ii.to(device=device, dtype=torch.long).reshape(-1) jj = jj.to(device=device, dtype=torch.long).reshape(-1) return ii, jj @torch.cuda.amp.autocast(enabled=True) def __context_encoder(self, images): # image must be normalized """ Context features """ context_maps, gru_input_maps = self.context_net(images).split([128,128], dim=2) return context_maps.tanh().squeeze(0), gru_input_maps.relu().squeeze(0) @torch.cuda.amp.autocast(enabled=True) def __feature_encoder(self, images): # image must be normalized """ Features for correlation volume """ return self.feature_net(images).squeeze(0) # images has shape [c, 3, h, w], where c is the number of images, h and w are the height and width of the images. # It outputs the images with an extra dimension at the front, the batch dimension b. # Further, it sends them to the GPU. def _normalize_imgs(self, images, droid_normalization=True): img_normalized = images[:,:,:3, ...] / 255.0 # Drop alpha channel if droid_normalization: mean = torch.as_tensor([0.485, 0.456, 0.406], device=self.device)[:, None, None] stdv = torch.as_tensor([0.229, 0.224, 0.225], device=self.device)[:, None, None] else: mean = img_normalized.mean(dim=(3,4), keepdim=True) stdv = img_normalized.std(dim=(3,4), keepdim=True) img_normalized = img_normalized.sub_(mean).div_(stdv) return img_normalized @torch.cuda.amp.autocast(enabled=True) @torch.no_grad() def has_enough_motion(self, current_imgs_features): # Only calculates if enough motion by looking at cam0 ic(self.last_kf_idx) last_img_features = self.features_imgs[self.last_kf_idx][0] current_img_features = current_imgs_features[0] last_img_context = self.contexts_imgs[self.last_kf_idx][0] last_img_gru_input = self.cst_contexts_imgs[self.last_kf_idx][0] # Viz hidden features if self.viz: self.viz_hidden_features(last_img_features, current_img_features) # Index correlation volume corr = CorrBlock(last_img_features[None, None], current_img_features[None, None])(self.coords0[None, None]) # TODO why not send the corr block? # Approximate flow magnitude using 1 update iteration _, delta, weight = self.update_net(last_img_context[None,None], last_img_gru_input[None,None], corr) # # Viz weight # if self.viz: # self.viz_weight(weight) # # Viz flow # if self.viz: # viz_flow(delta)#, flow_norm=100) # Check motion magnitude / add new frame to video has_enough_motion = delta.norm(dim=-1).mean().item() > self.motion_filter_thresh return has_enough_motion # Visulize the actual RGB image underneath! def viz_weight(self, weight, img_bg=None, ii=0, jj=0, write=True): weight = weight.permute(1,2,0) # Calculate the max norm of the weight image #weight_norm_max = weight.norm(dim=-1).max().item() #ic(weight_norm_max) #weight = weight / weight_norm_max weight_x = weight[...,0][None,None] weight_y = weight[...,1][None,None] # TODO should we multiply by 8 as in upflow8? new_size =(self.img_height, self.img_width) viz_weight_x = F.interpolate(weight_x, size=new_size, mode="bilinear", align_corners=True).squeeze() viz_weight_y = F.interpolate(weight_y, size=new_size, mode="bilinear", align_corners=True).squeeze() viz_weight_x = (viz_weight_x.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit viz_weight_y = (viz_weight_y.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit viz_weight_x = cv2.applyColorMap(viz_weight_x, cv2.COLORMAP_MAGMA) viz_weight_y = cv2.applyColorMap(viz_weight_y, cv2.COLORMAP_MAGMA) if img_bg is not None: img_bg = img_bg.permute(1,2,0).cpu().numpy() viz_weight_x = cv2.addWeighted(img_bg, 0.5, viz_weight_x, 0.5, 0) viz_weight_y = cv2.addWeighted(img_bg, 0.5, viz_weight_y, 0.5, 0) name_x = f'Weights X from {ii} to {jj}' name_y = f'Weights Y from {ii} to {jj}' cv2.imshow(name_x, viz_weight_x) cv2.imshow(name_y, viz_weight_y) if write: cv2.imwrite(name_x+".png", viz_weight_x) cv2.imwrite(name_y+".png", viz_weight_y) def viz_hidden_features(self, last_img_features, current_imgs_features): lif = last_img_features.norm(dim=0) cif = current_imgs_features.norm(dim=0) lif = lif / lif.max() cif = cif / cif.max() new_size =(self.img_height, self.img_width) lif = F.interpolate(lif[None,None], size=new_size, mode="bilinear", align_corners=True) cif = F.interpolate(cif[None,None], size=new_size, mode="bilinear", align_corners=True) lif = lif.squeeze() cif = cif.squeeze() cv2.imshow('Last Img Features', lif.cpu().numpy()) cv2.imshow('Current Img Features', cif.cpu().numpy()) def load_weights(self, weights_file): """ load trained model weights """ state_dict = OrderedDict([ (k.replace("module.", ""), v) for (k, v) in torch.load(weights_file).items()]) state_dict = OrderedDict([ (k.replace("fnet.", "feature_net."), v) for (k, v) in state_dict.items()]) state_dict = OrderedDict([ (k.replace("cnet.", "context_net."), v) for (k, v) in state_dict.items()]) state_dict = OrderedDict([ (k.replace("update.", "update_net."), v) for (k, v) in state_dict.items()]) state_dict["update_net.weight.2.weight"] = state_dict["update_net.weight.2.weight"][:2] state_dict["update_net.weight.2.bias"] = state_dict["update_net.weight.2.bias"][:2] state_dict["update_net.delta.2.weight"] = state_dict["update_net.delta.2.weight"][:2] state_dict["update_net.delta.2.bias"] = state_dict["update_net.delta.2.bias"][:2] return state_dict # TODO: I changed from kf0=1 to kf0=0 def ba(self, gru_estimated_flow, gru_estimated_flow_weight, damping, ii, jj, kf0=0, kf1=None, itrs=2, lm=1e-4, ep=0.1, motion_only=False, compute_covariances=True): """ dense bundle adjustment (DBA): gets reduced camera matrix, solves it using gtsam, calculates covariances """ if kf1 is None: # [kf0, kf1] window of bundle adjustment optimization kf1 = max(ii.max().item(), jj.max().item()) + 1 N = kf1 - kf0 HW = self.ht * self.wd kx = torch.unique(ii) kf_ids = [i+kf0 for i in range(kf1 - kf0)] f_ids = [self.kf_idx_to_f_idx[kf_id] for kf_id in kf_ids] Xii = np.array([X(f_id) for f_id in f_ids]) initial_priors = None if f_ids[0] == 0: # Add strong prior if self.world_T_cam0_t0: _, initial_priors = self.get_gt_priors_and_values(kf_ids[0], f_ids[0]) else: raise ("You need to add initial prior, or you'll have ill-cond hessian!") for _ in range(itrs): x0 = Values() linear_factor_graph = GaussianFactorGraph() for i in range(N): kf_id = i + kf0 x0.insert(Xii[i], gtsam.Pose3(SE3(self.world_T_body[kf_id]).matrix().cpu().numpy())) #ic(f"LINEARIZE iter: {iteration}") # At this point, cam0_poses and cam0_depths have been modified if ba() is called... # TODO: the saved self.* are to update the depths after GTSAM solve. # ic(target/weight.shape) ([42, 2, 44, 69]) # Get H_vision, linearize at current Pose, Depth H, v, Q, E, w = droid_backends.reduced_camera_matrix( self.cam0_T_world, self.world_T_body, # TODO(remove, unnecessary, given the math) self.cam0_idepths, self.cam0_intrinsics[0], self.cam0_T_body, self.cam0_idepths_sensed, gru_estimated_flow, gru_estimated_flow_weight, # TODO: we should pass as well previous pose covariances!, so to not lose information when optimizing damping, ii, jj, kf0, kf1) vision_factors = GaussianFactorGraph() H = torch.nn.functional.unfold(H[None,None], (6, 6), stride=6).permute(2,0,1).view(N, N, 6, 6) v = torch.nn.functional.unfold(v[None,None], (6, 1), stride=6).permute(2,0,1).view(N, 6) H[range(N), range(N)] /= N # because we add N times the diagonal factors... v[:] /= N upper_triangular_indices = torch.triu_indices(N, N) # TODO: this is quite slow, we should implement it in c++/CUDA for i, j in zip(upper_triangular_indices[0], upper_triangular_indices[1]): if i == j: vision_factors.add(HessianFactor(Xii[i], H[i, i].cpu().numpy(), v[i].cpu().numpy(), 0.0)) else: vision_factors.add(HessianFactor(Xii[i], Xii[j], H[i, i].cpu().numpy(), H[i, j].cpu().numpy(), v[i].cpu().numpy(), H[j, j].cpu().numpy(), v[j].cpu().numpy(), 0.0)) linear_factor_graph.push_back(vision_factors) # Add Factors if initial_priors is not None: ic("ADDING initial prior!") linear_factor_graph.push_back(initial_priors.linearize(x0)) # SOLVE # TODO: Droid does LM though, for which we should add to the diagonal # delta above has our delta, but I use GTSAM here because I don't want to convert delta to VectorValues... gtsam_delta = linear_factor_graph.optimizeDensely() # Calls Eigen Cholesky, without a particularly smart ordering (vs Eigen::SimplicialLLt...) self.last_state = x0.retract(gtsam_delta) # this retraction requires local right-hand deltas (chi) # RETRACT # Update world to Body Poses poses = gtsam.utilities.allPose3s(self.last_state) pose_keys = poses.keys() for i, key in enumerate(pose_keys): f_idx = gtsam.Symbol(key).index() kf_idx = self.f_idx_to_kf_idx[f_idx] self.world_T_body[kf_idx] = gtsam_pose_to_torch(poses.atPose3(key), device=self.device, dtype=torch.float) # Update Cam to world Poses self.cam0_T_world[:] = (SE3(self.cam0_T_body[None]) * SE3(self.world_T_body).inv()).vec() # Refresh cam0_poses given retracted body_poses # Retract Depths xi_delta = torch.as_tensor(gtsam_delta.vector(pose_keys), device=self.device, dtype=torch.float).view(-1, 6) droid_backends.solve_depth(xi_delta, self.cam0_idepths, Q, E, w, ii, jj, kf0, kf1) self.cam0_idepths.clamp_(min=0.001) if compute_covariances: H, v = linear_factor_graph.hessian() L = None try: L = torch.linalg.cholesky(torch.as_tensor(H, device=self.device, dtype=torch.float))# from double to float... except Exception as e: print(e) if L is not None: identity = torch.eye(L.shape[0], device=L.device) # L has shape (PD,PD) L_inv = torch.linalg.solve_triangular(L, identity, upper=False) if torch.isnan(L_inv).any(): print("NANs in L_inv!!") raise # We only care about block diagonals of sigma_g though, here we are calculating everything... sigma_gg = L_inv.transpose(-2,-1) @ L_inv # TODO: this is the same as optimizeDensely in gtsam.... # delta = sigma @ torch.as_tensor(v, device=self.device, dtype=torch.float) # Calculate sigmas # Extract only the block-diagonal of size D from sigma_g P = N D = L.shape[0] // P assert D == 6 sigma_gg = sigma_gg.view(P, D, P, D).permute(0,2,1,3) # P x P x D x D sigma_g = torch.diagonal(sigma_gg, dim1=0, dim2=1).permute(2,0,1).view(P, D, D) # P x D x D Ei = E[:P] Ejz = E[P:P+ii.shape[0]] M = Ejz.shape[0] assert M == ii.shape[0] kx, kk = torch.unique(ii, return_inverse=True) K = kx.shape[0] # !!!! Aixo es different de N o rather de P i probablement K = P + fixed_poses, so K>P! # Ejz contains all the psi(D)*z(HW) pairs of products (M in total) # These are populating the off-diagonal of E, # The diagonal is populated by Ei min_ii_jj = min(ii.min(),jj.min()) #ic(K) #ic(P) Ej = torch.zeros(K, K, D, HW, device=self.device) # HUGE MEMORY CONSUMPTION, this should be P, K, D, HW # The equation is E[jj[m], ii[m]] = Ejz[m] for m in M, but if we take into account the fixed poses, and indices, we get: Ej[jj - min_ii_jj, ii - min_ii_jj] = Ejz Ej = Ej[kf0-min_ii_jj:kf1-min_ii_jj].view(P,K,D,HW) # Keep only the keyframes we are optimizing over, and remove the fixed ones, but add all the depth-maps... # The diagonal is populated by Ei Ej[range(P), kf0-min_ii_jj:kf1-min_ii_jj, :, :] = Ei[range(P), :, :] E_sum = Ej E_sum = E_sum.view(P, K, D, HW) E_sum = E_sum.permute(0,2,1,3).reshape(P*D, K*HW) Q_ = Q.view(K*HW,1) F = torch.matmul(Q_ * E_sum.t(), L_inv) # K*HW x D*P F2 = torch.pow(F, 2) delta_cov = F2.sum(dim=-1) # K*HW z_cov = Q_.squeeze() + delta_cov # K*HW z_cov = z_cov.view(K, self.ht, self.wd) # Update body sigmas: for i, key in enumerate(pose_keys): f_idx = gtsam.Symbol(key).index() kf_idx = self.f_idx_to_kf_idx[f_idx] self.world_T_body_cov[kf_idx] = sigma_g[i] # WARNING: this assumes that poses and sigma_g are in the same order!! # Update idepths_sigma, clamp? self.cam0_idepths_cov[kx] = z_cov # Update depths_sigma, clamp? depth_cov = z_cov / self.cam0_idepths[kx]**4 self.cam0_depths_cov[kx] = depth_cov return x0, linear_factor_graph def get_gt_priors_and_values(self, kf_id, f_id): gt_pose = self.world_T_cam0_t0 #gt_pose = np.eye(4) # Get prior pose_key = X(f_id) pose_noise = gtsam.noiseModel.Diagonal.Sigmas( np.array([0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001])) # rot, trans pose_prior = PriorFactorPose3(pose_key, gtsam.Pose3(gt_pose), pose_noise) # Get guessed values x0 = Values() x0.insert(pose_key, gtsam.Pose3(gt_pose)) # Add factors graph = NonlinearFactorGraph() graph.push_back(pose_prior) return x0, graph def backend(self, steps=12): """ main update """ if not self.stereo and not torch.any(self.cam0_idepths_sensed): # TODO: we should normalize anyway if we have sensed idepths that are from the volume... #ic("Normalizing video.") self.normalize(self.kf_idx) self.max_factors = 16 * self.kf_idx self.corr_impl = "alt" self.use_uncertainty = False self.ii = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj = torch.as_tensor([], dtype=torch.long, device=self.device) self.age = torch.as_tensor([], dtype=torch.long, device=self.device) self.correlation_volumes = None self.gru_hidden_states = None # initialized as context, but evolves as hidden state self.gru_contexts_input = None # initialized as context, and remains as such self.damping = 1e-6 * torch.ones_like(self.cam0_idepths) self.gru_estimated_flow = torch.zeros([1, 0, self.ht, self.wd, 2], device=self.device, dtype=torch.float) self.gru_estimated_flow_weight = torch.zeros([1, 0, self.ht, self.wd, 2], device=self.device, dtype=torch.float) # inactive factors self.ii_inactive = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj_inactive = torch.as_tensor([], dtype=torch.long, device=self.device) self.ii_bad = torch.as_tensor([], dtype=torch.long, device=self.device) self.jj_bad = torch.as_tensor([], dtype=torch.long, device=self.device) self.gru_estimated_flow_inactive = torch.zeros([1, 0, self.ht, self.wd, 2], device=self.device, dtype=torch.float) self.gru_estimated_flow_weight_inactive = torch.zeros([1, 0, self.ht, self.wd, 2], device=self.device, dtype=torch.float) self.add_proximity_factors(rad=self.backend_radius, nms=self.backend_nms, thresh=self.backend_thresh, beta=self.beta) # uses self.max_factors to limit number of factors self.update_lowmem(steps=steps) self.clear_edges() self.viz_idx[:self.kf_idx] = True def clear_edges(self): self.rm_factors(self.ii >= 0) self.gru_hidden_states = None # initialized as context, but evolves as hidden state self.gru_contexts_input = None # initialized as context, and remains as such def normalize(self, last_kf=-1): """ normalize depth and poses """ s = self.cam0_idepths[:last_kf].mean() # mean inverse depth? whould it be the mean of the actual depth?? self.cam0_idepths[:last_kf] /= s self.cam0_T_world[:last_kf, :3] *= s self.viz_idx[:last_kf] = True def terminate(self, stream=None): """ terminate the visualization process, return poses [t, q] """ #del self.frontend # TODO: maybe return between backend calls to see the updated geometry? if self.global_ba: torch.cuda.empty_cache() print("#" * 32) self.backend(7) torch.cuda.empty_cache() print("#" * 32) self.backend(12) torch.cuda.empty_cache() print("#" * 32) print("Clear memory") self.backend(0) torch.cuda.empty_cache() print("DONE WITH GLOBAL BA") else: print("Not running global BA...") torch.cuda.empty_cache() self.stop = True def get_viz_out(self, batch): viz_index, = torch.where(self.viz_idx) viz_out = None if len(viz_index) != 0: cam0_T_world = torch.index_select(self.cam0_T_world, 0, viz_index) gt_poses = torch.index_select(self.gt_poses, 0, viz_index) gt_depths = torch.index_select(self.gt_depths, 0, viz_index) world_T_body = torch.index_select(self.world_T_body, 0, viz_index) world_T_body_cov = torch.index_select(self.world_T_body_cov, 0, viz_index) idepths = torch.index_select(self.cam0_idepths, 0, viz_index) idepths_up = torch.index_select(self.cam0_idepths_up, 0, viz_index) idepths_sensed = torch.index_select(self.cam0_idepths_sensed, 0, viz_index) idepths_cov = torch.index_select(self.cam0_idepths_cov, 0, viz_index) depths_cov = torch.index_select(self.cam0_depths_cov, 0, viz_index) depths_cov_up = torch.index_select(self.cam0_depths_cov_up, 0, viz_index) # do not use up images = torch.index_select(self.cam0_images, 0, viz_index) intrinsics = torch.index_select(self.cam0_intrinsics, 0, viz_index) # are these the up or down intrinsics? (down!) if self.args.multi_gpu: # We cannot send viz_out in our device, if the receiving end is in another device. # Need to cpu-transer, which is super slow. out_device = "cpu" else: out_device = self.device ic(out_device) viz_out = {"cam0_poses": cam0_T_world.to(device=out_device), "gt_poses": gt_poses.to(device=out_device), "gt_depths": gt_depths.to(device=out_device), "world_T_body": world_T_body.to(device=out_device), "world_T_body_cov": world_T_body_cov.to(device=out_device), "cam0_idepths": idepths.to(device=out_device), "cam0_idepths_up": idepths_up.to(device=out_device), "cam0_idepths_sensed": idepths_sensed.to(device=out_device), "cam0_idepths_cov": idepths_cov.to(device=out_device), "cam0_depths_cov": depths_cov.to(device=out_device), "cam0_depths_cov_up": depths_cov_up.to(device=out_device), "cam0_images": images.to(device=out_device), "cam0_intrinsics": intrinsics.to(device=out_device), # intrinsics are downsampled "calibs": batch["calibs"], # calibs are the original ones "viz_idx": viz_index.to(device=out_device), "kf_idx": self.kf_idx, "kf_idx_to_f_idx": self.kf_idx_to_f_idx, "is_last_frame": batch["is_last_frame"] } self.viz_idx[:] = False else: print("viz_index is empty, nothing to visualize") if batch["is_last_frame"]: ic("Last frame") # Leave all the entries null, this signals end of seq viz_out = {"is_last_frame": True} return viz_out