Spaces:
Running on Zero
Running on Zero
| """COLMAP scene loading + view selection. | |
| Verbatim helpers from the ReSplat reference inference script | |
| (https://github.com/cvg/resplat, scripts/infer_colmap.py, MIT), which shares | |
| the DepthSplat camera conventions used by CodecSplat. | |
| """ | |
| import collections | |
| import os | |
| import struct | |
| import numpy as np | |
| import torch | |
| CameraModel = collections.namedtuple( | |
| "CameraModel", ["model_id", "model_name", "num_params"] | |
| ) | |
| ColmapCamera = collections.namedtuple( | |
| "ColmapCamera", ["id", "model", "width", "height", "params"] | |
| ) | |
| BaseImage = collections.namedtuple( | |
| "ColmapImage", | |
| ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"], | |
| ) | |
| CAMERA_MODELS = { | |
| CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), | |
| CameraModel(model_id=1, model_name="PINHOLE", num_params=4), | |
| CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), | |
| CameraModel(model_id=3, model_name="RADIAL", num_params=5), | |
| CameraModel(model_id=4, model_name="OPENCV", num_params=8), | |
| CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), | |
| CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), | |
| CameraModel(model_id=7, model_name="FOV", num_params=5), | |
| CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), | |
| CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), | |
| CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), | |
| } | |
| CAMERA_MODEL_IDS = dict( | |
| [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS] | |
| ) | |
| class ColmapImage(BaseImage): | |
| def qvec2rotmat(self): | |
| return qvec2rotmat(self.qvec) | |
| def qvec2rotmat(qvec): | |
| return np.array( | |
| [ | |
| [ | |
| 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, | |
| 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], | |
| 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], | |
| ], | |
| [ | |
| 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], | |
| 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, | |
| 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], | |
| ], | |
| [ | |
| 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], | |
| 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], | |
| 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, | |
| ], | |
| ] | |
| ) | |
| def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): | |
| data = fid.read(num_bytes) | |
| return struct.unpack(endian_character + format_char_sequence, data) | |
| def read_extrinsics_binary(path_to_model_file): | |
| images = {} | |
| with open(path_to_model_file, "rb") as fid: | |
| num_reg_images = read_next_bytes(fid, 8, "Q")[0] | |
| for _ in range(num_reg_images): | |
| binary_image_properties = read_next_bytes( | |
| fid, num_bytes=64, format_char_sequence="idddddddi" | |
| ) | |
| image_id = binary_image_properties[0] | |
| qvec = np.array(binary_image_properties[1:5]) | |
| tvec = np.array(binary_image_properties[5:8]) | |
| camera_id = binary_image_properties[8] | |
| image_name = "" | |
| current_char = read_next_bytes(fid, 1, "c")[0] | |
| while current_char != b"\x00": | |
| image_name += current_char.decode("utf-8") | |
| current_char = read_next_bytes(fid, 1, "c")[0] | |
| num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ | |
| 0 | |
| ] | |
| x_y_id_s = read_next_bytes( | |
| fid, | |
| num_bytes=24 * num_points2D, | |
| format_char_sequence="ddq" * num_points2D, | |
| ) | |
| xys = np.column_stack( | |
| [ | |
| tuple(map(float, x_y_id_s[0::3])), | |
| tuple(map(float, x_y_id_s[1::3])), | |
| ] | |
| ) | |
| point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) | |
| images[image_id] = ColmapImage( | |
| id=image_id, | |
| qvec=qvec, | |
| tvec=tvec, | |
| camera_id=camera_id, | |
| name=image_name, | |
| xys=xys, | |
| point3D_ids=point3D_ids, | |
| ) | |
| return images | |
| def read_extrinsics_text(path): | |
| images = {} | |
| with open(path, "r") as fid: | |
| while True: | |
| line = fid.readline() | |
| if not line: | |
| break | |
| line = line.strip() | |
| if len(line) > 0 and line[0] != "#": | |
| elems = line.split() | |
| image_id = int(elems[0]) | |
| qvec = np.array(tuple(map(float, elems[1:5]))) | |
| tvec = np.array(tuple(map(float, elems[5:8]))) | |
| camera_id = int(elems[8]) | |
| image_name = elems[9] | |
| elems = fid.readline().split() | |
| xys = np.column_stack( | |
| [ | |
| tuple(map(float, elems[0::3])), | |
| tuple(map(float, elems[1::3])), | |
| ] | |
| ) | |
| point3D_ids = np.array(tuple(map(int, elems[2::3]))) | |
| images[image_id] = ColmapImage( | |
| id=image_id, | |
| qvec=qvec, | |
| tvec=tvec, | |
| camera_id=camera_id, | |
| name=image_name, | |
| xys=xys, | |
| point3D_ids=point3D_ids, | |
| ) | |
| return images | |
| def read_intrinsics_binary(path_to_model_file): | |
| cameras = {} | |
| with open(path_to_model_file, "rb") as fid: | |
| num_cameras = read_next_bytes(fid, 8, "Q")[0] | |
| for _ in range(num_cameras): | |
| camera_properties = read_next_bytes( | |
| fid, num_bytes=24, format_char_sequence="iiQQ" | |
| ) | |
| camera_id = camera_properties[0] | |
| model_id = camera_properties[1] | |
| model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name | |
| width = camera_properties[2] | |
| height = camera_properties[3] | |
| num_params = CAMERA_MODEL_IDS[model_id].num_params | |
| params = read_next_bytes( | |
| fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params | |
| ) | |
| cameras[camera_id] = ColmapCamera( | |
| id=camera_id, | |
| model=model_name, | |
| width=width, | |
| height=height, | |
| params=np.array(params), | |
| ) | |
| assert len(cameras) == num_cameras | |
| return cameras | |
| def read_intrinsics_text(path): | |
| cameras = {} | |
| with open(path, "r") as fid: | |
| while True: | |
| line = fid.readline() | |
| if not line: | |
| break | |
| line = line.strip() | |
| if len(line) > 0 and line[0] != "#": | |
| elems = line.split() | |
| camera_id = int(elems[0]) | |
| model = elems[1] | |
| width = int(elems[2]) | |
| height = int(elems[3]) | |
| params = np.array(tuple(map(float, elems[4:]))) | |
| cameras[camera_id] = ColmapCamera( | |
| id=camera_id, | |
| model=model, | |
| width=width, | |
| height=height, | |
| params=params, | |
| ) | |
| return cameras | |
| # ============================================================================= | |
| # Scene Loading | |
| # ============================================================================= | |
| def load_colmap_scene(scene_path, sparse_dir="sparse/0", images_dir="images"): | |
| """Load COLMAP reconstruction and convert to ReSplat format. | |
| Returns dict with: | |
| image_paths: list[str] - full paths to images | |
| image_names: list[str] - image filenames | |
| c2w: np.ndarray [N, 4, 4] - camera-to-world matrices (OpenCV convention) | |
| intrinsics: np.ndarray [N, 3, 3] - normalized intrinsic matrices | |
| image_sizes: list[tuple[int, int]] - (width, height) per image | |
| """ | |
| sparse_path = os.path.join(scene_path, sparse_dir) | |
| # Try binary format first, fall back to text | |
| cameras_bin = os.path.join(sparse_path, "cameras.bin") | |
| cameras_txt = os.path.join(sparse_path, "cameras.txt") | |
| images_bin = os.path.join(sparse_path, "images.bin") | |
| images_txt = os.path.join(sparse_path, "images.txt") | |
| if os.path.exists(cameras_bin) and os.path.exists(images_bin): | |
| cam_intrinsics = read_intrinsics_binary(cameras_bin) | |
| cam_extrinsics = read_extrinsics_binary(images_bin) | |
| print(f"Loaded COLMAP binary format from {sparse_path}") | |
| elif os.path.exists(cameras_txt) and os.path.exists(images_txt): | |
| cam_intrinsics = read_intrinsics_text(cameras_txt) | |
| cam_extrinsics = read_extrinsics_text(images_txt) | |
| print(f"Loaded COLMAP text format from {sparse_path}") | |
| else: | |
| raise FileNotFoundError( | |
| f"No COLMAP reconstruction found in {sparse_path}. " | |
| "Expected cameras.bin/txt and images.bin/txt" | |
| ) | |
| images_root = os.path.join(scene_path, images_dir) | |
| # Sort images by name for deterministic ordering | |
| sorted_images = sorted(cam_extrinsics.values(), key=lambda x: x.name) | |
| image_paths = [] | |
| image_names = [] | |
| c2w_list = [] | |
| intrinsics_list = [] | |
| image_sizes = [] | |
| point3D_ids_list = [] | |
| for img in sorted_images: | |
| # Check image exists | |
| img_path = os.path.join(images_root, img.name) | |
| if not os.path.exists(img_path): | |
| print(f"Warning: Image not found, skipping: {img_path}") | |
| continue | |
| # Get camera intrinsics | |
| cam = cam_intrinsics[img.camera_id] | |
| # Only support PINHOLE and SIMPLE_PINHOLE (undistorted images) | |
| if cam.model == "PINHOLE": | |
| fx, fy, cx, cy = cam.params[:4] | |
| elif cam.model == "SIMPLE_PINHOLE": | |
| f, cx, cy = cam.params[:3] | |
| fx = fy = f | |
| else: | |
| print( | |
| f"Warning: Unsupported camera model '{cam.model}' for image " | |
| f"{img.name}, skipping. Only PINHOLE and SIMPLE_PINHOLE are " | |
| f"supported. Please undistort images first using COLMAP." | |
| ) | |
| continue | |
| # Build normalized intrinsic matrix (fx/width, fy/height, cx/width, cy/height) | |
| K = np.eye(3, dtype=np.float32) | |
| K[0, 0] = fx / cam.width | |
| K[1, 1] = fy / cam.height | |
| K[0, 2] = cx / cam.width | |
| K[1, 2] = cy / cam.height | |
| # Build W2C matrix from COLMAP qvec + tvec | |
| # COLMAP: qvec is world-to-camera rotation, tvec is translation in camera frame | |
| R = qvec2rotmat(img.qvec) | |
| w2c = np.eye(4, dtype=np.float32) | |
| w2c[:3, :3] = R | |
| w2c[:3, 3] = img.tvec | |
| # Invert to get C2W (camera-to-world) | |
| # COLMAP is already in OpenCV convention (Y-down, Z-forward) | |
| c2w = np.linalg.inv(w2c).astype(np.float32) | |
| # Extract valid point3D_ids (filter out -1 which means unmatched) | |
| valid_pts = img.point3D_ids[img.point3D_ids >= 0] | |
| image_paths.append(img_path) | |
| image_names.append(img.name) | |
| c2w_list.append(c2w) | |
| intrinsics_list.append(K) | |
| image_sizes.append((cam.width, cam.height)) | |
| point3D_ids_list.append(set(valid_pts.tolist())) | |
| if len(image_paths) == 0: | |
| raise RuntimeError(f"No valid images found in {images_root}") | |
| return { | |
| "image_paths": image_paths, | |
| "image_names": image_names, | |
| "c2w": np.stack(c2w_list, axis=0), # [N, 4, 4] | |
| "intrinsics": np.stack(intrinsics_list, axis=0), # [N, 3, 3] | |
| "image_sizes": image_sizes, | |
| "point3D_ids": point3D_ids_list, # list of sets, one per image | |
| } | |
| # ============================================================================= | |
| # View Selection | |
| # ============================================================================= | |
| def farthest_point_sample(xyz, npoint): | |
| """Farthest point sampling on camera positions. | |
| Adapted from src/dataset/view_sampler/view_sampler_bounded_v2.py | |
| Args: | |
| xyz: [B, N, 3] point cloud data | |
| npoint: number of samples | |
| Returns: | |
| centroids: [B, npoint] sampled indices | |
| """ | |
| device = xyz.device | |
| B, N, C = xyz.shape | |
| centroids = torch.zeros(B, npoint, dtype=torch.long).to(device) | |
| distance = torch.ones(B, N).to(device) * 1e10 | |
| batch_indices = torch.arange(B, dtype=torch.long).to(device) | |
| barycenter = torch.sum(xyz, 1) | |
| barycenter = barycenter / xyz.shape[1] | |
| barycenter = barycenter.view(B, 1, 3) | |
| dist = torch.sum((xyz - barycenter) ** 2, -1) | |
| farthest = torch.max(dist, 1)[1] | |
| for i in range(npoint): | |
| centroids[:, i] = farthest | |
| centroid = xyz[batch_indices, farthest, :].view(B, 1, 3) | |
| dist = torch.sum((xyz - centroid) ** 2, -1) | |
| mask = dist < distance | |
| distance[mask] = dist[mask] | |
| farthest = torch.max(distance, -1)[1] | |
| return centroids | |
| def select_context_views(c2w, num_context, strategy="fps"): | |
| """Select context view indices from all available views. | |
| Args: | |
| c2w: [N, 4, 4] camera-to-world matrices | |
| num_context: number of context views to select | |
| strategy: "fps" | "uniform" | |
| Returns: | |
| context_indices: sorted numpy array of indices | |
| """ | |
| N = len(c2w) | |
| if num_context >= N: | |
| return np.arange(N) | |
| if strategy == "fps": | |
| positions = torch.tensor(c2w[:, :3, 3], dtype=torch.float32).unsqueeze(0) | |
| indices = farthest_point_sample(positions, num_context)[0].numpy() | |
| return np.sort(indices) | |
| elif strategy == "uniform": | |
| return np.linspace(0, N - 1, num_context, dtype=int) | |
| else: | |
| raise ValueError(f"Unknown context selection strategy: {strategy}") | |
| def select_target_views(num_images, context_indices, target_selection="remaining", | |
| num_target=None): | |
| """Select target views for rendering. | |
| Args: | |
| num_images: total number of images | |
| context_indices: indices used as context | |
| target_selection: "remaining" | "all" | |
| num_target: optional max number of target views | |
| Returns: | |
| target_indices: numpy array of indices | |
| """ | |
| if target_selection == "remaining": | |
| all_indices = np.arange(num_images) | |
| target_indices = np.setdiff1d(all_indices, context_indices) | |
| elif target_selection == "all": | |
| target_indices = np.arange(num_images) | |
| else: | |
| raise ValueError(f"Unknown target selection: {target_selection}") | |
| if num_target is not None and len(target_indices) > num_target: | |
| step = len(target_indices) / num_target | |
| selected = [int(i * step) for i in range(num_target)] | |
| target_indices = target_indices[selected] | |
| return target_indices | |