hugging-apps/resplat-recurrent-gs
1
1"""2Self-contained inference script for ReSplat on COLMAP-processed datasets.3 4Usage with --data_dir + --scene_name (for datasets with scene subdirectories):5 python scripts/infer_colmap.py \6 --model_preset dl3dv_8v_512x960 \7 --data_dir path/to/colmap_data \8 --scene_name SCENE_NAME \9 --output_dir path/to/output \10 --save_images --save_video --save_ply11 12Usage with --scene_path (single COLMAP scene directory):13 # 8-view high-res model14 python scripts/infer_colmap.py --model_preset dl3dv_8v_512x960 \15 --scene_path path/to/colmap_scene \16 --output_dir path/to/output --save_images --save_ply17 18 # 16-view high-res model19 python scripts/infer_colmap.py --model_preset dl3dv_16v_540x960 \20 --scene_path path/to/colmap_scene \21 --output_dir path/to/output --save_images --save_ply22 23Usage without presets (manual config):24 python scripts/infer_colmap.py \25 --scene_path path/to/colmap_scene \26 --checkpoint pretrained/resplat-base-dl3dv-512x960-view8-8179ed87.pth \27 --experiment dl3dv \28 --num_context 8 --num_refine 4 --max_resolution 960 \29 --output_dir path/to/output \30 --save_images --save_ply31 32Available presets:33 dl3dv_8v_512x960 - 8-view base model, high-res (recommended)34 dl3dv_16v_540x960 - 16-view base model, high-res35 dl3dv_8v_256x448 - 8-view base model, low-res36 dl3dv_16v_256x448 - 16-view base model, low-res37 dl3dv_32v_256x448 - 32-view base model, low-res38 dl3dv_8v_256x448_small - 8-view small (ViT-S) model39 dl3dv_8v_256x448_large - 8-view large (ViT-L) model, init-only40"""41 42import argparse43import collections44import json45import os46import struct47import warnings48import sys49from pathlib import Path50 51sys.path.insert(0, str(Path(__file__).resolve().parent.parent))52 53import numpy as np54import torch55import torchvision.transforms as tf56from PIL import Image57from src.misc.stablize_camera import render_stabilization_path58 59warnings.filterwarnings("ignore")60torch.set_float32_matmul_precision("high")61 62 63# =============================================================================64# Model Presets65# =============================================================================66 67MODEL_PRESETS = {68 # === High-resolution models (Phase 2+) ===69 "dl3dv_8v_512x960": {70 "overrides": [],71 "num_context": 8,72 "num_refine": 4,73 "max_resolution": 960,74 "checkpoint": "pretrained/resplat-base-dl3dv-512x960-view8-8179ed87.pth",75 },76 "dl3dv_16v_540x960": {77 "overrides": [78 "model.encoder.gaussian_adapter.gaussian_scale_max=3.",79 "model.encoder.depth_pred_half_res=true",80 ],81 "num_context": 16,82 "num_refine": 2,83 "max_resolution": 960,84 "checkpoint": "pretrained/resplat-base-dl3dv-540x960-view16-a72dc6d0.pth",85 },86 # === Low-resolution models (Phase 1) ===87 "dl3dv_8v_256x448": {88 "overrides": [],89 "num_context": 8,90 "num_refine": 4,91 "max_resolution": 448,92 "checkpoint": "pretrained/resplat-base-dl3dv-256x448-view8-1934a04c.pth",93 },94 "dl3dv_16v_256x448": {95 "overrides": [],96 "num_context": 16,97 "num_refine": 4,98 "max_resolution": 448,99 "checkpoint": "pretrained/resplat-base-dl3dv-256x448-view16-f38bf984.pth",100 },101 "dl3dv_32v_256x448": {102 "overrides": [],103 "num_context": 32,104 "num_refine": 4,105 "max_resolution": 448,106 "checkpoint": "pretrained/resplat-base-dl3dv-256x448-view32-439b63a6.pth",107 },108 # === Small/Large backbone variants ===109 "dl3dv_8v_256x448_small": {110 "overrides": [111 "model.encoder.monodepth_vit_type=vits",112 "model.encoder.gaussian_regressor_channels=256",113 ],114 "num_context": 8,115 "num_refine": 4,116 "max_resolution": 448,117 "checkpoint": "pretrained/resplat-small-dl3dv-256x448-view8-548993fe.pth",118 },119 "dl3dv_8v_256x448_large": {120 "overrides": [121 "model.encoder.monodepth_vit_type=vitl",122 "model.encoder.gaussian_regressor_channels=768",123 ],124 "num_context": 8,125 "num_refine": 0,126 "max_resolution": 448,127 "checkpoint": "pretrained/resplat-large-dl3dv-256x448-view8-62f1703a.pth",128 },129}130 131 132# =============================================================================133# COLMAP Loading Functions134# =============================================================================135 136CameraModel = collections.namedtuple(137 "CameraModel", ["model_id", "model_name", "num_params"]138)139ColmapCamera = collections.namedtuple(140 "ColmapCamera", ["id", "model", "width", "height", "params"]141)142BaseImage = collections.namedtuple(143 "ColmapImage",144 ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"],145)146 147CAMERA_MODELS = {148 CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),149 CameraModel(model_id=1, model_name="PINHOLE", num_params=4),150 CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),151 CameraModel(model_id=3, model_name="RADIAL", num_params=5),152 CameraModel(model_id=4, model_name="OPENCV", num_params=8),153 CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),154 CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),155 CameraModel(model_id=7, model_name="FOV", num_params=5),156 CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),157 CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),158 CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),159}160CAMERA_MODEL_IDS = dict(161 [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]162)163 164 165class ColmapImage(BaseImage):166 def qvec2rotmat(self):167 return qvec2rotmat(self.qvec)168 169 170def qvec2rotmat(qvec):171 return np.array(172 [173 [174 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,175 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],176 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],177 ],178 [179 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],180 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,181 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],182 ],183 [184 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],185 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],186 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,187 ],188 ]189 )190 191 192def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):193 data = fid.read(num_bytes)194 return struct.unpack(endian_character + format_char_sequence, data)195 196 197def read_extrinsics_binary(path_to_model_file):198 images = {}199 with open(path_to_model_file, "rb") as fid:200 num_reg_images = read_next_bytes(fid, 8, "Q")[0]201 for _ in range(num_reg_images):202 binary_image_properties = read_next_bytes(203 fid, num_bytes=64, format_char_sequence="idddddddi"204 )205 image_id = binary_image_properties[0]206 qvec = np.array(binary_image_properties[1:5])207 tvec = np.array(binary_image_properties[5:8])208 camera_id = binary_image_properties[8]209 image_name = ""210 current_char = read_next_bytes(fid, 1, "c")[0]211 while current_char != b"\x00":212 image_name += current_char.decode("utf-8")213 current_char = read_next_bytes(fid, 1, "c")[0]214 num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[215 0216 ]217 x_y_id_s = read_next_bytes(218 fid,219 num_bytes=24 * num_points2D,220 format_char_sequence="ddq" * num_points2D,221 )222 xys = np.column_stack(223 [224 tuple(map(float, x_y_id_s[0::3])),225 tuple(map(float, x_y_id_s[1::3])),226 ]227 )228 point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))229 images[image_id] = ColmapImage(230 id=image_id,231 qvec=qvec,232 tvec=tvec,233 camera_id=camera_id,234 name=image_name,235 xys=xys,236 point3D_ids=point3D_ids,237 )238 return images239 240 241def read_extrinsics_text(path):242 images = {}243 with open(path, "r") as fid:244 while True:245 line = fid.readline()246 if not line:247 break248 line = line.strip()249 if len(line) > 0 and line[0] != "#":250 elems = line.split()251 image_id = int(elems[0])252 qvec = np.array(tuple(map(float, elems[1:5])))253 tvec = np.array(tuple(map(float, elems[5:8])))254 camera_id = int(elems[8])255 image_name = elems[9]256 elems = fid.readline().split()257 xys = np.column_stack(258 [259 tuple(map(float, elems[0::3])),260 tuple(map(float, elems[1::3])),261 ]262 )263 point3D_ids = np.array(tuple(map(int, elems[2::3])))264 images[image_id] = ColmapImage(265 id=image_id,266 qvec=qvec,267 tvec=tvec,268 camera_id=camera_id,269 name=image_name,270 xys=xys,271 point3D_ids=point3D_ids,272 )273 return images274 275 276def read_intrinsics_binary(path_to_model_file):277 cameras = {}278 with open(path_to_model_file, "rb") as fid:279 num_cameras = read_next_bytes(fid, 8, "Q")[0]280 for _ in range(num_cameras):281 camera_properties = read_next_bytes(282 fid, num_bytes=24, format_char_sequence="iiQQ"283 )284 camera_id = camera_properties[0]285 model_id = camera_properties[1]286 model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name287 width = camera_properties[2]288 height = camera_properties[3]289 num_params = CAMERA_MODEL_IDS[model_id].num_params290 params = read_next_bytes(291 fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params292 )293 cameras[camera_id] = ColmapCamera(294 id=camera_id,295 model=model_name,296 width=width,297 height=height,298 params=np.array(params),299 )300 assert len(cameras) == num_cameras301 return cameras302 303 304def read_intrinsics_text(path):305 cameras = {}306 with open(path, "r") as fid:307 while True:308 line = fid.readline()309 if not line:310 break311 line = line.strip()312 if len(line) > 0 and line[0] != "#":313 elems = line.split()314 camera_id = int(elems[0])315 model = elems[1]316 width = int(elems[2])317 height = int(elems[3])318 params = np.array(tuple(map(float, elems[4:])))319 cameras[camera_id] = ColmapCamera(320 id=camera_id,321 model=model,322 width=width,323 height=height,324 params=params,325 )326 return cameras327 328 329# =============================================================================330# Scene Loading331# =============================================================================332 333 334def load_colmap_scene(scene_path, sparse_dir="sparse/0", images_dir="images"):335 """Load COLMAP reconstruction and convert to ReSplat format.336 337 Returns dict with:338 image_paths: list[str] - full paths to images339 image_names: list[str] - image filenames340 c2w: np.ndarray [N, 4, 4] - camera-to-world matrices (OpenCV convention)341 intrinsics: np.ndarray [N, 3, 3] - normalized intrinsic matrices342 image_sizes: list[tuple[int, int]] - (width, height) per image343 """344 sparse_path = os.path.join(scene_path, sparse_dir)345 346 # Try binary format first, fall back to text347 cameras_bin = os.path.join(sparse_path, "cameras.bin")348 cameras_txt = os.path.join(sparse_path, "cameras.txt")349 images_bin = os.path.join(sparse_path, "images.bin")350 images_txt = os.path.join(sparse_path, "images.txt")351 352 if os.path.exists(cameras_bin) and os.path.exists(images_bin):353 cam_intrinsics = read_intrinsics_binary(cameras_bin)354 cam_extrinsics = read_extrinsics_binary(images_bin)355 print(f"Loaded COLMAP binary format from {sparse_path}")356 elif os.path.exists(cameras_txt) and os.path.exists(images_txt):357 cam_intrinsics = read_intrinsics_text(cameras_txt)358 cam_extrinsics = read_extrinsics_text(images_txt)359 print(f"Loaded COLMAP text format from {sparse_path}")360 else:361 raise FileNotFoundError(362 f"No COLMAP reconstruction found in {sparse_path}. "363 "Expected cameras.bin/txt and images.bin/txt"364 )365 366 images_root = os.path.join(scene_path, images_dir)367 368 # Sort images by name for deterministic ordering369 sorted_images = sorted(cam_extrinsics.values(), key=lambda x: x.name)370 371 image_paths = []372 image_names = []373 c2w_list = []374 intrinsics_list = []375 image_sizes = []376 point3D_ids_list = []377 378 for img in sorted_images:379 # Check image exists380 img_path = os.path.join(images_root, img.name)381 if not os.path.exists(img_path):382 print(f"Warning: Image not found, skipping: {img_path}")383 continue384 385 # Get camera intrinsics386 cam = cam_intrinsics[img.camera_id]387 388 # Only support PINHOLE and SIMPLE_PINHOLE (undistorted images)389 if cam.model == "PINHOLE":390 fx, fy, cx, cy = cam.params[:4]391 elif cam.model == "SIMPLE_PINHOLE":392 f, cx, cy = cam.params[:3]393 fx = fy = f394 else:395 print(396 f"Warning: Unsupported camera model '{cam.model}' for image "397 f"{img.name}, skipping. Only PINHOLE and SIMPLE_PINHOLE are "398 f"supported. Please undistort images first using COLMAP."399 )400 continue401 402 # Build normalized intrinsic matrix (fx/width, fy/height, cx/width, cy/height)403 K = np.eye(3, dtype=np.float32)404 K[0, 0] = fx / cam.width405 K[1, 1] = fy / cam.height406 K[0, 2] = cx / cam.width407 K[1, 2] = cy / cam.height408 409 # Build W2C matrix from COLMAP qvec + tvec410 # COLMAP: qvec is world-to-camera rotation, tvec is translation in camera frame411 R = qvec2rotmat(img.qvec)412 w2c = np.eye(4, dtype=np.float32)413 w2c[:3, :3] = R414 w2c[:3, 3] = img.tvec415 416 # Invert to get C2W (camera-to-world)417 # COLMAP is already in OpenCV convention (Y-down, Z-forward)418 c2w = np.linalg.inv(w2c).astype(np.float32)419 420 # Extract valid point3D_ids (filter out -1 which means unmatched)421 valid_pts = img.point3D_ids[img.point3D_ids >= 0]422 423 image_paths.append(img_path)424 image_names.append(img.name)425 c2w_list.append(c2w)426 intrinsics_list.append(K)427 image_sizes.append((cam.width, cam.height))428 point3D_ids_list.append(set(valid_pts.tolist()))429 430 if len(image_paths) == 0:431 raise RuntimeError(f"No valid images found in {images_root}")432 433 return {434 "image_paths": image_paths,435 "image_names": image_names,436 "c2w": np.stack(c2w_list, axis=0), # [N, 4, 4]437 "intrinsics": np.stack(intrinsics_list, axis=0), # [N, 3, 3]438 "image_sizes": image_sizes,439 "point3D_ids": point3D_ids_list, # list of sets, one per image440 }441 442 443# =============================================================================444# Frame Subsetting445# =============================================================================446 447 448def subset_scene_data(scene_data, start_frame, frame_distance):449 """Subset scene_data to frames [start_frame, start_frame + frame_distance).450 451 Returns a new scene_data dict with only the selected frames.452 Indices are 0-based in the returned data (transparent to downstream code).453 """454 end_frame = min(start_frame + frame_distance, len(scene_data["image_paths"]))455 indices = list(range(start_frame, end_frame))456 457 return {458 "image_paths": [scene_data["image_paths"][i] for i in indices],459 "image_names": [scene_data["image_names"][i] for i in indices],460 "c2w": scene_data["c2w"][indices],461 "intrinsics": scene_data["intrinsics"][indices],462 "image_sizes": [scene_data["image_sizes"][i] for i in indices],463 "point3D_ids": [scene_data["point3D_ids"][i] for i in indices],464 }465 466 467# =============================================================================468# View Selection469# =============================================================================470 471 472def farthest_point_sample(xyz, npoint):473 """Farthest point sampling on camera positions.474 475 Adapted from src/dataset/view_sampler/view_sampler_bounded_v2.py476 477 Args:478 xyz: [B, N, 3] point cloud data479 npoint: number of samples480 Returns:481 centroids: [B, npoint] sampled indices482 """483 device = xyz.device484 B, N, C = xyz.shape485 486 centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)487 distance = torch.ones(B, N).to(device) * 1e10488 489 batch_indices = torch.arange(B, dtype=torch.long).to(device)490 491 barycenter = torch.sum(xyz, 1)492 barycenter = barycenter / xyz.shape[1]493 barycenter = barycenter.view(B, 1, 3)494 495 dist = torch.sum((xyz - barycenter) ** 2, -1)496 farthest = torch.max(dist, 1)[1]497 498 for i in range(npoint):499 centroids[:, i] = farthest500 centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)501 dist = torch.sum((xyz - centroid) ** 2, -1)502 mask = dist < distance503 distance[mask] = dist[mask]504 farthest = torch.max(distance, -1)[1]505 506 return centroids507 508 509def select_context_views(c2w, num_context, strategy="fps"):510 """Select context view indices from all available views.511 512 Args:513 c2w: [N, 4, 4] camera-to-world matrices514 num_context: number of context views to select515 strategy: "fps" | "uniform"516 Returns:517 context_indices: sorted numpy array of indices518 """519 N = len(c2w)520 if num_context >= N:521 return np.arange(N)522 523 if strategy == "fps":524 positions = torch.tensor(c2w[:, :3, 3], dtype=torch.float32).unsqueeze(0)525 indices = farthest_point_sample(positions, num_context)[0].numpy()526 return np.sort(indices)527 elif strategy == "uniform":528 return np.linspace(0, N - 1, num_context, dtype=int)529 else:530 raise ValueError(f"Unknown context selection strategy: {strategy}")531 532 533def select_target_views(num_images, context_indices, target_selection="remaining",534 num_target=None):535 """Select target views for rendering.536 537 Args:538 num_images: total number of images539 context_indices: indices used as context540 target_selection: "remaining" | "all"541 num_target: optional max number of target views542 Returns:543 target_indices: numpy array of indices544 """545 if target_selection == "remaining":546 all_indices = np.arange(num_images)547 target_indices = np.setdiff1d(all_indices, context_indices)548 elif target_selection == "all":549 target_indices = np.arange(num_images)550 else:551 raise ValueError(f"Unknown target selection: {target_selection}")552 553 if num_target is not None and len(target_indices) > num_target:554 step = len(target_indices) / num_target555 selected = [int(i * step) for i in range(num_target)]556 target_indices = target_indices[selected]557 558 return target_indices559 560 561# =============================================================================562# Image Loading and Preprocessing563# =============================================================================564 565 566def compute_target_shape(orig_h, orig_w, max_resolution=960, image_shape=None):567 """Compute target image shape, ensuring divisibility by 64.568 569 The encoder requires H and W to be divisible by570 shim_patch_size * downscale_factor = 16 * 4 = 64.571 """572 DIVISOR = 64573 574 if image_shape is not None:575 h, w = image_shape576 else:577 scale = max_resolution / max(orig_h, orig_w)578 if scale < 1.0:579 h = int(orig_h * scale)580 w = int(orig_w * scale)581 else:582 h, w = orig_h, orig_w583 584 # Round down to nearest multiple of DIVISOR585 h = (h // DIVISOR) * DIVISOR586 w = (w // DIVISOR) * DIVISOR587 588 assert h > 0 and w > 0, (589 f"Resolution too small after rounding to multiple of {DIVISOR}: "590 f"{h}x{w}. Increase --max_resolution."591 )592 return h, w593 594 595def load_and_preprocess_images(image_paths, target_h, target_w):596 """Load images and resize to target resolution.597 598 Args:599 image_paths: list of image file paths600 target_h, target_w: target dimensions (must be divisible by 64)601 Returns:602 images: [V, 3, H, W] float32 tensor in [0, 1]603 """604 to_tensor = tf.ToTensor()605 images = []606 for path in image_paths:607 img = Image.open(path).convert("RGB")608 img = img.resize((target_w, target_h), Image.LANCZOS)609 images.append(to_tensor(img))610 return torch.stack(images)611 612 613# =============================================================================614# Pose Utilities615# =============================================================================616 617 618def camera_normalization(pivotal_pose, poses):619 """Align all poses relative to a reference pose.620 621 Adapted from src/dataset/dataset_dl3dv.py:camera_normalization622 623 Args:624 pivotal_pose: [1, 4, 4] reference camera pose625 poses: [N, 4, 4] all camera poses626 Returns:627 normalized poses [N, 4, 4]628 """629 camera_norm_matrix = torch.inverse(pivotal_pose)630 poses = torch.bmm(camera_norm_matrix.repeat(poses.shape[0], 1, 1), poses)631 return poses632 633 634# =============================================================================635# Model Construction636# =============================================================================637 638 639def build_model(experiment, checkpoint, num_refine, image_shape, overrides,640 device, no_strict_load=True):641 """Build model using Hydra compose API and load checkpoint.642 643 Returns: (encoder, decoder, data_shim) all on device644 """645 from hydra import compose, initialize_config_dir646 from hydra.core.global_hydra import GlobalHydra647 648 # Clear any existing Hydra state649 GlobalHydra.instance().clear()650 651 config_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config")652 653 with initialize_config_dir(config_dir=config_dir, version_base=None):654 hydra_overrides = [655 f"+experiment={experiment}",656 "mode=test",657 f"model.encoder.num_refine={num_refine}",658 f"dataset.image_shape=[{image_shape[0]},{image_shape[1]}]",659 f"dataset.ori_image_shape=[{image_shape[0]},{image_shape[1]}]",660 f"output_dir=outputs/colmap_inference",661 ]662 hydra_overrides.extend(overrides)663 cfg_dict = compose(config_name="main", overrides=hydra_overrides)664 665 # Import src modules (outside beartype hook is fine for inference)666 from src.config import load_typed_root_config667 from src.dataset.data_module import get_data_shim668 from src.global_cfg import set_cfg669 from src.model.decoder import get_decoder670 from src.model.encoder import get_encoder671 from src.model.model_wrapper import ModelWrapper672 673 set_cfg(cfg_dict)674 cfg = load_typed_root_config(cfg_dict)675 676 # Build encoder and decoder677 encoder, _ = get_encoder(cfg.model.encoder)678 decoder = get_decoder(cfg.model.decoder, cfg.dataset)679 680 # Build ModelWrapper to load checkpoint with correct key prefixes681 model_wrapper = ModelWrapper(682 cfg.optimizer,683 cfg.test,684 cfg.train,685 encoder,686 None, # encoder_visualizer687 decoder,688 [], # losses (not needed for inference)689 None, # step_tracker690 )691 692 # Load checkpoint693 print(f"Loading checkpoint: {checkpoint}")694 ckpt = torch.load(checkpoint, map_location="cpu")695 if "state_dict" in ckpt:696 ckpt = ckpt["state_dict"]697 model_wrapper.load_state_dict(ckpt, strict=not no_strict_load)698 model_wrapper = model_wrapper.to(device).eval()699 700 data_shim = get_data_shim(model_wrapper.encoder)701 702 return model_wrapper.encoder, model_wrapper.decoder, data_shim703 704 705# =============================================================================706# Batch Construction707# =============================================================================708 709 710def move_to_device(data, device):711 """Recursively move tensors in a nested dict to device."""712 if isinstance(data, torch.Tensor):713 return data.to(device)714 elif isinstance(data, dict):715 return {k: move_to_device(v, device) for k, v in data.items()}716 elif isinstance(data, list):717 return [move_to_device(v, device) for v in data]718 return data719 720 721def build_batch(722 context_images,723 target_images,724 context_c2w,725 target_c2w,726 context_K,727 target_K,728 near,729 far,730 scene_name,731 device,732):733 """Build a batch dict matching ReSplat's BatchedExample format.734 735 All tensors get batch dimension B=1. Poses are aligned to the middle736 context view.737 """738 Vc = len(context_c2w)739 Vt = len(target_c2w)740 741 # Align all poses to middle context view742 all_c2w = torch.cat([context_c2w, target_c2w], dim=0)743 mid_idx = Vc // 2744 all_c2w = camera_normalization(context_c2w[mid_idx : mid_idx + 1], all_c2w)745 context_c2w_aligned = all_c2w[:Vc]746 target_c2w_aligned = all_c2w[Vc:]747 748 batch = {749 "context": {750 "image": context_images.unsqueeze(0), # [1, Vc, 3, H, W]751 "extrinsics": context_c2w_aligned.unsqueeze(0), # [1, Vc, 4, 4]752 "intrinsics": context_K.unsqueeze(0), # [1, Vc, 3, 3]753 "near": torch.full((1, Vc), near),754 "far": torch.full((1, Vc), far),755 "index": torch.arange(Vc).unsqueeze(0),756 },757 "target": {758 "image": target_images.unsqueeze(0), # [1, Vt, 3, H, W]759 "extrinsics": target_c2w_aligned.unsqueeze(0), # [1, Vt, 4, 4]760 "intrinsics": target_K.unsqueeze(0), # [1, Vt, 3, 3]761 "near": torch.full((1, Vt), near),762 "far": torch.full((1, Vt), far),763 "index": torch.arange(Vc, Vc + Vt).unsqueeze(0),764 },765 "scene": [scene_name],766 }767 768 return move_to_device(batch, device)769 770 771# =============================================================================772# Inference773# =============================================================================774 775 776@torch.no_grad()777def run_inference(encoder, decoder, batch, num_refine, render_chunk_size=10,778 save_depth=False):779 """Run ReSplat inference.780 781 Following the pattern from src/model/model_wrapper.py test_step.782 783 Returns:784 gaussians: Gaussians dataclass785 rendered: [Vt, 3, H, W] rendered images tensor786 visualization_dump: dict with depth data (if save_depth=True), else None787 """788 _, _, _, h, w = batch["target"]["image"].shape789 790 visualization_dump = {} if save_depth else None791 792 # 1. Initial forward pass (encoder)793 print("Running encoder forward pass...")794 gaussians_out = encoder(batch["context"], global_step=0, deterministic=False,795 visualization_dump=visualization_dump)796 797 if isinstance(gaussians_out, dict):798 condition_features = gaussians_out.get("condition_features", None)799 gaussians = gaussians_out["gaussians"]800 else:801 gaussians = gaussians_out802 condition_features = None803 804 # 2. Refinement (if enabled)805 if num_refine > 0 and condition_features is not None:806 print(f"Running refinement ({num_refine} iterations)...")807 refine_output = encoder.forward_update(808 batch["context"],809 batch["target"],810 condition_features,811 gaussians,812 decoder,813 None, # context_remain814 )815 gaussians = refine_output["gaussian"][-1]816 817 # 3. Render target views (chunked for memory)818 Vt = batch["target"]["extrinsics"].shape[1]819 print(f"Rendering {Vt} target views...")820 all_colors = []821 all_depths = []822 for i in range(0, Vt, render_chunk_size):823 end = min(i + render_chunk_size, Vt)824 output = decoder.forward(825 gaussians,826 batch["target"]["extrinsics"][:, i:end],827 batch["target"]["intrinsics"][:, i:end],828 batch["target"]["near"][:, i:end],829 batch["target"]["far"][:, i:end],830 (h, w),831 depth_mode=None,832 )833 all_colors.append(output.color[0]) # [chunk, 3, H, W]834 all_depths.append(output.depth[0]) # [chunk, H, W]835 836 rendered = torch.cat(all_colors, dim=0) # [Vt, 3, H, W]837 rendered_depth = torch.cat(all_depths, dim=0) # [Vt, H, W]838 print(f"Rendered {Vt} views at {h}x{w}")839 840 return gaussians, rendered, rendered_depth, visualization_dump841 842 843# =============================================================================844# Output Saving845# =============================================================================846 847 848def save_outputs(849 rendered, gaussians, batch, output_dir, image_names,850 save_images=True, save_ply=False, save_depth=False,851 context_image_names=None, context_images=None,852 visualization_dump=None, rendered_depth=None, max_save_images=10,853):854 """Save rendered images, depth maps, and Gaussians (PLY)."""855 os.makedirs(output_dir, exist_ok=True)856 857 # Save rendered images (at most max_save_images, evenly spaced)858 if save_images:859 images_dir = os.path.join(output_dir, "rendered")860 os.makedirs(images_dir, exist_ok=True)861 total = len(rendered)862 if max_save_images > 0 and total > max_save_images:863 save_indices = np.linspace(0, total - 1, max_save_images, dtype=int)864 else:865 save_indices = range(total)866 for i in save_indices:867 img_np = (rendered[i].clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(868 np.uint8869 )870 stem = Path(image_names[i]).stem871 Image.fromarray(img_np).save(os.path.join(images_dir, f"{stem}.png"))872 print(f"Saved {len(save_indices)} rendered images to {images_dir}")873 874 # Save context (input) images for reference875 if save_images and context_images is not None and context_image_names is not None:876 input_dir = os.path.join(output_dir, "input")877 os.makedirs(input_dir, exist_ok=True)878 for img_tensor, name in zip(context_images, context_image_names):879 img_np = (img_tensor.clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(880 np.uint8881 )882 stem = Path(name).stem883 Image.fromarray(img_np).save(os.path.join(input_dir, f"{stem}.png"))884 print(f"Saved {len(context_images)} input images to {input_dir}")885 886 # Save input view depth maps887 if save_depth and visualization_dump is not None:888 from src.visualization.vis_depth import viz_depth_tensor889 890 # Latent-resolution depth (used for Gaussian unprojection)891 if "depth" in visualization_dump:892 depth_dir = os.path.join(output_dir, "depth")893 os.makedirs(depth_dir, exist_ok=True)894 895 # [B, V, H, W, srf, s] -> [V, H, W]896 depth = visualization_dump["depth"][0, :, :, :, 0, 0].cpu().detach()897 898 for depth_i, name in zip(depth, context_image_names or []):899 stem = Path(name).stem900 depth_viz = viz_depth_tensor(1.0 / depth_i, return_numpy=True)901 Image.fromarray(depth_viz).save(os.path.join(depth_dir, f"{stem}.png"))902 903 print(f"Saved {len(depth)} latent-res depth maps to {depth_dir}")904 905 # Full-resolution depth (before latent downsampling)906 if "depth_fullres" in visualization_dump:907 depth_fullres_dir = os.path.join(output_dir, "depth_fullres")908 os.makedirs(depth_fullres_dir, exist_ok=True)909 910 # [B, V, H, W] -> [V, H, W]911 depth_fullres = visualization_dump["depth_fullres"][0].cpu().detach()912 913 for depth_i, name in zip(depth_fullres, context_image_names or []):914 stem = Path(name).stem915 depth_viz = viz_depth_tensor(1.0 / depth_i, return_numpy=True)916 Image.fromarray(depth_viz).save(917 os.path.join(depth_fullres_dir, f"{stem}.png")918 )919 920 print(f"Saved {len(depth_fullres)} full-res depth maps to {depth_fullres_dir}")921 922 # Save rendered (target view) depth maps923 if save_depth and rendered_depth is not None:924 from src.visualization.vis_depth import viz_depth_tensor925 926 rendered_depth_dir = os.path.join(output_dir, "rendered_depth")927 os.makedirs(rendered_depth_dir, exist_ok=True)928 total = len(rendered_depth)929 if max_save_images > 0 and total > max_save_images:930 save_indices = np.linspace(0, total - 1, max_save_images, dtype=int)931 else:932 save_indices = range(total)933 for i in save_indices:934 depth_i = rendered_depth[i].cpu().detach()935 stem = Path(image_names[i]).stem936 depth_viz = viz_depth_tensor(1.0 / depth_i, return_numpy=True)937 Image.fromarray(depth_viz).save(os.path.join(rendered_depth_dir, f"{stem}.png"))938 print(f"Saved {len(save_indices)} rendered depth maps to {rendered_depth_dir}")939 940 # Save PLY941 if save_ply:942 try:943 from src.model.ply_export import export_ply944 945 ply_path = Path(output_dir) / "gaussians.ply"946 Vc = batch["context"]["extrinsics"].shape[1]947 mid_idx = Vc // 2948 export_ply(949 batch["context"]["extrinsics"][0, mid_idx],950 gaussians.means[0],951 gaussians.scales[0],952 gaussians.rotations[0],953 gaussians.harmonics[0],954 gaussians.opacities[0],955 ply_path,956 align_to_view=True,957 )958 print(f"Saved Gaussian PLY to {ply_path}")959 except Exception as e:960 print(f"Warning: Failed to save PLY: {e}")961 962 963# =============================================================================964# Smooth Video Rendering965# =============================================================================966 967 968@torch.no_grad()969def render_smooth_video(970 gaussians, decoder, all_c2w_np, context_c2w, intrinsic_np,971 near, far, image_shape, output_dir,972 render_chunk_size=10, fps=30, smooth_kernel=45,973 device="cuda",974):975 """Render a smooth video by smoothing all scene poses and rendering each frame.976 977 Args:978 gaussians: Gaussians dataclass from inference979 decoder: the gsplat decoder980 all_c2w_np: numpy [N, 4, 4] all scene camera-to-world poses in order981 context_c2w: torch [Vc, 4, 4] unnormalized context poses (for normalization pivot)982 intrinsic_np: numpy [3, 3] single normalized intrinsic matrix983 near, far: float near/far plane distances984 image_shape: (H, W) target image resolution985 output_dir: directory to save the video986 render_chunk_size: number of views to render at once987 fps: video frame rate988 smooth_kernel: Gaussian kernel size for trajectory smoothing989 device: torch device990 """991 import imageio992 993 N = len(all_c2w_np)994 h, w = image_shape995 print(f"Rendering smooth video with {N} frames (kernel={smooth_kernel})...")996 997 # 1. Smooth the trajectory998 poses_3x4 = all_c2w_np[:, :3, :] # [N, 3, 4]999 smoothed_list = render_stabilization_path(poses_3x4, k_size=smooth_kernel)1000 1001 # Reconstruct [N, 4, 4] from smoothed [3, 4] matrices1002 smoothed_c2w = np.zeros((N, 4, 4), dtype=np.float32)1003 for i, pose in enumerate(smoothed_list):1004 smoothed_c2w[i, :3, :] = pose1005 smoothed_c2w[i, 3, 3] = 1.01006 1007 # 2. Normalize poses using the same pivot as build_batch()1008 smoothed_c2w_t = torch.tensor(smoothed_c2w, dtype=torch.float32)1009 mid_idx = len(context_c2w) // 21010 smoothed_c2w_aligned = camera_normalization(1011 context_c2w[mid_idx:mid_idx + 1], smoothed_c2w_t1012 )1013 1014 # 3. Prepare intrinsics (single intrinsic for all frames)1015 intrinsic_t = torch.tensor(intrinsic_np, dtype=torch.float32)1016 1017 # 4. Render in chunks1018 all_frames = []1019 for i in range(0, N, render_chunk_size):1020 end = min(i + render_chunk_size, N)1021 chunk_size = end - i1022 extrinsics = smoothed_c2w_aligned[i:end].unsqueeze(0).to(device) # [1, chunk, 4, 4]1023 intrinsics = intrinsic_t.unsqueeze(0).unsqueeze(0).expand(1, chunk_size, -1, -1).to(device) # [1, chunk, 3, 3]1024 near_t = torch.full((1, chunk_size), near, device=device)1025 far_t = torch.full((1, chunk_size), far, device=device)1026 1027 output = decoder.forward(1028 gaussians, extrinsics, intrinsics, near_t, far_t,1029 (h, w), depth_mode=None,1030 )1031 # output.color[0] is [chunk, 3, H, W]1032 for img_tensor in output.color[0]:1033 img_np = (img_tensor.clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)1034 all_frames.append(img_np)1035 1036 # 5. Write video1037 video_path = os.path.join(output_dir, "video.mp4")1038 imageio.mimwrite(video_path, all_frames, fps=fps, quality=8)1039 print(f"Saved smooth video ({N} frames) to {video_path}")1040 1041 1042# =============================================================================1043# Evaluation Metrics1044# =============================================================================1045 1046 1047@torch.no_grad()1048def compute_metrics(rendered, target_images, target_names, output_dir, device="cuda:0",1049 chunk_size=4):1050 """Compute PSNR, SSIM, LPIPS between rendered and ground truth target views.1051 1052 Args:1053 rendered: [Vt, 3, H, W] rendered images tensor (on GPU)1054 target_images: [Vt, 3, H, W] ground truth images tensor1055 target_names: list of image names for per-view results1056 output_dir: directory to save metrics.json1057 device: device for LPIPS computation1058 chunk_size: chunk size for LPIPS to avoid OOM1059 Returns:1060 dict with mean metrics1061 """1062 from src.evaluation.metrics import compute_psnr, compute_ssim, compute_lpips1063 1064 rendered = rendered.clamp(0, 1)1065 target_images = target_images.clamp(0, 1).to(rendered.device)1066 1067 Vt = rendered.shape[0]1068 1069 # Compute PSNR and SSIM (lightweight, can do all at once)1070 psnr_vals = compute_psnr(target_images, rendered) # [Vt]1071 ssim_vals = compute_ssim(target_images, rendered) # [Vt]1072 1073 # Compute LPIPS in chunks (VGG network uses memory)1074 lpips_vals = []1075 for i in range(0, Vt, chunk_size):1076 end = min(i + chunk_size, Vt)1077 lpips_chunk = compute_lpips(target_images[i:end], rendered[i:end])1078 lpips_vals.append(lpips_chunk)1079 lpips_vals = torch.cat(lpips_vals, dim=0) # [Vt]1080 1081 # Aggregate1082 mean_psnr = psnr_vals.mean().item()1083 mean_ssim = ssim_vals.mean().item()1084 mean_lpips = lpips_vals.mean().item()1085 1086 # Build per-view results1087 per_view = []1088 for i in range(Vt):1089 per_view.append({1090 "name": target_names[i],1091 "psnr": round(psnr_vals[i].item(), 4),1092 "ssim": round(ssim_vals[i].item(), 4),1093 "lpips": round(lpips_vals[i].item(), 4),1094 })1095 1096 results = {1097 "mean": {1098 "psnr": round(mean_psnr, 3),1099 "ssim": round(mean_ssim, 3),1100 "lpips": round(mean_lpips, 3),1101 },1102 "per_view": per_view,1103 }1104 1105 # Save to JSON1106 os.makedirs(output_dir, exist_ok=True)1107 metrics_path = os.path.join(output_dir, "metrics.json")1108 with open(metrics_path, "w") as f:1109 json.dump(results, f, indent=2)1110 1111 print(f"\n Evaluation Metrics ({Vt} target views):")1112 print(f" PSNR: {mean_psnr:.3f}")1113 print(f" SSIM: {mean_ssim:.3f}")1114 print(f" LPIPS: {mean_lpips:.3f}")1115 print(f" Saved to {metrics_path}")1116 1117 return results1118 1119 1120# =============================================================================1121# Main1122# =============================================================================1123 1124 1125def parse_args():1126 parser = argparse.ArgumentParser(1127 description="ReSplat inference on COLMAP-processed datasets"1128 )1129 1130 # Scene specification (either --scene_path OR --data_dir + --scene_name/--scene_list)1131 parser.add_argument(1132 "--scene_path",1133 type=str,1134 default=None,1135 help="Path to COLMAP scene directory (must contain sparse/0/ and images/)",1136 )1137 parser.add_argument(1138 "--data_dir",1139 type=str,1140 default=None,1141 help="Base directory containing scene subdirectories "1142 "(e.g., datasets/dl3dv-evaluation/images_tar)",1143 )1144 parser.add_argument(1145 "--scene_name",1146 type=str,1147 default=None,1148 help="Scene name (hash). Used with --data_dir to construct scene_path",1149 )1150 parser.add_argument(1151 "--scene_list",1152 type=str,1153 default=None,1154 help="Path to text file with one scene name per line, or 'all' for all "1155 "subdirs in --data_dir. Used with --data_dir for batch processing",1156 )1157 1158 # Frame range (for long video scenes)1159 parser.add_argument(1160 "--start_frame",1161 type=int,1162 default=0,1163 help="First frame index in the evaluation subset",1164 )1165 parser.add_argument(1166 "--frame_distance",1167 type=int,1168 default=60,1169 help="Number of frames from start_frame to include",1170 )1171 parser.add_argument(1172 "--checkpoint",1173 type=str,1174 default=None,1175 help="Path to pretrained model .pth file (required unless --model_preset is used)",1176 )1177 1178 # Model preset1179 parser.add_argument(1180 "--model_preset",1181 type=str,1182 default=None,1183 choices=list(MODEL_PRESETS.keys()),1184 help="Model preset that auto-sets checkpoint, overrides, num_context, "1185 "num_refine, and max_resolution. Explicit CLI args override preset values.",1186 )1187 1188 # Model config1189 parser.add_argument(1190 "--experiment",1191 type=str,1192 default="dl3dv",1193 help="Hydra experiment config name (dl3dv or re10k), must match checkpoint",1194 )1195 parser.add_argument(1196 "--num_refine",1197 type=int,1198 default=None,1199 help="Number of refinement iterations (0=init only). Default: from preset, or 4",1200 )