import torch import lietorch import numpy as np import matplotlib.pyplot as plt from lietorch import SE3 from .modules.corr import CorrBlock, AltCorrBlock from .geom import projective_ops as pops class FactorGraph: def __init__(self, video, update_net, device="cuda:0", corr_impl="volume", max_factors=-1): self.video = video self.update_net = update_net self.device = device self.max_factors = max_factors self.corr_impl = corr_impl # operator at 1/8 resolution self.ht = ht = video.ht // 8 self.wd = wd = video.wd // 8 self.coords0 = pops.coords_grid(ht, wd, device=device) self.ii = torch.as_tensor([], dtype=torch.long, device=device) self.jj = torch.as_tensor([], dtype=torch.long, device=device) self.age = torch.as_tensor([], dtype=torch.long, device=device) self.correlation_volumes, self.gru_hidden_states, self.gru_contexts_input = None, None, None self.damping = 1e-6 * torch.ones_like(self.video.disps) self.gru_estimated_flow = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float) self.gru_estimated_flow_weight = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float) # inactive factors self.ii_inac = torch.as_tensor([], dtype=torch.long, device=device) self.jj_inac = torch.as_tensor([], dtype=torch.long, device=device) self.ii_bad = torch.as_tensor([], dtype=torch.long, device=device) self.jj_bad = torch.as_tensor([], dtype=torch.long, device=device) self.target_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float) self.weight_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float) 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_inac, self.jj_inac)]) for k, (i, j) in enumerate(zip(ii, jj)): keep[k] = (i.item(), j.item()) not in eset return ii[keep], jj[keep] 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() def filter_edges(self): """ remove bad edges """ conf = torch.mean(self.gru_estimated_flow_weight, dim=[0,2,3,4]) mask = (torch.abs(self.ii-self.jj) > 2) & (conf < 0.001) self.ii_bad = torch.cat([self.ii_bad, self.ii[mask]]) self.jj_bad = torch.cat([self.jj_bad, self.jj[mask]]) self.rm_factors(mask, store=False) def clear_edges(self): self.rm_factors(self.ii >= 0) self.gru_hidden_states = None self.gru_contexts_input = None @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()] 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": c = (ii == jj).long() # TODO: does this still hold for multi-camera? fmap1 = self.video.fmaps[ii,0].to(self.device).unsqueeze(0) fmap2 = self.video.fmaps[jj,c].to(self.device).unsqueeze(0) corr = CorrBlock(fmap1, fmap2) self.correlation_volumes = corr if self.correlation_volumes is None \ else self.correlation_volumes.cat(corr) ### Gru hidden states are initialized to the context features, and then they evolve gru_hidden_state = self.video.nets[ii].to(self.device).unsqueeze(0) self.gru_hidden_states = gru_hidden_state if self.gru_hidden_states is None else \ torch.cat([self.gru_hidden_states, gru_hidden_state], 1) ### Gru estimated flow is initialized to the reprojected flow, and then it evolves with torch.cuda.amp.autocast(enabled=False): target, _ = self.video.reproject(ii, jj) weight = torch.zeros_like(target) # Initialize weights to 0 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? @torch.cuda.amp.autocast(enabled=True) def rm_factors(self, mask, store=False): """ drop edges from factor graph """ # store estimated factors if store: self.ii_inac = torch.cat([self.ii_inac, self.ii[mask]], 0) self.jj_inac = torch.cat([self.jj_inac, self.jj[mask]], 0) self.target_inac = torch.cat([self.target_inac, self.gru_estimated_flow[:,mask]], 1) self.weight_inac = torch.cat([self.weight_inac, self.gru_estimated_flow_weight[:,mask]], 1) 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] @torch.cuda.amp.autocast(enabled=True) def rm_keyframe(self, ix): """ drop nodes from factor graph """ with self.video.get_lock(): self.video.images[ix] = self.video.images[ix+1] self.video.poses[ix] = self.video.poses[ix+1] self.video.disps[ix] = self.video.disps[ix+1] self.video.disps_sens[ix] = self.video.disps_sens[ix+1] self.video.intrinsics[ix] = self.video.intrinsics[ix+1] self.video.nets[ix] = self.video.nets[ix+1] self.video.inps[ix] = self.video.inps[ix+1] self.video.fmaps[ix] = self.video.fmaps[ix+1] m = (self.ii_inac == ix) | (self.jj_inac == ix) self.ii_inac[self.ii_inac >= ix] -= 1 self.jj_inac[self.jj_inac >= ix] -= 1 if torch.any(m): self.ii_inac = self.ii_inac[~m] self.jj_inac = self.jj_inac[~m] self.target_inac = self.target_inac[:,~m] self.weight_inac = self.weight_inac[:,~m] m = (self.ii == ix) | (self.jj == ix) self.ii[self.ii >= ix] -= 1 self.jj[self.jj >= ix] -= 1 self.rm_factors(m, store=False) @torch.cuda.amp.autocast(enabled=True) def update(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, motion_only=False): """ run update operator on factor graph """ # motion features with torch.cuda.amp.autocast(enabled=False): # try mixed precision? coords1, mask = self.video.reproject(self.ii, self.jj) # "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 = coords0 + 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) self.gru_hidden_states, flow_delta, gru_estimated_flow_weight, damping, upmask = \ self.update_net(self.gru_hidden_states, self.gru_contexts_input, corr, motion, self.ii, self.jj) if t0 is None: t0 = max(1, self.ii.min().item()+1) with torch.cuda.amp.autocast(enabled=False): 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: # TODO What is this doing? m = (self.ii_inac >= t0 - 3) & (self.jj_inac >= t0 - 3) ii = torch.cat([self.ii_inac[m], self.ii], 0) jj = torch.cat([self.jj_inac[m], self.jj], 0) gru_estimated_flow = torch.cat([self.target_inac[:,m], self.gru_estimated_flow], 1) gru_estimated_flow_weight = torch.cat([self.weight_inac[:,m], 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? ht, wd = self.coords0.shape[0:2] gru_estimated_flow = gru_estimated_flow.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous() gru_estimated_flow_weight = gru_estimated_flow_weight.view(-1, ht, 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 self.video.ba(gru_estimated_flow, gru_estimated_flow_weight, damping, ii, jj, t0, t1, itrs=itrs, lm=1e-4, ep=0.1, motion_only=motion_only) self.age += 1 @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 t = self.video.counter.value num, rig, ch, ht, wd = self.video.fmaps.shape corr_op = AltCorrBlock(self.video.fmaps.view(1, num*rig, ch, ht, wd)) for step in range(steps): print("Global BA Iteration #{}".format(step+1)) with torch.cuda.amp.autocast(enabled=False): coords1, mask = self.video.reproject(self.ii, self.jj) motn = torch.cat([coords1 - self.coords0, self.gru_estimated_flow - coords1], dim=-1) motn = motn.permute(0,1,4,2,3).clamp(-64.0, 64.0) s = 8 for i in range(0, self.jj.max()+1, s): v = (self.ii >= i) & (self.ii < i + s) iis = self.ii[v] jjs = self.jj[v] ht, wd = self.coords0.shape[0:2] corr1 = corr_op(coords1[:,v], rig * iis, rig * jjs + (iis == jjs).long()) with torch.cuda.amp.autocast(enabled=True): net, delta, weight, damping, _ = \ self.update_net(self.gru_hidden_states[:,v], self.video.inps[None,iis], corr1, motn[:,v], iis, jjs) self.gru_hidden_states[:,v] = net self.gru_estimated_flow[:,v] = coords1[:,v] + delta.float() self.gru_estimated_flow_weight[:,v] = weight.float() self.damping[torch.unique(iis)] = damping damping = .2 * self.damping[torch.unique(self.ii)].contiguous() + EP target = self.gru_estimated_flow.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous() weight = self.gru_estimated_flow_weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous() # dense bundle adjustment self.video.ba(target, weight, damping, self.ii, self.jj, 1, t, itrs=itrs, lm=1e-5, ep=1e-2, motion_only=False) self.video.dirty[:t] = True def add_neighborhood_factors(self, t0, t1, r=3): """ add edges between neighboring frames within radius r """ # Build dense adjacency matrix ii, jj = torch.meshgrid(torch.arange(t0,t1), torch.arange(t0,t1)) 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.video.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. keep = ((ii - jj).abs() > c) & ((ii - jj).abs() <= r) # Computes the correlation between these frames self.add_factors(ii[keep], jj[keep]) def add_proximity_factors(self, t0=0, t1=0, rad=2, nms=2, beta=0.25, thresh=16.0, remove=False): """ add edges to the factor graph based on distance """ t = self.video.counter.value ix = torch.arange(t0, t) jx = torch.arange(t1, t) ii, jj = torch.meshgrid(ix, jx) ii = ii.reshape(-1) jj = jj.reshape(-1) d = self.video.distance(ii, jj, beta=beta) d[ii - rad < jj] = np.inf d[d > 100] = np.inf ii1 = torch.cat([self.ii, self.ii_bad, self.ii_inac], 0) jj1 = torch.cat([self.jj, self.jj_bad, self.jj_inac], 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 (t0 <= i1 < t) and (t1 <= j1 < t): d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf es = [] for i in range(t0, t): if self.video.stereo: es.append((i, i)) d[(i-t0)*(t-t1) + (i-t1)] = np.inf for j in range(max(i-rad-1,0), i): es.append((i,j)) es.append((j,i)) d[(i-t0)*(t-t1) + (j-t1)] = 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 (t0 <= i1 < t) and (t1 <= j1 < t): d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf ii, jj = torch.as_tensor(es, device=self.device).unbind(dim=-1) self.add_factors(ii, jj, remove)