import sys sys.path.append('../') sys.path.append("../submodules") sys.path.append('../submodules/RoMa') from matplotlib import pyplot as plt from PIL import Image import torch import numpy as np #from tqdm import tqdm_notebook as tqdm from tqdm import tqdm from scipy.cluster.vq import kmeans, vq from scipy.spatial.distance import cdist import torch.nn.functional as F from romatch import roma_outdoor, roma_indoor from utils.sh_utils import RGB2SH from romatch.utils import get_tuple_transform_ops def pairwise_distances(matrix): """ Computes the pairwise Euclidean distances between all vectors in the input matrix. Args: matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality. Returns: torch.Tensor: Pairwise distance matrix of shape [N, N]. """ # Compute squared pairwise distances squared_diff = torch.cdist(matrix, matrix, p=2) return squared_diff def k_closest_vectors(matrix, k): """ Finds the k-closest vectors for each vector in the input matrix based on Euclidean distance. Args: matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality. k (int): Number of closest vectors to return for each vector. Returns: torch.Tensor: Indices of the k-closest vectors for each vector, excluding the vector itself. """ # Compute pairwise distances distances = pairwise_distances(matrix) # For each vector, sort distances and get the indices of the k-closest vectors (excluding itself) # Set diagonal distances to infinity to exclude the vector itself from the nearest neighbors distances.fill_diagonal_(float('inf')) # Get the indices of the k smallest distances (k-closest vectors) _, indices = torch.topk(distances, k, largest=False, dim=1) return indices def select_cameras_kmeans(cameras, K): """ Selects K cameras from a set using K-means clustering. Args: cameras: NumPy array of shape (N, 16), representing N cameras with their 4x4 homogeneous matrices flattened. K: Number of clusters (cameras to select). Returns: selected_indices: List of indices of the cameras closest to the cluster centers. """ # Ensure input is a NumPy array if not isinstance(cameras, np.ndarray): cameras = np.asarray(cameras) if cameras.shape[1] != 16: raise ValueError("Each camera must have 16 values corresponding to a flattened 4x4 matrix.") # Perform K-means clustering cluster_centers, _ = kmeans(cameras, K) # Assign each camera to a cluster and find distances to cluster centers cluster_assignments, _ = vq(cameras, cluster_centers) # Find the camera nearest to each cluster center selected_indices = [] for k in range(K): cluster_members = cameras[cluster_assignments == k] distances = cdist([cluster_centers[k]], cluster_members)[0] nearest_camera_idx = np.where(cluster_assignments == k)[0][np.argmin(distances)] selected_indices.append(nearest_camera_idx) return selected_indices def compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, device="cuda", verbose=False, output_dict={}): """ Computes the warp and confidence between two viewpoint cameras using the roma_model. Args: viewpoint_cam1: Source viewpoint camera. viewpoint_cam2: Target viewpoint camera. roma_model: Pre-trained Roma model for correspondence matching. device: Device to run the computation on. verbose: If True, displays the images. Returns: certainty: Confidence tensor. warp: Warp tensor. imB: Processed image B as numpy array. """ # Prepare images imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0) imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0) imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8)) imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8)) if verbose: fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(16, 8)) cax1 = ax[0].imshow(imA) ax[0].set_title("Image 1") cax2 = ax[1].imshow(imB) ax[1].set_title("Image 2") fig.colorbar(cax1, ax=ax[0]) fig.colorbar(cax2, ax=ax[1]) for axis in ax: axis.axis('off') # Save the figure into the dictionary output_dict[f'image_pair'] = fig # Transform images ws, hs = roma_model.w_resized, roma_model.h_resized test_transform = get_tuple_transform_ops(resize=(hs, ws), normalize=True) im_A, im_B = test_transform((imA, imB)) batch = {"im_A": im_A[None].to(device), "im_B": im_B[None].to(device)} # Forward pass through Roma model corresps = roma_model.forward(batch) if not roma_model.symmetric else roma_model.forward_symmetric(batch) finest_scale = 1 hs, ws = roma_model.upsample_res if roma_model.upsample_preds else (hs, ws) # Process certainty and warp certainty = corresps[finest_scale]["certainty"] im_A_to_im_B = corresps[finest_scale]["flow"] if roma_model.attenuate_cert: low_res_certainty = F.interpolate( corresps[16]["certainty"], size=(hs, ws), align_corners=False, mode="bilinear" ) certainty -= 0.5 * low_res_certainty * (low_res_certainty < 0) # Upsample predictions if needed if roma_model.upsample_preds: im_A_to_im_B = F.interpolate( im_A_to_im_B, size=(hs, ws), align_corners=False, mode="bilinear" ) certainty = F.interpolate( certainty, size=(hs, ws), align_corners=False, mode="bilinear" ) # Convert predictions to final format im_A_to_im_B = im_A_to_im_B.permute(0, 2, 3, 1) im_A_coords = torch.stack(torch.meshgrid( torch.linspace(-1 + 1 / hs, 1 - 1 / hs, hs, device=device), torch.linspace(-1 + 1 / ws, 1 - 1 / ws, ws, device=device), indexing='ij' ), dim=0).permute(1, 2, 0).unsqueeze(0).expand(im_A_to_im_B.size(0), -1, -1, -1) warp = torch.cat((im_A_coords, im_A_to_im_B), dim=-1) certainty = certainty.sigmoid() return certainty[0, 0], warp[0], np.array(imB) def resize_batch(tensors_3d, tensors_4d, target_shape): """ Resizes a batch of tensors with shapes [B, H, W] and [B, H, W, 4] to the target spatial dimensions. Args: tensors_3d: Tensor of shape [B, H, W]. tensors_4d: Tensor of shape [B, H, W, 4]. target_shape: Tuple (target_H, target_W) specifying the target spatial dimensions. Returns: resized_tensors_3d: Tensor of shape [B, target_H, target_W]. resized_tensors_4d: Tensor of shape [B, target_H, target_W, 4]. """ target_H, target_W = target_shape # Resize [B, H, W] tensor resized_tensors_3d = F.interpolate( tensors_3d.unsqueeze(1), size=(target_H, target_W), mode="bilinear", align_corners=False ).squeeze(1) # Resize [B, H, W, 4] tensor B, _, _, C = tensors_4d.shape resized_tensors_4d = F.interpolate( tensors_4d.permute(0, 3, 1, 2), size=(target_H, target_W), mode="bilinear", align_corners=False ).permute(0, 2, 3, 1) return resized_tensors_3d, resized_tensors_4d def aggregate_confidences_and_warps(viewpoint_stack, closest_indices, roma_model, source_idx, verbose=False, output_dict={}): """ Aggregates confidences and warps by iterating over the nearest neighbors of the source viewpoint. Args: viewpoint_stack: Stack of viewpoint cameras. closest_indices: Indices of the nearest neighbors for each viewpoint. roma_model: Pre-trained Roma model. source_idx: Index of the source viewpoint. verbose: If True, displays intermediate results. Returns: certainties_max: Aggregated maximum confidences. warps_max: Aggregated warps corresponding to maximum confidences. certainties_max_idcs: Pixel-wise index of the image from which we taken the best matching. imB_compound: List of the neighboring images. """ certainties_all, warps_all, imB_compound = [], [], [] for nn in tqdm(closest_indices[source_idx]): viewpoint_cam1 = viewpoint_stack[source_idx] viewpoint_cam2 = viewpoint_stack[nn] certainty, warp, imB = compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, verbose=verbose, output_dict=output_dict) certainties_all.append(certainty) warps_all.append(warp) imB_compound.append(imB) certainties_all = torch.stack(certainties_all, dim=0) target_shape = imB_compound[0].shape[:2] if verbose: print("certainties_all.shape:", certainties_all.shape) print("torch.stack(warps_all, dim=0).shape:", torch.stack(warps_all, dim=0).shape) print("target_shape:", target_shape) certainties_all_resized, warps_all_resized = resize_batch(certainties_all, torch.stack(warps_all, dim=0), target_shape ) if verbose: print("warps_all_resized.shape:", warps_all_resized.shape) for n, cert in enumerate(certainties_all): fig, ax = plt.subplots() cax = ax.imshow(cert.cpu().numpy(), cmap='viridis') fig.colorbar(cax, ax=ax) ax.set_title("Pixel-wise Confidence") output_dict[f'certainty_{n}'] = fig for n, warp in enumerate(warps_all): fig, ax = plt.subplots() cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis') fig.colorbar(cax, ax=ax) ax.set_title("Pixel-wise warp") output_dict[f'warp_resized_{n}'] = fig for n, cert in enumerate(certainties_all_resized): fig, ax = plt.subplots() cax = ax.imshow(cert.cpu().numpy(), cmap='viridis') fig.colorbar(cax, ax=ax) ax.set_title("Pixel-wise Confidence resized") output_dict[f'certainty_resized_{n}'] = fig for n, warp in enumerate(warps_all_resized): fig, ax = plt.subplots() cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis') fig.colorbar(cax, ax=ax) ax.set_title("Pixel-wise warp resized") output_dict[f'warp_resized_{n}'] = fig certainties_max, certainties_max_idcs = torch.max(certainties_all_resized, dim=0) H, W = certainties_max.shape warps_max = warps_all_resized[certainties_max_idcs, torch.arange(H).unsqueeze(1), torch.arange(W)] return certainties_max, warps_max, certainties_max_idcs, imB_compound, certainties_all_resized, warps_all_resized def extract_keypoints_and_colors(imA, imB_compound, certainties_max, certainties_max_idcs, matches, roma_model, verbose=False, output_dict={}): """ Extracts keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound). Args: imA: Source image as a NumPy array (H_A, W_A, C). imB_compound: List of target images as NumPy arrays [(H_B, W_B, C), ...]. certainties_max: Tensor of pixel-wise maximum confidences. certainties_max_idcs: Tensor of pixel-wise indices for the best matches. matches: Matches in normalized coordinates. roma_model: Roma model instance for keypoint operations. verbose: if to show intermediate outputs and visualize results Returns: kptsA_np: Keypoints in imA in normalized coordinates. kptsB_np: Keypoints in imB in normalized coordinates. kptsA_color: Colors of keypoints in imA. kptsB_color: Colors of keypoints in imB based on certainties_max_idcs. """ H_A, W_A, _ = imA.shape H, W = certainties_max.shape # Convert matches to pixel coordinates kptsA, kptsB = roma_model.to_pixel_coordinates( matches, W_A, H_A, H, W # W, H ) kptsA_np = kptsA.detach().cpu().numpy() kptsB_np = kptsB.detach().cpu().numpy() kptsA_np = kptsA_np[:, [1, 0]] if verbose: fig, ax = plt.subplots(figsize=(12, 6)) cax = ax.imshow(imA) ax.set_title("Reference image, imA") output_dict[f'reference_image'] = fig fig, ax = plt.subplots(figsize=(12, 6)) cax = ax.imshow(imB_compound[0]) ax.set_title("Image to compare to image, imB_compound") output_dict[f'imB_compound'] = fig fig, ax = plt.subplots(figsize=(12, 6)) cax = ax.imshow(np.flipud(imA)) cax = ax.scatter(kptsA_np[:, 0], H_A - kptsA_np[:, 1], s=.03) ax.set_title("Keypoints in imA") ax.set_xlim(0, W_A) ax.set_ylim(0, H_A) output_dict[f'kptsA'] = fig fig, ax = plt.subplots(figsize=(12, 6)) cax = ax.imshow(np.flipud(imB_compound[0])) cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03) ax.set_title("Keypoints in imB") ax.set_xlim(0, W_A) ax.set_ylim(0, H_A) output_dict[f'kptsB'] = fig # Keypoints are in format (row, column) so the first value is alwain in range [0;height] and second is in range[0;width] kptsA_np = kptsA.detach().cpu().numpy() kptsB_np = kptsB.detach().cpu().numpy() # Extract colors for keypoints in imA (vectorized) # New experimental version kptsA_x = np.round(kptsA_np[:, 0] / 1.).astype(int) kptsA_y = np.round(kptsA_np[:, 1] / 1.).astype(int) kptsA_color = imA[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)] # Create a composite image from imB_compound imB_compound_np = np.stack(imB_compound, axis=0) H_B, W_B, _ = imB_compound[0].shape # Extract colors for keypoints in imB using certainties_max_idcs imB_np = imB_compound_np[ certainties_max_idcs.detach().cpu().numpy(), np.arange(H).reshape(-1, 1), np.arange(W) ] if verbose: print("imB_np.shape:", imB_np.shape) print("imB_np:", imB_np) fig, ax = plt.subplots(figsize=(12, 6)) cax = ax.imshow(np.flipud(imB_np)) cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03) ax.set_title("np.flipud(imB_np[0]") ax.set_xlim(0, W_A) ax.set_ylim(0, H_A) output_dict[f'np.flipud(imB_np[0]'] = fig kptsB_x = np.round(kptsB_np[:, 0]).astype(int) kptsB_y = np.round(kptsB_np[:, 1]).astype(int) certainties_max_idcs_np = certainties_max_idcs.detach().cpu().numpy() kptsB_proj_matrices_idx = certainties_max_idcs_np[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)] kptsB_color = imB_compound_np[kptsB_proj_matrices_idx, np.clip(kptsB_y, 0, H - 1), np.clip(kptsB_x, 0, W - 1)] # Normalize keypoints in both images kptsA_np[:, 0] = kptsA_np[:, 0] / H * 2.0 - 1.0 kptsA_np[:, 1] = kptsA_np[:, 1] / W * 2.0 - 1.0 kptsB_np[:, 0] = kptsB_np[:, 0] / W_B * 2.0 - 1.0 kptsB_np[:, 1] = kptsB_np[:, 1] / H_B * 2.0 - 1.0 return kptsA_np[:, [1, 0]], kptsB_np, kptsB_proj_matrices_idx, kptsA_color, kptsB_color def prepare_tensor(input_array, device): """ Converts an input array to a torch tensor, clones it, and detaches it for safe computation. Args: input_array (array-like): The input array to convert. device (str or torch.device): The device to move the tensor to. Returns: torch.Tensor: A detached tensor clone of the input array on the specified device. """ if not isinstance(input_array, torch.Tensor): return torch.tensor(input_array, dtype=torch.float32).to(device).clone().detach() return input_array.clone().detach().to(device).to(torch.float32) def triangulate_points(P1, P2, k1_x, k1_y, k2_x, k2_y, device="cuda"): """ Solves for a batch of 3D points given batches of projection matrices and corresponding image points. Parameters: - P1, P2: Tensors of projection matrices of size (batch_size, 4, 4) or (4, 4) - k1_x, k1_y: Tensors of shape (batch_size,) - k2_x, k2_y: Tensors of shape (batch_size,) Returns: - X: A tensor containing the 3D homogeneous coordinates, shape (batch_size, 4) """ EPS = 1e-4 # Ensure inputs are tensors P1 = prepare_tensor(P1, device) P2 = prepare_tensor(P2, device) k1_x = prepare_tensor(k1_x, device) k1_y = prepare_tensor(k1_y, device) k2_x = prepare_tensor(k2_x, device) k2_y = prepare_tensor(k2_y, device) batch_size = k1_x.shape[0] # Expand P1 and P2 if they are not batched if P1.ndim == 2: P1 = P1.unsqueeze(0).expand(batch_size, -1, -1) if P2.ndim == 2: P2 = P2.unsqueeze(0).expand(batch_size, -1, -1) # Extract columns from P1 and P2 P1_0 = P1[:, :, 0] # Shape: (batch_size, 4) P1_1 = P1[:, :, 1] P1_2 = P1[:, :, 2] P2_0 = P2[:, :, 0] P2_1 = P2[:, :, 1] P2_2 = P2[:, :, 2] # Reshape kx and ky to (batch_size, 1) k1_x = k1_x.view(-1, 1) k1_y = k1_y.view(-1, 1) k2_x = k2_x.view(-1, 1) k2_y = k2_y.view(-1, 1) # Construct the equations for each batch # For camera 1 A1 = P1_0 - k1_x * P1_2 # Shape: (batch_size, 4) A2 = P1_1 - k1_y * P1_2 # For camera 2 A3 = P2_0 - k2_x * P2_2 A4 = P2_1 - k2_y * P2_2 # Stack the equations A = torch.stack([A1, A2, A3, A4], dim=1) # Shape: (batch_size, 4, 4) # Right-hand side (constants) b = -A[:, :, 3] # Shape: (batch_size, 4) A_reduced = A[:, :, :3] # Coefficients of x, y, z # Solve using torch.linalg.lstsq (supports batching) X_xyz = torch.linalg.lstsq(A_reduced, b.unsqueeze(2)).solution.squeeze(2) # Shape: (batch_size, 3) # Append 1 to get homogeneous coordinates ones = torch.ones((batch_size, 1), dtype=torch.float32, device=X_xyz.device) X = torch.cat([X_xyz, ones], dim=1) # Shape: (batch_size, 4) # Now compute the errors of projections. seeked_splats_proj1 = (X.unsqueeze(1) @ P1).squeeze(1) seeked_splats_proj1 = seeked_splats_proj1 / (EPS + seeked_splats_proj1[:, [3]]) seeked_splats_proj2 = (X.unsqueeze(1) @ P2).squeeze(1) seeked_splats_proj2 = seeked_splats_proj2 / (EPS + seeked_splats_proj2[:, [3]]) proj1_target = torch.concat([k1_x, k1_y], dim=1) proj2_target = torch.concat([k2_x, k2_y], dim=1) errors_proj1 = torch.abs(seeked_splats_proj1[:, :2] - proj1_target).sum(1).detach().cpu().numpy() errors_proj2 = torch.abs(seeked_splats_proj2[:, :2] - proj2_target).sum(1).detach().cpu().numpy() return X, errors_proj1, errors_proj2 def select_best_keypoints( NNs_triangulated_points, NNs_errors_proj1, NNs_errors_proj2, device="cuda"): """ From all the points fitted to keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound). Args: NNs_triangulated_points: torch tensor with keypoints coordinates (num_nns, num_points, dim). dim can be arbitrary, usually 3 or 4(for homogeneous representation). NNs_errors_proj1: numpy array with projection error of the estimated keypoint on the reference frame (num_nns, num_points). NNs_errors_proj2: numpy array with projection error of the estimated keypoint on the neighbor frame (num_nns, num_points). Returns: selected_keypoints: keypoints with the best score. """ NNs_errors_proj = np.maximum(NNs_errors_proj1, NNs_errors_proj2) # Convert indices to PyTorch tensor indices = torch.from_numpy(np.argmin(NNs_errors_proj, axis=0)).long().to(device) # Create index tensor for the second dimension n_indices = torch.arange(NNs_triangulated_points.shape[1]).long().to(device) # Use advanced indexing to select elements NNs_triangulated_points_selected = NNs_triangulated_points[indices, n_indices, :] # Shape: [N, k] return NNs_triangulated_points_selected, np.min(NNs_errors_proj, axis=0) import time from collections import defaultdict from tqdm import tqdm # def init_gaussians_with_corr_profiled(gaussians, scene, cfg, device, verbose=False, roma_model=None): # timings = defaultdict(list) # To accumulate timings # if roma_model is None: # if cfg.roma_model == "indoors": # roma_model = roma_indoor(device=device) # else: # roma_model = roma_outdoor(device=device) # roma_model.upsample_preds = False # roma_model.symmetric = False # M = cfg.matches_per_ref # upper_thresh = roma_model.sample_thresh # scaling_factor = cfg.scaling_factor # expansion_factor = 1 # keypoint_fit_error_tolerance = cfg.proj_err_tolerance # visualizations = {} # viewpoint_stack = scene.getTrainCameras().copy() # NUM_REFERENCE_FRAMES = min(cfg.num_refs, len(viewpoint_stack)) # NUM_NNS_PER_REFERENCE = min(cfg.nns_per_ref, len(viewpoint_stack)) # viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0) # selected_indices = select_cameras_kmeans(cameras=viewpoint_cam_all.detach().cpu().numpy(), K=NUM_REFERENCE_FRAMES) # selected_indices = sorted(selected_indices) # viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0) # closest_indices = k_closest_vectors(viewpoint_cam_all, NUM_NNS_PER_REFERENCE) # closest_indices_selected = closest_indices[:, :].detach().cpu().numpy() # all_new_xyz = [] # all_new_features_dc = [] # all_new_features_rest = [] # all_new_opacities = [] # all_new_scaling = [] # all_new_rotation = [] # # Dummy first pass to initialize model # with torch.no_grad(): # viewpoint_cam1 = viewpoint_stack[0] # viewpoint_cam2 = viewpoint_stack[1] # imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0) # imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0) # imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8)) # imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8)) # warp, certainty_warp = roma_model.match(imA, imB, device=device) # del warp, certainty_warp # torch.cuda.empty_cache() # # Main Loop over source_idx # for source_idx in tqdm(sorted(selected_indices), desc="Profiling source frames"): # # =================== Step 1: Aggregate Confidences and Warps =================== # start = time.time() # viewpoint_cam1 = viewpoint_stack[source_idx] # viewpoint_cam2 = viewpoint_stack[closest_indices_selected[source_idx,0]] # imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0) # imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0) # imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8)) # imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8)) # warp, certainty_warp = roma_model.match(imA, imB, device=device) # certainties_max, warps_max, certainties_max_idcs, imB_compound, certainties_all, warps_all = aggregate_confidences_and_warps( # viewpoint_stack=viewpoint_stack, # closest_indices=closest_indices_selected, # roma_model=roma_model, # source_idx=source_idx, # verbose=verbose, # output_dict=visualizations # ) # certainties_max = certainty_warp # with torch.no_grad(): # warps_all = warps.unsqueeze(0) # timings['aggregation_warp_certainty'].append(time.time() - start) # # =================== Step 2: Good Samples Selection =================== # start = time.time() # certainty = certainties_max.reshape(-1).clone() # certainty[certainty > upper_thresh] = 1 # good_samples = torch.multinomial(certainty, num_samples=min(expansion_factor * M, len(certainty)), replacement=False) # timings['good_samples_selection'].append(time.time() - start) # # =================== Step 3: Triangulate Keypoints for Each NN =================== # reference_image_dict = { # "triangulated_points": [], # "triangulated_points_errors_proj1": [], # "triangulated_points_errors_proj2": [] # } # start = time.time() # for NN_idx in range(len(warps_all)): # matches_NN = warps_all[NN_idx].reshape(-1, 4)[good_samples] # # Extract keypoints and colors # kptsA_np, kptsB_np, kptsB_proj_matrices_idcs, kptsA_color, kptsB_color = extract_keypoints_and_colors( # imA, imB_compound, certainties_max, certainties_max_idcs, matches_NN, roma_model # ) # proj_matrices_A = viewpoint_stack[source_idx].full_proj_transform # proj_matrices_B = viewpoint_stack[closest_indices_selected[source_idx, NN_idx]].full_proj_transform # triangulated_points, triangulated_points_errors_proj1, triangulated_points_errors_proj2 = triangulate_points( # P1=torch.stack([proj_matrices_A] * M, axis=0), # P2=torch.stack([proj_matrices_B] * M, axis=0), # k1_x=kptsA_np[:M, 0], k1_y=kptsA_np[:M, 1], # k2_x=kptsB_np[:M, 0], k2_y=kptsB_np[:M, 1]) # reference_image_dict["triangulated_points"].append(triangulated_points) # reference_image_dict["triangulated_points_errors_proj1"].append(triangulated_points_errors_proj1) # reference_image_dict["triangulated_points_errors_proj2"].append(triangulated_points_errors_proj2) # timings['triangulation_per_NN'].append(time.time() - start) # # =================== Step 4: Select Best Triangulated Points =================== # start = time.time() # NNs_triangulated_points_selected, NNs_triangulated_points_selected_proj_errors = select_best_keypoints( # NNs_triangulated_points=torch.stack(reference_image_dict["triangulated_points"], dim=0), # NNs_errors_proj1=np.stack(reference_image_dict["triangulated_points_errors_proj1"], axis=0), # NNs_errors_proj2=np.stack(reference_image_dict["triangulated_points_errors_proj2"], axis=0)) # timings['select_best_keypoints'].append(time.time() - start) # # =================== Step 5: Create New Gaussians =================== # start = time.time() # viewpoint_cam1 = viewpoint_stack[source_idx] # N = len(NNs_triangulated_points_selected) # new_xyz = NNs_triangulated_points_selected[:, :-1] # all_new_xyz.append(new_xyz) # all_new_features_dc.append(RGB2SH(torch.tensor(kptsA_color.astype(np.float32) / 255.)).unsqueeze(1)) # all_new_features_rest.append(torch.stack([gaussians._features_rest[-1].clone().detach() * 0.] * N, dim=0)) # mask_bad_points = torch.tensor( # NNs_triangulated_points_selected_proj_errors > keypoint_fit_error_tolerance, # dtype=torch.float32).unsqueeze(1).to(device) # all_new_opacities.append(torch.stack([gaussians._opacity[-1].clone().detach()] * N, dim=0) * 0. - mask_bad_points * (1e1)) # dist_points_to_cam1 = torch.linalg.norm(viewpoint_cam1.camera_center.clone().detach() - new_xyz, dim=1, ord=2) # all_new_scaling.append(gaussians.scaling_inverse_activation((dist_points_to_cam1 * scaling_factor).unsqueeze(1).repeat(1, 3))) # all_new_rotation.append(torch.stack([gaussians._rotation[-1].clone().detach()] * N, dim=0)) # timings['save_gaussians'].append(time.time() - start) # # =================== Final Densification Postfix =================== # start = time.time() # all_new_xyz = torch.cat(all_new_xyz, dim=0) # all_new_features_dc = torch.cat(all_new_features_dc, dim=0) # new_tmp_radii = torch.zeros(all_new_xyz.shape[0]) # prune_mask = torch.ones(all_new_xyz.shape[0], dtype=torch.bool) # gaussians.densification_postfix( # all_new_xyz[prune_mask].to(device), # all_new_features_dc[prune_mask].to(device), # torch.cat(all_new_features_rest, dim=0)[prune_mask].to(device), # torch.cat(all_new_opacities, dim=0)[prune_mask].to(device), # torch.cat(all_new_scaling, dim=0)[prune_mask].to(device), # torch.cat(all_new_rotation, dim=0)[prune_mask].to(device), # new_tmp_radii[prune_mask].to(device) # ) # timings['final_densification_postfix'].append(time.time() - start) # # =================== Print Profiling Results =================== # print("\n=== Profiling Summary (average per frame) ===") # for key, times in timings.items(): # print(f"{key:35s}: {sum(times) / len(times):.4f} sec (total {sum(times):.2f} sec)") # return viewpoint_stack, closest_indices_selected, visualizations def extract_keypoints_and_colors_single(imA, imB, matches, roma_model, verbose=False, output_dict={}): """ Extracts keypoints and corresponding colors from a source image (imA) and a single target image (imB). Args: imA: Source image as a NumPy array (H_A, W_A, C). imB: Target image as a NumPy array (H_B, W_B, C). matches: Matches in normalized coordinates (torch.Tensor). roma_model: Roma model instance for keypoint operations. verbose: If True, outputs intermediate visualizations. Returns: kptsA_np: Keypoints in imA (normalized). kptsB_np: Keypoints in imB (normalized). kptsA_color: Colors of keypoints in imA. kptsB_color: Colors of keypoints in imB. """ H_A, W_A, _ = imA.shape H_B, W_B, _ = imB.shape # Convert matches to pixel coordinates # Matches format: (B, 4) = (x1_norm, y1_norm, x2_norm, y2_norm) kptsA = matches[:, :2] # [N, 2] kptsB = matches[:, 2:] # [N, 2] # Scale normalized coordinates [-1,1] to pixel coordinates kptsA_pix = torch.zeros_like(kptsA) kptsB_pix = torch.zeros_like(kptsB) # Important! [Normalized to pixel space] kptsA_pix[:, 0] = (kptsA[:, 0] + 1) * (W_A - 1) / 2 kptsA_pix[:, 1] = (kptsA[:, 1] + 1) * (H_A - 1) / 2 kptsB_pix[:, 0] = (kptsB[:, 0] + 1) * (W_B - 1) / 2 kptsB_pix[:, 1] = (kptsB[:, 1] + 1) * (H_B - 1) / 2 kptsA_np = kptsA_pix.detach().cpu().numpy() kptsB_np = kptsB_pix.detach().cpu().numpy() # Extract colors kptsA_x = np.round(kptsA_np[:, 0]).astype(int) kptsA_y = np.round(kptsA_np[:, 1]).astype(int) kptsB_x = np.round(kptsB_np[:, 0]).astype(int) kptsB_y = np.round(kptsB_np[:, 1]).astype(int) kptsA_color = imA[np.clip(kptsA_y, 0, H_A-1), np.clip(kptsA_x, 0, W_A-1)] kptsB_color = imB[np.clip(kptsB_y, 0, H_B-1), np.clip(kptsB_x, 0, W_B-1)] # Normalize keypoints into [-1, 1] for downstream triangulation kptsA_np_norm = np.zeros_like(kptsA_np) kptsB_np_norm = np.zeros_like(kptsB_np) kptsA_np_norm[:, 0] = kptsA_np[:, 0] / (W_A - 1) * 2.0 - 1.0 kptsA_np_norm[:, 1] = kptsA_np[:, 1] / (H_A - 1) * 2.0 - 1.0 kptsB_np_norm[:, 0] = kptsB_np[:, 0] / (W_B - 1) * 2.0 - 1.0 kptsB_np_norm[:, 1] = kptsB_np[:, 1] / (H_B - 1) * 2.0 - 1.0 return kptsA_np_norm, kptsB_np_norm, kptsA_color, kptsB_color def init_gaussians_with_corr_profiled(gaussians, scene, cfg, device, verbose=False, roma_model=None): timings = defaultdict(list) if roma_model is None: if cfg.roma_model == "indoors": roma_model = roma_indoor(device=device) else: roma_model = roma_outdoor(device=device) roma_model.upsample_preds = False roma_model.symmetric = False M = cfg.matches_per_ref upper_thresh = roma_model.sample_thresh scaling_factor = cfg.scaling_factor expansion_factor = 1 keypoint_fit_error_tolerance = cfg.proj_err_tolerance visualizations = {} viewpoint_stack = scene.getTrainCameras().copy() NUM_REFERENCE_FRAMES = min(cfg.num_refs, len(viewpoint_stack)) NUM_NNS_PER_REFERENCE = 1 # Only ONE neighbor now! viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0) selected_indices = select_cameras_kmeans(cameras=viewpoint_cam_all.detach().cpu().numpy(), K=NUM_REFERENCE_FRAMES) selected_indices = sorted(selected_indices) viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0) closest_indices = k_closest_vectors(viewpoint_cam_all, NUM_NNS_PER_REFERENCE) closest_indices_selected = closest_indices[:, :].detach().cpu().numpy() all_new_xyz = [] all_new_features_dc = [] all_new_features_rest = [] all_new_opacities = [] all_new_scaling = [] all_new_rotation = [] # Dummy first pass to initialize model with torch.no_grad(): viewpoint_cam1 = viewpoint_stack[0] viewpoint_cam2 = viewpoint_stack[1] imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0) imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0) imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8)) imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8)) warp, certainty_warp = roma_model.match(imA, imB, device=device) del warp, certainty_warp torch.cuda.empty_cache() # Main Loop over source_idx for source_idx in tqdm(sorted(selected_indices), desc="Profiling source frames"): # =================== Step 1: Compute Warp and Certainty =================== start = time.time() viewpoint_cam1 = viewpoint_stack[source_idx] NNs=closest_indices_selected.shape[1] viewpoint_cam2 = viewpoint_stack[closest_indices_selected[source_idx, np.random.randint(NNs)]] imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0) imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0) imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8)) imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8)) warp, certainty_warp = roma_model.match(imA, imB, device=device) certainties_max = certainty_warp # New manual sampling timings['aggregation_warp_certainty'].append(time.time() - start) # =================== Step 2: Good Samples Selection =================== start = time.time() certainty = certainties_max.reshape(-1).clone() certainty[certainty > upper_thresh] = 1 good_samples = torch.multinomial(certainty, num_samples=min(expansion_factor * M, len(certainty)), replacement=False) timings['good_samples_selection'].append(time.time() - start) # =================== Step 3: Triangulate Keypoints =================== reference_image_dict = { "triangulated_points": [], "triangulated_points_errors_proj1": [], "triangulated_points_errors_proj2": [] } start = time.time() matches_NN = warp.reshape(-1, 4)[good_samples] # Convert matches to pixel coordinates kptsA_np, kptsB_np, kptsA_color, kptsB_color = extract_keypoints_and_colors_single( np.array(imA).astype(np.uint8), np.array(imB).astype(np.uint8), matches_NN, roma_model ) proj_matrices_A = viewpoint_stack[source_idx].full_proj_transform proj_matrices_B = viewpoint_stack[closest_indices_selected[source_idx, 0]].full_proj_transform triangulated_points, triangulated_points_errors_proj1, triangulated_points_errors_proj2 = triangulate_points( P1=torch.stack([proj_matrices_A] * M, axis=0), P2=torch.stack([proj_matrices_B] * M, axis=0), k1_x=kptsA_np[:M, 0], k1_y=kptsA_np[:M, 1], k2_x=kptsB_np[:M, 0], k2_y=kptsB_np[:M, 1]) reference_image_dict["triangulated_points"].append(triangulated_points) reference_image_dict["triangulated_points_errors_proj1"].append(triangulated_points_errors_proj1) reference_image_dict["triangulated_points_errors_proj2"].append(triangulated_points_errors_proj2) timings['triangulation_per_NN'].append(time.time() - start) # =================== Step 4: Select Best Triangulated Points =================== start = time.time() NNs_triangulated_points_selected, NNs_triangulated_points_selected_proj_errors = select_best_keypoints( NNs_triangulated_points=torch.stack(reference_image_dict["triangulated_points"], dim=0), NNs_errors_proj1=np.stack(reference_image_dict["triangulated_points_errors_proj1"], axis=0), NNs_errors_proj2=np.stack(reference_image_dict["triangulated_points_errors_proj2"], axis=0)) timings['select_best_keypoints'].append(time.time() - start) # =================== Step 5: Create New Gaussians =================== start = time.time() viewpoint_cam1 = viewpoint_stack[source_idx] N = len(NNs_triangulated_points_selected) new_xyz = NNs_triangulated_points_selected[:, :-1] all_new_xyz.append(new_xyz) all_new_features_dc.append(RGB2SH(torch.tensor(kptsA_color.astype(np.float32) / 255.)).unsqueeze(1)) all_new_features_rest.append(torch.stack([gaussians._features_rest[-1].clone().detach() * 0.] * N, dim=0)) mask_bad_points = torch.tensor( NNs_triangulated_points_selected_proj_errors > keypoint_fit_error_tolerance, dtype=torch.float32).unsqueeze(1).to(device) all_new_opacities.append(torch.stack([gaussians._opacity[-1].clone().detach()] * N, dim=0) * 0. - mask_bad_points * (1e1)) dist_points_to_cam1 = torch.linalg.norm(viewpoint_cam1.camera_center.clone().detach() - new_xyz, dim=1, ord=2) all_new_scaling.append(gaussians.scaling_inverse_activation((dist_points_to_cam1 * scaling_factor).unsqueeze(1).repeat(1, 3))) all_new_rotation.append(torch.stack([gaussians._rotation[-1].clone().detach()] * N, dim=0)) timings['save_gaussians'].append(time.time() - start) # =================== Final Densification Postfix =================== start = time.time() all_new_xyz = torch.cat(all_new_xyz, dim=0) all_new_features_dc = torch.cat(all_new_features_dc, dim=0) new_tmp_radii = torch.zeros(all_new_xyz.shape[0]) prune_mask = torch.ones(all_new_xyz.shape[0], dtype=torch.bool) gaussians.densification_postfix( all_new_xyz[prune_mask].to(device), all_new_features_dc[prune_mask].to(device), torch.cat(all_new_features_rest, dim=0)[prune_mask].to(device), torch.cat(all_new_opacities, dim=0)[prune_mask].to(device), torch.cat(all_new_scaling, dim=0)[prune_mask].to(device), torch.cat(all_new_rotation, dim=0)[prune_mask].to(device), new_tmp_radii[prune_mask].to(device) ) timings['final_densification_postfix'].append(time.time() - start) # =================== Print Profiling Results =================== print("\n=== Profiling Summary (average per frame) ===") for key, times in timings.items(): print(f"{key:35s}: {sum(times) / len(times):.4f} sec (total {sum(times):.2f} sec)") return viewpoint_stack, closest_indices_selected, visualizations