#!/usr/bin/env python3 """SplatAnything: single-scene feed-forward Gaussian inference.""" import argparse import json import time from pathlib import Path ROOT = Path(__file__).resolve().parent def parse_args(argv=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--config_path', default=str(ROOT / 'configs/inference.yaml')) inputs = parser.add_mutually_exclusive_group() inputs.add_argument('--example', help='Bundled scene, e.g. object/1. See --list_examples.') inputs.add_argument('--input_path', help='Directory of images for one scene.') parser.add_argument('--list_examples', action='store_true', help='List bundled scenes and exit; no weights needed.') parser.add_argument('--output_path', help='Output root; default outputs, or outputs/CATEGORY/MODE for examples.') parser.add_argument('--resolution', type=int, nargs=2, metavar=('W', 'H'), help='Positive multiples of 14; example preset or 1022 1022.') parser.add_argument('--max_images', type=int, help='First N sorted images; example preset or 64.') parser.add_argument('--cameras', '--prior_cam_path', dest='cameras', help='prior_camera.json: OpenCV 4x4 camera-to-world poses and ' '3x3 pixel intrinsics at the original image resolution.') mode = parser.add_mutually_exclusive_group() mode.add_argument('--mode', choices=['rgb', 'pose'], help='RGB-only or pose-conditioned inference (requires cameras).') mode.add_argument('--pose_condition', dest='pose_condition', action='store_true', help=argparse.SUPPRESS) mode.add_argument('--no_pose_condition', dest='pose_condition', action='store_false', help=argparse.SUPPRESS) parser.set_defaults(pose_condition=None) parser.add_argument('--sky_seg', action='store_true', help='Segment and remove sky before voxel merging.') parser.add_argument('--sky_onnx_path', default=str(ROOT / 'checkpoints/skyseg.onnx')) parser.add_argument('--sky_max_fraction', type=float, default=0.95) parser.add_argument('--ckpt_path', help='Complete Gaussian student checkpoint (.ckpt/.pth/.safetensors).') parser.add_argument('--geometry_ckpt_path', help='WorldMirror-2 geometry checkpoint.') parser.add_argument('--opacity_thresh', type=float, help='Override the exported Gaussian opacity threshold.') parser.add_argument('--no_cam_align', action='store_true', help='Keep supplied camera frame; default anchors to camera 0.') parser.add_argument('--no_depth_rgb', action='store_true', help='Skip per-view RGB and depth export.') parser.add_argument('--device', choices=['cuda', 'cpu'], default='cuda') args = parser.parse_args(argv) if args.list_examples: manifest = json.loads((ROOT / 'examples/manifest.json').read_text()) for name, preset in manifest['examples'].items(): modes = 'rgb, pose' if preset['cameras'] else 'rgb (pose requires --cameras)' width, height = preset['resolution'] print(f'{name:12} {width}x{height}, up to {preset["max_images"]:3} views; {modes}') parser.exit() if args.mode is not None: args.pose_condition = args.mode == 'pose' preset = {} if args.example: manifest = json.loads((ROOT / 'examples/manifest.json').read_text()) if args.example not in manifest['examples']: parser.error(f'Unknown example {args.example!r}. Use --list_examples.') preset = manifest['examples'][args.example] args.input_path = str(ROOT / 'examples' / preset['images']) if args.pose_condition is None: args.pose_condition = False args.mode = 'pose' if args.pose_condition else 'rgb' if args.pose_condition and not args.cameras: if not preset['cameras']: parser.error(f'{args.example} has no bundled calibration. Use --mode rgb or supply --cameras.') args.cameras = str(ROOT / 'examples' / preset['cameras']) if args.output_path is None: args.output_path = str(ROOT / 'outputs' / args.example.split('/')[0] / args.mode) if not args.input_path: parser.error('Provide --example NAME or --input_path IMAGES.') args.resolution = args.resolution if args.resolution is not None else preset.get('resolution', [1022, 1022]) args.max_images = args.max_images if args.max_images is not None else preset.get('max_images', 64) args.output_path = args.output_path or str(ROOT / 'outputs') if any(n <= 0 or n % 14 for n in args.resolution): parser.error('--resolution W H must be positive multiples of 14 (e.g. 518 518 or 1022 756).') if args.max_images <= 0: parser.error('--max_images must be positive.') if not 0 <= args.sky_max_fraction <= 1: parser.error('--sky_max_fraction must be in [0, 1].') if args.opacity_thresh is not None and not 0 <= args.opacity_thresh <= 1: parser.error('--opacity_thresh must be in [0, 1].') return args def load_config(args): from omegaconf import OmegaConf cfg = OmegaConf.load(args.config_path) if not OmegaConf.is_dict(cfg) or set(cfg) != {'wrapper'}: raise ValueError('The config must contain only the wrapper mapping.') fields = { 'wrapper': {'worldmirror2_pretrained', 'SplatAnything_pretrained', 'enable_cond', 'infer_opacity_thresh', 'model'}, 'wrapper.model': {'use_gt_camera', 'filter_flying_points_infer', 'filter_large_scale_infer', 'max_scale_abs', 'max_scale_pct', 'gs_voxel_size'}, } for path, expected in fields.items(): section = OmegaConf.select(cfg, path) if not OmegaConf.is_dict(section): raise ValueError(f'{path} must be a mapping.') missing = expected - set(section) extra = set(section) - expected if missing or extra: raise ValueError(f'Invalid config fields in {path}: ' f'missing={sorted(missing)}, unsupported={sorted(extra)}.') if args.pose_condition is not None: cfg.wrapper.enable_cond = args.pose_condition cfg.wrapper.model.use_gt_camera = args.pose_condition if bool(cfg.wrapper.enable_cond) != bool(cfg.wrapper.model.use_gt_camera): raise ValueError('enable_cond and use_gt_camera must both be true (pose-conditioned) ' 'or both false (RGB-only).') if cfg.wrapper.enable_cond and not args.cameras: raise ValueError('Pose conditioning requires --cameras with poses AND intrinsics for every input image.') if not cfg.wrapper.enable_cond and args.cameras: raise ValueError('--cameras requires pose conditioning. Add --mode pose or set both config flags true.') if args.cameras and not Path(args.cameras).is_file(): raise FileNotFoundError(f'Camera JSON not found: {args.cameras}') if args.ckpt_path: cfg.wrapper.SplatAnything_pretrained = args.ckpt_path if args.geometry_ckpt_path: cfg.wrapper.worldmirror2_pretrained = args.geometry_ckpt_path if args.opacity_thresh is not None: cfg.wrapper.infer_opacity_thresh = args.opacity_thresh if not 0 <= cfg.wrapper.infer_opacity_thresh <= 1: raise ValueError('infer_opacity_thresh must be in [0, 1].') if cfg.wrapper.model.gs_voxel_size <= 0: raise ValueError('gs_voxel_size must be positive.') for name in ('filter_flying_points_infer', 'filter_large_scale_infer'): if not isinstance(cfg.wrapper.model[name], bool): raise ValueError(f'{name} must be true or false.') max_scale_abs = cfg.wrapper.model.max_scale_abs if max_scale_abs is not None and max_scale_abs <= 0: raise ValueError('max_scale_abs must be null or positive.') pct = cfg.wrapper.model.max_scale_pct if pct is not None and not 0 < pct <= 100: raise ValueError('max_scale_pct must be null or in (0, 100].') # Paths in YAML are relative to the repository root; explicit CLI paths are relative to cwd. for key, override in [('SplatAnything_pretrained', args.ckpt_path), ('worldmirror2_pretrained', args.geometry_ckpt_path)]: path = Path(cfg.wrapper[key]).expanduser() if not path.is_absolute(): path = (Path.cwd() if override else ROOT) / path if not path.is_file(): raise FileNotFoundError(f'{key} checkpoint not found: {path}. See README.md for weight setup.') cfg.wrapper[key] = str(path.resolve()) return cfg def main(argv=None): args = parse_args(argv) import torch from omegaconf import OmegaConf from src.models.models.gaussian_worldmirror import GaussianWorldMirror from src.models.models.geometry_worldmirror import GeometryWorldMirror from src.utils.checkpoints import load_model_weights from src.utils.scene_io import build_views, export_gaussians, export_points, export_depth_rgb from src.utils.sky_segmentation import SkySegmenter cfg = load_config(args) device = torch.device(args.device) if device.type == 'cuda' and not torch.cuda.is_available(): raise RuntimeError('CUDA is unavailable. Use a working CUDA/PyTorch environment, or --device cpu for a slow smoke run.') segmenter = SkySegmenter(args.sky_onnx_path, device=args.device) if args.sky_seg else None views, scene_name, img_paths = build_views( args.input_path, args.cameras, tuple(args.resolution), device, cam_align=not args.no_cam_align, max_images=args.max_images, sky_segmenter=segmenter, sky_max_fraction=args.sky_max_fraction) cond_flags = [1, 0, 1] if cfg.wrapper.enable_cond else [0, 0, 0] # Architecture, positional encoding, spacing normalization, and chunking use code defaults. model = GaussianWorldMirror(**OmegaConf.to_container(cfg.wrapper.model), enable_cond=bool(cfg.wrapper.enable_cond), filter_sky_infer=args.sky_seg) teacher = GeometryWorldMirror(enable_cond=bool(cfg.wrapper.enable_cond)) load_model_weights(model, cfg.wrapper.SplatAnything_pretrained) load_model_weights(teacher, cfg.wrapper.worldmirror2_pretrained) model.to(device).eval().requires_grad_(False) teacher.to(device).eval().requires_grad_(False) print(f'[Input] {scene_name}: {views["img"].shape[1]} images, resolution={args.resolution}, ' f'enable_cond={cfg.wrapper.enable_cond}, use_gt_camera={cfg.wrapper.model.use_gt_camera}, sky_seg={args.sky_seg}') started = time.perf_counter() with torch.inference_mode(): geometry = teacher(views, cond_flags=cond_flags) preds = model(views, geometry, cond_flags=cond_flags) if device.type == 'cuda': torch.cuda.synchronize() print(f'[Inference] {time.perf_counter() - started:.2f}s') outdir = Path(args.output_path) / scene_name outdir.mkdir(parents=True, exist_ok=True) # An empty filtered scene must not be mistaken for a successful reconstruction. if not len(preds['splats']['means'][0]): raise RuntimeError('All Gaussians were filtered. Check the input, sky masks, and scale threshold.') if not (preds['splats']['opacities'][0] >= cfg.wrapper.infer_opacity_thresh).any(): raise RuntimeError('All Gaussians are below opacity_thresh. Lower the threshold and check the checkpoint.') export_gaussians(preds, outdir, cfg.wrapper.infer_opacity_thresh, None) export_points(preds, outdir) if not args.no_depth_rgb: export_depth_rgb(preds, views, img_paths, outdir) cameras = {'extrinsics': [], 'intrinsics': []} for index, image in enumerate(img_paths): camera_id = Path(image).stem for name, key in [('extrinsics', 'export_camera_poses'), ('intrinsics', 'export_camera_intrs')]: cameras[name].append({'camera_id': camera_id, 'matrix': preds[key][0, index].cpu().tolist()}) (outdir / 'cameras.json').write_text(json.dumps(cameras, indent=2) + '\n') OmegaConf.save(cfg, outdir / 'config.yaml', resolve=True) (outdir / 'run.json').write_text(json.dumps({**vars(args), 'images': img_paths, 'cond_flags': cond_flags}, indent=2) + '\n') print(f'[Done] {outdir}') if __name__ == '__main__': main()