"""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