import os import open3d as o3d import torch import torch.nn.functional as F import numpy as np import matplotlib.pyplot as plt from tqdm import tqdm from utils.recon_helpers import setup_camera from utils.slam_external import build_rotation,calc_psnr from diff_gaussian_rasterization import GaussianRasterizer as Renderer from pytorch_msssim import ms_ssim from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity loss_fn_alex = LearnedPerceptualImagePatchSimilarity(net_type='alex', normalize=True).cuda() def l1_loss_v1(x, y): return torch.abs((x - y)).mean() def l1_loss_v2(x, y): return (torch.abs(x - y).sum(-1)).mean() def weighted_l2_loss_v1(x, y, w): return torch.sqrt(((x - y) ** 2) * w + 1e-20).mean() def weighted_l2_loss_v2(x, y, w): return torch.sqrt(((x - y) ** 2).sum(-1) * w + 1e-20).mean() def align(model, data): """Align two trajectories using the method of Horn (closed-form). Args: model -- first trajectory (3xn) data -- second trajectory (3xn) Returns: rot -- rotation matrix (3x3) trans -- translation vector (3x1) trans_error -- translational error per point (1xn) """ np.set_printoptions(precision=3, suppress=True) model_zerocentered = model - model.mean(1).reshape((3,-1)) data_zerocentered = data - data.mean(1).reshape((3,-1)) W = np.zeros((3, 3)) for column in range(model.shape[1]): W += np.outer(model_zerocentered[:, column], data_zerocentered[:, column]) U, d, Vh = np.linalg.linalg.svd(W.transpose()) S = np.matrix(np.identity(3)) if (np.linalg.det(U) * np.linalg.det(Vh) < 0): S[2, 2] = -1 rot = U*S*Vh trans = data.mean(1).reshape((3,-1)) - rot * model.mean(1).reshape((3,-1)) model_aligned = rot * model + trans alignment_error = model_aligned - data trans_error = np.sqrt(np.sum(np.multiply( alignment_error, alignment_error), 0)).A[0] return rot, trans, trans_error def evaluate_ate(gt_traj, est_traj): """ Input : gt_traj: list of 4x4 matrices est_traj: list of 4x4 matrices len(gt_traj) == len(est_traj) """ gt_traj_pts = [gt_traj[idx][:3,3] for idx in range(len(gt_traj))] est_traj_pts = [est_traj[idx][:3,3] for idx in range(len(est_traj))] gt_traj_pts = torch.stack(gt_traj_pts).detach().cpu().numpy().T est_traj_pts = torch.stack(est_traj_pts).detach().cpu().numpy().T _, _, trans_error = align(gt_traj_pts, est_traj_pts) avg_trans_error = trans_error.mean() return avg_trans_error def quat_mult(q1, q2): w1, x1, y1, z1 = q1.T w2, x2, y2, z2 = q2.T w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 return torch.stack([w, x, y, z]).T def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: """ Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0. Source: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#matrix_to_quaternion """ ret = torch.zeros_like(x) positive_mask = x > 0 ret[positive_mask] = torch.sqrt(x[positive_mask]) return ret def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: """ Convert rotations given as rotation matrices to quaternions. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). Returns: quaternions with real part first, as tensor of shape (..., 4). Source: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#matrix_to_quaternion """ if matrix.size(-1) != 3 or matrix.size(-2) != 3: raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") batch_dim = matrix.shape[:-2] m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( matrix.reshape(batch_dim + (9,)), dim=-1 ) q_abs = _sqrt_positive_part( torch.stack( [ 1.0 + m00 + m11 + m22, 1.0 + m00 - m11 - m22, 1.0 - m00 + m11 - m22, 1.0 - m00 - m11 + m22, ], dim=-1, ) ) # we produce the desired quaternion multiplied by each of r, i, j, k quat_by_rijk = torch.stack( [ # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and # `int`. torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and # `int`. torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and # `int`. torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and # `int`. torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), ], dim=-2, ) # We floor here at 0.1 but the exact level is not important; if q_abs is small, # the candidate won't be picked. flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) # if not for numerical problems, quat_candidates[i] should be same (up to a sign), # forall i; we pick the best-conditioned one (with the largest denominator) return quat_candidates[ F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : ].reshape(batch_dim + (4,)) def o3d_knn(pts, num_knn): indices = [] sq_dists = [] pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(np.ascontiguousarray(pts, np.float64)) pcd_tree = o3d.geometry.KDTreeFlann(pcd) for p in pcd.points: [_, i, d] = pcd_tree.search_knn_vector_3d(p, num_knn + 1) indices.append(i[1:]) sq_dists.append(d[1:]) return np.array(sq_dists), np.array(indices) def params2rendervar(params): rendervar = { 'means3D': params['means3D'], 'colors_precomp': params['rgb_colors'], 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def transformed_params2rendervar(params, transformed_pts): rendervar = { 'means3D': transformed_pts, 'colors_precomp': params['rgb_colors'], 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def project_points(points_3d, intrinsics): """ Function to project 3D points to image plane. params: points_3d: [num_gaussians, 3] intrinsics: [3, 3] out: [num_gaussians, 2] """ points_2d = torch.matmul(intrinsics, points_3d.transpose(0, 1)) points_2d = points_2d.transpose(0, 1) points_2d = points_2d / points_2d[:, 2:] points_2d = points_2d[:, :2] return points_2d def params2silhouette(params): sil_color = torch.zeros_like(params['rgb_colors']) sil_color[:, 0] = 1.0 rendervar = { 'means3D': params['means3D'], 'colors_precomp': sil_color, 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def transformed_params2silhouette(params, transformed_pts): sil_color = torch.zeros_like(params['rgb_colors']) sil_color[:, 0] = 1.0 rendervar = { 'means3D': transformed_pts, 'colors_precomp': sil_color, 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def get_depth_and_silhouette(pts_3D, w2c): """ Function to compute depth and silhouette for each gaussian. These are evaluated at gaussian center. """ # Depth of each gaussian center in camera frame pts4 = torch.cat((pts_3D, torch.ones_like(pts_3D[:, :1])), dim=-1) pts_in_cam = (w2c @ pts4.transpose(0, 1)).transpose(0, 1) depth_z = pts_in_cam[:, 2].unsqueeze(-1) # [num_gaussians, 1] depth_z_sq = torch.square(depth_z) # [num_gaussians, 1] # Depth and Silhouette depth_silhouette = torch.zeros((pts_3D.shape[0], 3)).cuda().float() depth_silhouette[:, 0] = depth_z.squeeze(-1) depth_silhouette[:, 1] = 1.0 depth_silhouette[:, 2] = depth_z_sq.squeeze(-1) return depth_silhouette def params2depthplussilhouette(params, w2c): rendervar = { 'means3D': params['means3D'], 'colors_precomp': get_depth_and_silhouette(params['means3D'], w2c), 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def transformed_params2depthplussilhouette(params, w2c, transformed_pts): rendervar = { 'means3D': transformed_pts, 'colors_precomp': get_depth_and_silhouette(transformed_pts, w2c), 'rotations': F.normalize(params['unnorm_rotations']), 'opacities': torch.sigmoid(params['logit_opacities']), 'scales': torch.exp(torch.tile(params['log_scales'], (1, 3))), 'means2D': torch.zeros_like(params['means3D'], requires_grad=True, device="cuda") + 0 } return rendervar def transform_to_frame(params, time_idx, gaussians_grad, camera_grad): """ Function to transform Isotropic Gaussians from world frame to camera frame. Args: params: dict of parameters time_idx: time index to transform to gaussians_grad: enable gradients for Gaussians camera_grad: enable gradients for camera pose Returns: transformed_pts: Transformed Centers of Gaussians """ # Get Frame Camera Pose if camera_grad: cam_rot = F.normalize(params['cam_unnorm_rots'][..., time_idx]) cam_tran = params['cam_trans'][..., time_idx] else: cam_rot = F.normalize(params['cam_unnorm_rots'][..., time_idx].detach()) cam_tran = params['cam_trans'][..., time_idx].detach() rel_w2c = torch.eye(4).cuda().float() rel_w2c[:3, :3] = build_rotation(cam_rot) rel_w2c[:3, 3] = cam_tran # Get Centers and norm Rots of Gaussians in World Frame if gaussians_grad: pts = params['means3D'] else: pts = params['means3D'].detach() # Transform Centers and Unnorm Rots of Gaussians to Camera Frame pts_ones = torch.ones(pts.shape[0], 1).cuda().float() pts4 = torch.cat((pts, pts_ones), dim=1) transformed_pts = (rel_w2c @ pts4.T).T[:, :3] return transformed_pts def report_loss(losses, wandb_run, wandb_step, tracking=False, mapping=False): # Update loss dict loss_dict = {'Loss': losses['loss'].item(), 'Image Loss': losses['im'].item(), 'Depth Loss': losses['depth'].item(),} if tracking: tracking_loss_dict = {} for k, v in loss_dict.items(): tracking_loss_dict[f"Tracking {k}"] = v wandb_run.log(tracking_loss_dict, step=wandb_step) elif mapping: mapping_loss_dict = {} for k, v in loss_dict.items(): mapping_loss_dict[f"Mapping {k}"] = v wandb_run.log(mapping_loss_dict, step=wandb_step) else: frame_opt_loss_dict = {} for k, v in loss_dict.items(): frame_opt_loss_dict[f"Current Frame Optimization {k}"] = v wandb_run.log(frame_opt_loss_dict, step=wandb_step) # Increment wandb step wandb_step += 1 return wandb_step def plot_rgbd_silhouette(color, depth, rastered_color, rastered_depth, presence_sil_mask, diff_depth_rmse, psnr, rmse, fig_title, plot_dir=None, plot_name=None, save_plot=False, wandb_run=None, wandb_step=None, wandb_title=None): # Determine Plot Aspect Ratio aspect_ratio = color.shape[2] / color.shape[1] fig_height = 8 fig_width = 14/1.55 fig_width = fig_width * aspect_ratio # Plot the Ground Truth and Rasterized RGB & Depth, along with Diff Depth & Silhouette fig, axs = plt.subplots(2, 3, figsize=(fig_width, fig_height)) axs[0, 0].imshow(color.cpu().permute(1, 2, 0)) axs[0, 0].set_title("Ground Truth RGB") axs[0, 1].imshow(depth[0, :, :].cpu(), cmap='jet', vmin=0, vmax=6) axs[0, 1].set_title("Ground Truth Depth") rastered_color = torch.clamp(rastered_color, 0, 1) axs[1, 0].imshow(rastered_color.cpu().permute(1, 2, 0)) axs[1, 0].set_title("Rasterized RGB, PSNR: {:.2f}".format(psnr)) axs[1, 1].imshow(rastered_depth[0, :, :].cpu(), cmap='jet', vmin=0, vmax=6) axs[1, 1].set_title("Rasterized Depth, RMSE: {:.2f}".format(rmse)) axs[0, 2].imshow(presence_sil_mask, cmap='gray') axs[0, 2].set_title("Rasterized Silhouette") diff_depth_rmse = diff_depth_rmse.cpu().squeeze(0) axs[1, 2].imshow(diff_depth_rmse, cmap='jet', vmin=0, vmax=1) axs[1, 2].set_title("Diff Depth RMSE") for ax in axs.flatten(): ax.axis('off') fig.suptitle(fig_title, y=0.95, fontsize=16) fig.tight_layout() if save_plot: save_path = os.path.join(plot_dir, f"{plot_name}.png") plt.savefig(save_path, bbox_inches='tight') if wandb_run is not None: if wandb_step is None: wandb_run.log({wandb_title: fig}) else: wandb_run.log({wandb_title: fig}, step=wandb_step) plt.close() def report_progress(params, data, i, progress_bar, iter_time_idx, sil_thres, every_i=1, qual_every_i=1, tracking=False, mapping=False, wandb_run=None, wandb_step=None, wandb_save_qual=False, online_time_idx=None): if i % every_i == 0 or i == 1: if wandb_run is not None: if tracking: stage = "Tracking" elif mapping: stage = "Mapping" else: stage = "Current Frame Optimization" # Initialize Render Variables rendervar = params2rendervar(params) depth_sil_rendervar = params2depthplussilhouette(params, data['w2c']) # Initialize Render Variables depth_sil, _, _, = Renderer(raster_settings=data['cam'])(**depth_sil_rendervar) rastered_depth = depth_sil[0, :, :].unsqueeze(0) valid_depth_mask = (data['depth'] > 0) silhouette = depth_sil[1, :, :] presence_sil_mask = (silhouette > sil_thres) im, _, _, = Renderer(raster_settings=data['cam'])(**rendervar) if tracking: psnr = calc_psnr(im * presence_sil_mask, data['im'] * presence_sil_mask).mean() else: psnr = calc_psnr(im, data['im']).mean() if tracking: diff_depth_rmse = torch.sqrt((((rastered_depth - data['depth']) * presence_sil_mask) ** 2)) diff_depth_rmse = diff_depth_rmse * valid_depth_mask rmse = diff_depth_rmse.sum() / valid_depth_mask.sum() else: diff_depth_rmse = torch.sqrt(((rastered_depth - data['depth']) ** 2)) diff_depth_rmse = diff_depth_rmse * valid_depth_mask rmse = diff_depth_rmse.sum() / valid_depth_mask.sum() if not mapping: progress_bar.set_postfix({f"Time-Step: {iter_time_idx} | Frame {data['id']} | PSNR: {psnr:.{7}} | RMSE": f"{rmse:.{7}}"}) progress_bar.update(every_i) else: progress_bar.set_postfix({f"Time-Step: {online_time_idx} | Frame {data['id']} | PSNR: {psnr:.{7}} | RMSE": f"{rmse:.{7}}"}) progress_bar.update(every_i) if wandb_run is not None: wandb_run.log({f"{stage} PSNR": psnr, f"{stage} RMSE": rmse}, step=wandb_step) if wandb_save_qual and (i % qual_every_i == 0 or i == 1): # Silhouette Mask presence_sil_mask = presence_sil_mask.detach().cpu().numpy() # Log plot to wandb if not mapping: fig_title = f"Time-Step: {iter_time_idx} | Iter: {i} | Frame: {data['id']}" else: fig_title = f"Time-Step: {online_time_idx} | Iter: {i} | Frame: {data['id']}" plot_rgbd_silhouette(data['im'], data['depth'], im, rastered_depth, presence_sil_mask, diff_depth_rmse, psnr, rmse, fig_title, wandb_run=wandb_run, wandb_step=wandb_step, wandb_title=f"{stage} Qual Viz") def eval(dataset, final_params, num_frames, eval_dir, sil_thres, mapping_iters, add_new_gaussians, wandb_run=None, wandb_save_qual=False): print("Evaluating Final Parameters ...") psnr_list = [] rmse_list = [] lpips_list = [] ssim_list = [] plot_dir = os.path.join(eval_dir, "plots") os.makedirs(plot_dir, exist_ok=True) gt_w2c_list = [] for time_idx in tqdm(range(num_frames)): # Get RGB-D Data & Camera Parameters color, depth, intrinsics, pose = dataset[time_idx] gt_w2c = torch.linalg.inv(pose) gt_w2c_list.append(gt_w2c) intrinsics = intrinsics[:3, :3] # Process RGB-D Data color = color.permute(2, 0, 1) / 255 # (H, W, C) -> (C, H, W) depth = depth.permute(2, 0, 1) # (H, W, C) -> (C, H, W) # Process Camera Parameters w2c = torch.linalg.inv(pose) if time_idx == 0: first_frame_w2c = w2c # Setup Camera cam = setup_camera(color.shape[2], color.shape[1], intrinsics.cpu().numpy(), w2c.detach().cpu().numpy()) # Define current frame data curr_data = {'cam': cam, 'im': color, 'depth': depth, 'id': time_idx, 'intrinsics': intrinsics, 'w2c': w2c} # Initialize Render Variables rendervar = params2rendervar(final_params) depth_sil_rendervar = params2depthplussilhouette(final_params, w2c) # Render Depth & Silhouette depth_sil, _, _, = Renderer(raster_settings=curr_data['cam'])(**depth_sil_rendervar) rastered_depth = depth_sil[0, :, :].unsqueeze(0) valid_depth_mask = (curr_data['depth'] > 0) silhouette = depth_sil[1, :, :] presence_sil_mask = (silhouette > sil_thres) # Render RGB and Calculate PSNR im, radius, _, = Renderer(raster_settings=curr_data['cam'])(**rendervar) if mapping_iters==0 and not add_new_gaussians: weighted_im = im * presence_sil_mask weighted_gt_im = curr_data['im'] * presence_sil_mask psnr = calc_psnr(weighted_im, weighted_gt_im).mean() ssim = ms_ssim(weighted_im.unsqueeze(0).cpu(), weighted_gt_im.unsqueeze(0).cpu(), data_range=1.0, size_average=True) lpips_score = loss_fn_alex(torch.clamp(weighted_im.unsqueeze(0), 0.0, 1.0), torch.clamp(weighted_gt_im.unsqueeze(0), 0.0, 1.0)).item() else: psnr = calc_psnr(im, curr_data['im']).mean() ssim = ms_ssim(im.unsqueeze(0).cpu(), curr_data['im'].unsqueeze(0).cpu(), data_range=1.0, size_average=True) lpips_score = loss_fn_alex(torch.clamp(im.unsqueeze(0), 0.0, 1.0), torch.clamp(curr_data['im'].unsqueeze(0), 0.0, 1.0)).item() psnr_list.append(psnr.cpu().numpy()) ssim_list.append(ssim.cpu().numpy()) lpips_list.append(lpips_score) # Compute Depth RMSE if mapping_iters==0 and not add_new_gaussians: diff_depth_rmse = torch.sqrt((((rastered_depth - curr_data['depth']) * presence_sil_mask) ** 2)) diff_depth_rmse = diff_depth_rmse * valid_depth_mask rmse = diff_depth_rmse.sum() / valid_depth_mask.sum() else: diff_depth_rmse = torch.sqrt(((rastered_depth - curr_data['depth']) ** 2)) diff_depth_rmse = diff_depth_rmse * valid_depth_mask rmse = diff_depth_rmse.sum() / valid_depth_mask.sum() rmse_list.append(rmse.cpu().numpy()) # Plot the Ground Truth and Rasterized RGB & Depth, along with Silhouette fig_title = "Time Step: {}".format(time_idx) plot_name = "%04d" % time_idx presence_sil_mask = presence_sil_mask.detach().cpu().numpy() if wandb_run is None: plot_rgbd_silhouette(color, depth, im, rastered_depth, presence_sil_mask, diff_depth_rmse, psnr, rmse, fig_title, plot_dir, plot_name=plot_name, save_plot=True) elif wandb_save_qual: plot_rgbd_silhouette(color, depth, im, rastered_depth, presence_sil_mask, diff_depth_rmse, psnr, rmse, fig_title, plot_dir, plot_name=plot_name, save_plot=True, wandb_run=wandb_run, wandb_step=None, wandb_title="Eval Qual Viz") # Compute Average Metrics psnr_list = np.array(psnr_list) rmse_list = np.array(rmse_list) ssim_list = np.array(ssim_list) lpips_list = np.array(lpips_list) avg_psnr = psnr_list.mean() avg_rmse = rmse_list.mean() avg_ssim = ssim_list.mean() avg_lpips = lpips_list.mean() print("Average PSNR: {:.2f}".format(avg_psnr)) print("Average Depth RMSE: {:.2f}".format(avg_rmse)) print("Average MS-SSIM: {:.2f}".format(avg_ssim)) print("Average LPIPS: {:.2f}".format(avg_lpips)) if wandb_run is not None: wandb_run.log({"Average PSNR": avg_psnr, "Average Depth RMSE": avg_rmse, "Average MS-SSIM": avg_ssim, "Average LPIPS": avg_lpips}) # # Save metric lists as text files # np.savetxt(os.path.join(eval_dir, "psnr.txt"), psnr_list) # np.savetxt(os.path.join(eval_dir, "rmse.txt"), rmse_list) # np.savetxt(os.path.join(eval_dir, "ssim.txt"), ssim_list) # np.savetxt(os.path.join(eval_dir, "lpips.txt"), lpips_list) # # Plot PSNR & RMSE as line plots # fig, axs = plt.subplots(1, 2, figsize=(12, 4)) # axs[0].plot(np.arange(num_frames), psnr_list) # axs[0].set_title("RGB PSNR") # axs[0].set_xlabel("Time Step") # axs[0].set_ylabel("PSNR") # axs[1].plot(np.arange(num_frames), rmse_list) # axs[1].set_title("Depth RMSE") # axs[1].set_xlabel("Time Step") # axs[1].set_ylabel("RMSE") # fig.suptitle("Average PSNR: {:.2f}, Average Depth RMSE: {:.2f}".format(avg_psnr, avg_rmse), y=1.05, fontsize=16) # plt.savefig(os.path.join(eval_dir, "metrics.png"), bbox_inches='tight') # if wandb_run is not None: # wandb_run.log({"Eval Metrics": fig}) # plt.close()