import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm import contextlib import cv2 from dust3r.cloud_opt_flow.base_opt import BasePCOptimizer, edge_str from dust3r.cloud_opt_flow.pair_viewer import PairViewer from dust3r.utils.geometry import xy_grid, geotrf, depthmap_to_pts3d from dust3r.utils.device import to_cpu, to_numpy from dust3r.utils.goem_opt import DepthBasedWarping, OccMask, WarpImage, depth_regularization_si_weighted, tum_to_pose_matrix from third_party.raft import load_RAFT # from sam2.build_sam import build_sam2_video_predictor # sam2_checkpoint = "third_party/sam2/checkpoints/sam2.1_hiera_large.pt" # model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml" def smooth_L1_loss_fn(estimate, gt, mask, beta=1.0, per_pixel_thre=50.): loss_raw_shape = F.smooth_l1_loss(estimate*mask, gt*mask, beta=beta, reduction='none') if per_pixel_thre > 0: per_pixel_mask = (loss_raw_shape < per_pixel_thre) * mask else: per_pixel_mask = mask return torch.sum(loss_raw_shape * per_pixel_mask) / torch.sum(per_pixel_mask) def mse_loss_fn(estimate, gt, mask): v = torch.sum((estimate*mask-gt*mask)**2) / torch.sum(mask) return v # , v.item() class PointCloudOptimizer(BasePCOptimizer): """ Optimize a global scene, given a list of pairwise observations. Graph node: images Graph edges: observations = (pred1, pred2) """ def __init__(self, *args, optimize_pp=False, focal_break=20, shared_focal=False, flow_loss_fn='smooth_l1', flow_loss_weight=0.0, depth_regularize_weight=0.0, num_total_iter=300, temporal_smoothing_weight=0, translation_weight=0.1, flow_loss_start_epoch=0.15, flow_loss_thre=50, sintel_ckpt=False, use_self_mask=False, pxl_thre=50, sam2_mask_refine=True, motion_mask_thre=0.35, **kwargs): super().__init__(*args, **kwargs) self.has_im_poses = True # by definition of this class self.focal_break = focal_break self.num_total_iter = num_total_iter self.temporal_smoothing_weight = temporal_smoothing_weight self.translation_weight = translation_weight self.flow_loss_flag = False self.flow_loss_start_epoch = flow_loss_start_epoch self.flow_loss_thre = flow_loss_thre self.optimize_pp = optimize_pp self.pxl_thre = pxl_thre self.motion_mask_thre = motion_mask_thre # adding thing to optimize self.im_depthmaps = nn.ParameterList(torch.randn(H, W)/10-3 for H, W in self.imshapes) # log(depth) self.im_poses = nn.ParameterList(self.rand_pose(self.POSE_DIM) for _ in range(self.n_imgs)) # camera poses self.shared_focal = shared_focal if self.shared_focal: self.im_focals = nn.ParameterList(torch.FloatTensor( [self.focal_break*np.log(max(H, W))]) for H, W in self.imshapes[:1]) # camera intrinsics else: self.im_focals = nn.ParameterList(torch.FloatTensor( [self.focal_break*np.log(max(H, W))]) for H, W in self.imshapes) # camera intrinsics self.im_pp = nn.ParameterList(torch.zeros((2,)) for _ in range(self.n_imgs)) # camera intrinsics self.im_pp.requires_grad_(optimize_pp) self.imshape = self.imshapes[0] im_areas = [h*w for h, w in self.imshapes] self.max_area = max(im_areas) # adding thing to optimize self.im_depthmaps = ParameterStack(self.im_depthmaps, is_param=True, fill=self.max_area) #(num_imgs, H*W) self.im_poses = ParameterStack(self.im_poses, is_param=True) self.im_focals = ParameterStack(self.im_focals, is_param=True) self.im_pp = ParameterStack(self.im_pp, is_param=True) self.register_buffer('_pp', torch.tensor([(w/2, h/2) for h, w in self.imshapes])) self.register_buffer('_grid', ParameterStack( [xy_grid(W, H, device=self.device) for H, W in self.imshapes], fill=self.max_area)) # pre-compute pixel weights self.register_buffer('_weight_i', ParameterStack( [self.conf_trf(self.conf_i[i_j]) for i_j in self.str_edges], fill=self.max_area)) self.register_buffer('_weight_j', ParameterStack( [self.conf_trf(self.conf_j[i_j]) for i_j in self.str_edges], fill=self.max_area)) # precompute aa self.register_buffer('_stacked_pred_i', ParameterStack(self.pred_i, self.str_edges, fill=self.max_area)) self.register_buffer('_stacked_pred_j', ParameterStack(self.pred_j, self.str_edges, fill=self.max_area)) self.register_buffer('_ei', torch.tensor([i for i, j in self.edges])) self.register_buffer('_ej', torch.tensor([j for i, j in self.edges])) self.total_area_i = sum([im_areas[i] for i, j in self.edges]) self.total_area_j = sum([im_areas[j] for i, j in self.edges]) self.depth_wrapper = DepthBasedWarping() self.backward_warper = WarpImage() self.depth_regularizer = depth_regularization_si_weighted if flow_loss_fn == 'smooth_l1': self.flow_loss_fn = smooth_L1_loss_fn elif flow_loss_fn == 'mse': self.low_loss_fn = mse_loss_fn self.flow_loss_weight = flow_loss_weight self.depth_regularize_weight = depth_regularize_weight if self.flow_loss_weight > 0: self.flow_ij, self.flow_ji, self.flow_valid_mask_i, self.flow_valid_mask_j = self.get_flow(sintel_ckpt) # (num_pairs, 2, H, W) if use_self_mask: self.get_motion_mask_from_pairs(*args) # turn off the gradient for the flow self.flow_ij.requires_grad_(False) self.flow_ji.requires_grad_(False) self.flow_valid_mask_i.requires_grad_(False) self.flow_valid_mask_j.requires_grad_(False) sam2_mask_refine = False if sam2_mask_refine: with torch.no_grad(): self.refine_motion_mask_w_sam2() else: self.sam2_dynamic_masks = None def get_flow(self, sintel_ckpt=False): #TODO: test with gt flow print('precomputing flow...') device = 'cuda' if torch.cuda.is_available() else 'cpu' get_valid_flow_mask = OccMask(th=3.0) pair_imgs = [np.stack(self.imgs)[self._ei], np.stack(self.imgs)[self._ej]] sintel_ckpt=False flow_net = load_RAFT() if sintel_ckpt else load_RAFT("third_party/RAFT/models/Tartan-C-T432x960-M.pth") flow_net = flow_net.to(device) flow_net.eval() if len(pair_imgs[0].shape)==3: pair_imgs = [pair_imgs[0][None], pair_imgs[1][None]] #print(self._ei) with torch.no_grad(): chunk_size = 12 flow_ij = [] flow_ji = [] num_pairs = len(pair_imgs[0]) for i in tqdm(range(0, num_pairs, chunk_size)): end_idx = min(i + chunk_size, num_pairs) imgs_ij = [torch.tensor(pair_imgs[0][i:end_idx]).float().to(device), torch.tensor(pair_imgs[1][i:end_idx]).float().to(device)] #print(imgs_ij[0].shape) flow_ij.append(flow_net(imgs_ij[0].permute(0, 3, 1, 2) * 255, imgs_ij[1].permute(0, 3, 1, 2) * 255, iters=20, test_mode=True)[1]) flow_ji.append(flow_net(imgs_ij[1].permute(0, 3, 1, 2) * 255, imgs_ij[0].permute(0, 3, 1, 2) * 255, iters=20, test_mode=True)[1]) flow_ij = torch.cat(flow_ij, dim=0) flow_ji = torch.cat(flow_ji, dim=0) valid_mask_i = get_valid_flow_mask(flow_ij, flow_ji) valid_mask_j = get_valid_flow_mask(flow_ji, flow_ij) print('flow precomputed') # delete the flow net if flow_net is not None: del flow_net return flow_ij, flow_ji, valid_mask_i, valid_mask_j def get_motion_mask_from_pairs(self, view1, view2, pred1, pred2): assert self.is_symmetrized, 'only support symmetric case' symmetry_pairs_idx = [(i, i+len(self.edges)//2) for i in range(len(self.edges)//2)] intrinsics_i = [] intrinsics_j = [] R_i = [] R_j = [] T_i = [] T_j = [] depth_maps_i = [] depth_maps_j = [] for i, j in tqdm(symmetry_pairs_idx): new_view1 = {} new_view2 = {} for key in view1.keys(): if isinstance(view1[key], list): new_view1[key] = [view1[key][i], view1[key][j]] new_view2[key] = [view2[key][i], view2[key][j]] elif isinstance(view1[key], torch.Tensor): new_view1[key] = torch.stack([view1[key][i], view1[key][j]]) new_view2[key] = torch.stack([view2[key][i], view2[key][j]]) new_view1['idx'] = [0, 1] new_view2['idx'] = [1, 0] new_pred1 = {} new_pred2 = {} for key in pred1.keys(): if isinstance(pred1[key], list): new_pred1[key] = [pred1[key][i], pred1[key][j]] elif isinstance(pred1[key], torch.Tensor): new_pred1[key] = torch.stack([pred1[key][i], pred1[key][j]]) for key in pred2.keys(): if isinstance(pred2[key], list): new_pred2[key] = [pred2[key][i], pred2[key][j]] elif isinstance(pred2[key], torch.Tensor): new_pred2[key] = torch.stack([pred2[key][i], pred2[key][j]]) pair_viewer = PairViewer(new_view1, new_view2, new_pred1, new_pred2, verbose=False) intrinsics_i.append(pair_viewer.get_intrinsics()[0]) intrinsics_j.append(pair_viewer.get_intrinsics()[1]) R_i.append(pair_viewer.get_im_poses()[0][:3, :3]) R_j.append(pair_viewer.get_im_poses()[1][:3, :3]) T_i.append(pair_viewer.get_im_poses()[0][:3, 3:]) T_j.append(pair_viewer.get_im_poses()[1][:3, 3:]) depth_maps_i.append(pair_viewer.get_depthmaps()[0]) depth_maps_j.append(pair_viewer.get_depthmaps()[1]) self.intrinsics_i = torch.stack(intrinsics_i).to(self.flow_ij.device) self.intrinsics_j = torch.stack(intrinsics_j).to(self.flow_ij.device) self.R_i = torch.stack(R_i).to(self.flow_ij.device) self.R_j = torch.stack(R_j).to(self.flow_ij.device) self.T_i = torch.stack(T_i).to(self.flow_ij.device) self.T_j = torch.stack(T_j).to(self.flow_ij.device) self.depth_maps_i = torch.stack(depth_maps_i).unsqueeze(1).to(self.flow_ij.device) self.depth_maps_j = torch.stack(depth_maps_j).unsqueeze(1).to(self.flow_ij.device) # self.depth_maps_i[self.depth_maps_i>0.7] = 0.7 # self.depth_maps_j[self.depth_maps_j>0.7] = 0.7 #cv2.imwrite('1.png', self.depth_maps_i[0,0].cpu().numpy()*255) #print(self.depth_maps_i,self.depth_maps_i.shape) try: ego_flow_1_2, _ = self.depth_wrapper(self.R_i, self.T_i, self.R_j, self.T_j, 1 / (self.depth_maps_i + 1e-6), self.intrinsics_j, torch.linalg.inv(self.intrinsics_i)) except Exception as e: ego_flow_1_2, _ = self.depth_wrapper(self.R_i, self.T_i, self.R_j, self.T_j, 1 / (self.depth_maps_i + 1e-6), self.intrinsics_j, torch.linalg.pinv(self.intrinsics_i)) try: ego_flow_2_1, _ = self.depth_wrapper(self.R_j, self.T_j, self.R_i, self.T_i, 1 / (self.depth_maps_j + 1e-6), self.intrinsics_i, torch.linalg.inv(self.intrinsics_j)) except Exception as e: ego_flow_2_1, _ = self.depth_wrapper(self.R_j, self.T_j, self.R_i, self.T_i, 1 / (self.depth_maps_j + 1e-6), self.intrinsics_i, torch.linalg.pinv(self.intrinsics_j)) err_map_i = torch.norm(ego_flow_1_2[:, :2, ...] - self.flow_ij[:len(symmetry_pairs_idx)], dim=1) err_map_j = torch.norm(ego_flow_2_1[:, :2, ...] - self.flow_ji[:len(symmetry_pairs_idx)], dim=1) # normalize the error map for each pair err_map_i = (err_map_i - err_map_i.amin(dim=(1, 2), keepdim=True)) / (err_map_i.amax(dim=(1, 2), keepdim=True) - err_map_i.amin(dim=(1, 2), keepdim=True)) err_map_j = (err_map_j - err_map_j.amin(dim=(1, 2), keepdim=True)) / (err_map_j.amax(dim=(1, 2), keepdim=True) - err_map_j.amin(dim=(1, 2), keepdim=True)) self.dynamic_masks = [[] for _ in range(self.n_imgs)] for i, j in symmetry_pairs_idx: i_idx = self._ei[i] j_idx = self._ej[i] self.dynamic_masks[i_idx].append(err_map_i[i]) self.dynamic_masks[j_idx].append(err_map_j[i]) for i in range(self.n_imgs): self.dynamic_masks[i] = torch.stack(self.dynamic_masks[i]).mean(dim=0) > self.motion_mask_thre def refine_motion_mask_w_sam2(self): device = 'cuda' if torch.cuda.is_available() else 'cpu' # Save previous TF32 settings if device == 'cuda': prev_allow_tf32 = torch.backends.cuda.matmul.allow_tf32 prev_allow_cudnn_tf32 = torch.backends.cudnn.allow_tf32 # Enable TF32 for Ampere GPUs if torch.cuda.get_device_properties(0).major >= 8: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True try: autocast_dtype = torch.bfloat16 if device == 'cuda' else torch.float32 with torch.autocast(device_type=device, dtype=autocast_dtype): predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device=device) frame_tensors = torch.from_numpy(np.array((self.imgs))).permute(0, 3, 1, 2).to(device) inference_state = predictor.init_state(video_path=frame_tensors) mask_list = [self.dynamic_masks[i] for i in range(self.n_imgs)] ann_obj_id = 1 self.sam2_dynamic_masks = [[] for _ in range(self.n_imgs)] # Process even frames predictor.reset_state(inference_state) for idx, mask in enumerate(mask_list): if idx % 2 == 1: _, out_obj_ids, out_mask_logits = predictor.add_new_mask( inference_state, frame_idx=idx, obj_id=ann_obj_id, mask=mask, ) video_segments = {} for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state, start_frame_idx=0): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } for out_frame_idx in range(self.n_imgs): if out_frame_idx % 2 == 0: self.sam2_dynamic_masks[out_frame_idx] = video_segments[out_frame_idx][ann_obj_id] # Process odd frames predictor.reset_state(inference_state) for idx, mask in enumerate(mask_list): if idx % 2 == 0: _, out_obj_ids, out_mask_logits = predictor.add_new_mask( inference_state, frame_idx=idx, obj_id=ann_obj_id, mask=mask, ) video_segments = {} for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state, start_frame_idx=0): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } for out_frame_idx in range(self.n_imgs): if out_frame_idx % 2 == 1: self.sam2_dynamic_masks[out_frame_idx] = video_segments[out_frame_idx][ann_obj_id] # Update dynamic masks for i in range(self.n_imgs): self.sam2_dynamic_masks[i] = torch.from_numpy(self.sam2_dynamic_masks[i][0]).to(device) self.dynamic_masks[i] = self.dynamic_masks[i].to(device) self.dynamic_masks[i] = self.dynamic_masks[i] | self.sam2_dynamic_masks[i] # Clean up del predictor finally: # Restore previous TF32 settings if device == 'cuda': torch.backends.cuda.matmul.allow_tf32 = prev_allow_tf32 torch.backends.cudnn.allow_tf32 = prev_allow_cudnn_tf32 def _check_all_imgs_are_selected(self, msk): self.msk = torch.from_numpy(np.array(msk, dtype=bool)).to(self.device) assert np.all(self._get_msk_indices(msk) == np.arange(self.n_imgs)), 'incomplete mask!' pass def preset_pose(self, known_poses, pose_msk=None, requires_grad=False): # cam-to-world self._check_all_imgs_are_selected(pose_msk) if isinstance(known_poses, torch.Tensor) and known_poses.ndim == 2: known_poses = [known_poses] if known_poses.shape[-1] == 7: # xyz wxyz known_poses = [tum_to_pose_matrix(pose) for pose in known_poses] for idx, pose in zip(self._get_msk_indices(pose_msk), known_poses): if self.verbose: print(f' (setting pose #{idx} = {pose[:3,3]})') self._no_grad(self._set_pose(self.im_poses, idx, torch.tensor(pose))) # normalize scale if there's less than 1 known pose n_known_poses = sum((p.requires_grad is False) for p in self.im_poses) self.norm_pw_scale = (n_known_poses <= 1) if len(known_poses) == self.n_imgs: if requires_grad: self.im_poses.requires_grad_(True) else: self.im_poses.requires_grad_(False) self.norm_pw_scale = False def preset_intrinsics(self, known_intrinsics, msk=None): if isinstance(known_intrinsics, torch.Tensor) and known_intrinsics.ndim == 2: known_intrinsics = [known_intrinsics] for K in known_intrinsics: assert K.shape == (3, 3) self.preset_focal([K.diagonal()[:2].mean() for K in known_intrinsics], msk) if self.optimize_pp: self.preset_principal_point([K[:2, 2] for K in known_intrinsics], msk) def preset_focal(self, known_focals, msk=None, requires_grad=False): self._check_all_imgs_are_selected(msk) for idx, focal in zip(self._get_msk_indices(msk), known_focals): if self.verbose: print(f' (setting focal #{idx} = {focal})') self._no_grad(self._set_focal(idx, focal)) if len(known_focals) == self.n_imgs: if requires_grad: self.im_focals.requires_grad_(True) else: self.im_focals.requires_grad_(False) def preset_principal_point(self, known_pp, msk=None): self._check_all_imgs_are_selected(msk) for idx, pp in zip(self._get_msk_indices(msk), known_pp): if self.verbose: print(f' (setting principal point #{idx} = {pp})') self._no_grad(self._set_principal_point(idx, pp)) self.im_pp.requires_grad_(False) def _get_msk_indices(self, msk): if msk is None: return range(self.n_imgs) elif isinstance(msk, int): return [msk] elif isinstance(msk, (tuple, list)): return self._get_msk_indices(np.array(msk)) elif msk.dtype in (bool, torch.bool, np.bool_): assert len(msk) == self.n_imgs return np.where(msk)[0] elif np.issubdtype(msk.dtype, np.integer): return msk else: raise ValueError(f'bad {msk=}') def _no_grad(self, tensor): assert tensor.requires_grad, 'it must be True at this point, otherwise no modification occurs' def _set_focal(self, idx, focal, force=False): param = self.im_focals[idx] if param.requires_grad or force: # can only init a parameter not already initialized param.data[:] = self.focal_break * np.log(focal) return param def get_focals(self): if self.shared_focal: log_focals = torch.stack([self.im_focals[0]] * self.n_imgs, dim=0) else: log_focals = torch.stack(list(self.im_focals), dim=0) return (log_focals / self.focal_break).exp() def get_known_focal_mask(self): return torch.tensor([not (p.requires_grad) for p in self.im_focals]) def _set_principal_point(self, idx, pp, force=False): param = self.im_pp[idx] H, W = self.imshapes[idx] if param.requires_grad or force: # can only init a parameter not already initialized param.data[:] = to_cpu(to_numpy(pp) - (W/2, H/2)) / 10 return param def get_principal_points(self): return self._pp + 10 * self.im_pp def get_intrinsics(self): K = torch.zeros((self.n_imgs, 3, 3), device=self.device) focals = self.get_focals().flatten() K[:, 0, 0] = K[:, 1, 1] = focals K[:, :2, 2] = self.get_principal_points() K[:, 2, 2] = 1 return K def get_im_poses(self): # cam to world cam2world = self._get_poses(self.im_poses) return cam2world def _set_depthmap(self, idx, depth, force=False): depth = _ravel_hw(depth, self.max_area) param = self.im_depthmaps[idx] if param.requires_grad or force: # can only init a parameter not already initialized param.data[:] = depth.log().nan_to_num(neginf=0) return param def preset_depthmap(self, known_depthmaps, msk=None, requires_grad=False): self._check_all_imgs_are_selected(msk) for idx, depth in zip(self._get_msk_indices(msk), known_depthmaps): if self.verbose: print(f' (setting depthmap #{idx})') self._no_grad(self._set_depthmap(idx, depth)) if len(known_depthmaps) == self.n_imgs: if requires_grad: self.im_depthmaps.requires_grad_(True) else: self.im_depthmaps.requires_grad_(False) def _set_init_depthmap(self): depth_maps = self.get_depthmaps(raw=True) self.init_depthmap = [dm.detach().clone() for dm in depth_maps] def get_init_depthmaps(self, raw=False): res = self.init_depthmap if not raw: res = [dm[:h*w].view(h, w) for dm, (h, w) in zip(res, self.imshapes)] return res def get_depthmaps(self, raw=False): res = self.im_depthmaps.exp() if not raw: res = [dm[:h*w].view(h, w) for dm, (h, w) in zip(res, self.imshapes)] return res def depth_to_pts3d(self): # Get depths and projection params if not provided focals = self.get_focals() pp = self.get_principal_points() im_poses = self.get_im_poses() depth = self.get_depthmaps(raw=True) # get pointmaps in camera frame rel_ptmaps = _fast_depthmap_to_pts3d(depth, self._grid, focals, pp=pp) # project to world frame return geotrf(im_poses, rel_ptmaps) def depth_to_pts3d_partial(self): # Get depths and projection params if not provided focals = self.get_focals() pp = self.get_principal_points() im_poses = self.get_im_poses() depth = self.get_depthmaps() # convert focal to (1,2,H,W) constant field def focal_ex(i): return focals[i][..., None, None].expand(1, *focals[i].shape, *self.imshapes[i]) # get pointmaps in camera frame rel_ptmaps = [depthmap_to_pts3d(depth[i][None], focal_ex(i), pp=pp[i:i+1])[0] for i in range(im_poses.shape[0])] # project to world frame return [geotrf(pose, ptmap) for pose, ptmap in zip(im_poses, rel_ptmaps)] def get_pts3d(self, raw=False, **kwargs): res = self.depth_to_pts3d() if not raw: res = [dm[:h*w].view(h, w, 3) for dm, (h, w) in zip(res, self.imshapes)] return res def forward(self, epoch=9999): pw_poses = self.get_pw_poses() # cam-to-world pw_adapt = self.get_adaptors().unsqueeze(1) proj_pts3d = self.get_pts3d(raw=True) # rotate pairwise prediction according to pw_poses aligned_pred_i = geotrf(pw_poses, pw_adapt * self._stacked_pred_i) aligned_pred_j = geotrf(pw_poses, pw_adapt * self._stacked_pred_j) # compute the less li = self.dist(proj_pts3d[self._ei], aligned_pred_i, weight=self._weight_i).sum() / self.total_area_i lj = self.dist(proj_pts3d[self._ej], aligned_pred_j, weight=self._weight_j).sum() / self.total_area_j # camera temporal loss if self.temporal_smoothing_weight > 0: temporal_smoothing_loss = self.relative_pose_loss(self.get_im_poses()[:-1], self.get_im_poses()[1:]).sum() else: temporal_smoothing_loss = 0 if self.flow_loss_weight > 0 and epoch >= self.num_total_iter * self.flow_loss_start_epoch: # enable flow loss after certain epoch R_all, T_all = self.get_im_poses()[:,:3].split([3, 1], dim=-1) R1, T1 = R_all[self._ei], T_all[self._ei] R2, T2 = R_all[self._ej], T_all[self._ej] K_all = self.get_intrinsics() inv_K_all = torch.linalg.inv(K_all) K_1, inv_K_1 = K_all[self._ei], inv_K_all[self._ei] K_2, inv_K_2 = K_all[self._ej], inv_K_all[self._ej] depth_all = torch.stack(self.get_depthmaps(raw=False)).unsqueeze(1) depth1, depth2 = depth_all[self._ei], depth_all[self._ej] disp_1, disp_2 = 1 / (depth1 + 1e-6), 1 / (depth2 + 1e-6) ego_flow_1_2, _ = self.depth_wrapper(R1, T1, R2, T2, disp_1, K_2, inv_K_1) ego_flow_2_1, _ = self.depth_wrapper(R2, T2, R1, T1, disp_2, K_1, inv_K_2) dynamic_masks_all = torch.stack(self.dynamic_masks).to(self.device).unsqueeze(1) dynamic_mask1, dynamic_mask2 = dynamic_masks_all[self._ei], dynamic_masks_all[self._ej] flow_loss_i = self.flow_loss_fn(ego_flow_1_2[:, :2, ...], self.flow_ij, ~dynamic_mask1, per_pixel_thre=self.pxl_thre) flow_loss_j = self.flow_loss_fn(ego_flow_2_1[:, :2, ...], self.flow_ji, ~dynamic_mask2, per_pixel_thre=self.pxl_thre) flow_loss = flow_loss_i + flow_loss_j print(f'flow loss: {flow_loss.item()}') if flow_loss.item() > self.flow_loss_thre and self.flow_loss_thre > 0: flow_loss = 0 self.flow_loss_flag = True else: flow_loss = 0 if self.depth_regularize_weight > 0: init_depthmaps = torch.stack(self.get_init_depthmaps(raw=False)).unsqueeze(1) depthmaps = torch.stack(self.get_depthmaps(raw=False)).unsqueeze(1) dynamic_masks_all = torch.stack(self.dynamic_masks).to(self.device).unsqueeze(1) depth_prior_loss = self.depth_regularizer(depthmaps, init_depthmaps, dynamic_masks_all) else: depth_prior_loss = 0 loss = (li + lj) * 1 + self.temporal_smoothing_weight * temporal_smoothing_loss + \ self.flow_loss_weight * flow_loss + self.depth_regularize_weight * depth_prior_loss return loss def relative_pose_loss(self, RT1, RT2): relative_RT = torch.matmul(torch.inverse(RT1), RT2) rotation_diff = relative_RT[:, :3, :3] translation_diff = relative_RT[:, :3, 3] # Frobenius norm for rotation difference rotation_loss = torch.norm(rotation_diff - (torch.eye(3, device=RT1.device)), dim=(1, 2)) # L2 norm for translation difference translation_loss = torch.norm(translation_diff, dim=1) # Combined loss (one can weigh these differently if needed) pose_loss = rotation_loss + translation_loss * self.translation_weight return pose_loss def _fast_depthmap_to_pts3d(depth, pixel_grid, focal, pp): pp = pp.unsqueeze(1) focal = focal.unsqueeze(1) assert focal.shape == (len(depth), 1, 1) assert pp.shape == (len(depth), 1, 2) assert pixel_grid.shape == depth.shape + (2,) depth = depth.unsqueeze(-1) return torch.cat((depth * (pixel_grid - pp) / focal, depth), dim=-1) def ParameterStack(params, keys=None, is_param=None, fill=0): if keys is not None: params = [params[k] for k in keys] if fill > 0: params = [_ravel_hw(p, fill) for p in params] requires_grad = params[0].requires_grad assert all(p.requires_grad == requires_grad for p in params) params = torch.stack(list(params)).float().detach() if is_param or requires_grad: params = nn.Parameter(params) params.requires_grad_(requires_grad) return params def _ravel_hw(tensor, fill=0): # ravel H,W tensor = tensor.view((tensor.shape[0] * tensor.shape[1],) + tensor.shape[2:]) if len(tensor) < fill: tensor = torch.cat((tensor, tensor.new_zeros((fill - len(tensor),)+tensor.shape[1:]))) return tensor def acceptable_focal_range(H, W, minf=0.5, maxf=3.5): focal_base = max(H, W) / (2 * np.tan(np.deg2rad(60) / 2)) # size / 1.1547005383792515 return minf*focal_base, maxf*focal_base def apply_mask(img, msk): img = img.copy() img[msk] = 0 return img def ordered_ratio(disp_a, disp_b, mask=None): ratio_a = torch.maximum(disp_a, disp_b) / \ (torch.minimum(disp_a, disp_b)+1e-5) if mask is not None: ratio_a = ratio_a[mask] return ratio_a - 1