vrscan_pipeline.py

#!/usr/bin/env python3
"""
VRScan Pipeline
Automates: OpenQuestCapture zip -> COLMAP -> LichtFeld Studio -> .ply for VRChat

Usage:
  python vrscan_pipeline.py --input capture.zip --name my_room
  python vrscan_pipeline.py --input capture.zip --name my_room --skip-to train
"""

import argparse
import os
import shutil
import subprocess
import sys
import zipfile
from pathlib import Path

# -- CONFIG: update these paths for your machine -----------------------------------------------

# Path to cloned quest-3d-reconstruction repo
QUEST_RECON_DIR = Path("./quest-3d-reconstruction")

# Path to LichtFeld Studio binary
# Windows: download from github.com/MrNeRF/LichtFeld-Studio/releases
# Linux:   ./LichtFeld-Studio/build/LichtFeld-Studio
LICHTFELD_BIN = Path("./LichtFeld-Studio/bin/LichtFeld-Studio.exe")

# Output directory for .ply files ready for ALCOM / Unity import
VRCHAT_OUTPUT_DIR = Path("./vrchat_ready")

# -- STEP 1: Unpack zip from OpenQuestCapture --------------------------------------------------
# Export via: Y button > Export Data inside OpenQuestCapture app
# Pull via USB: /Quest 3/Internal Shared Storage/data/com.samusynth.OpenQuestCapture/files

def step1_unpack(zip_path: Path, work_dir: Path) -> Path:
    print(f"\n[1/4] Unpacking capture: {zip_path}")
    project_dir = work_dir / "raw_capture"
    project_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(project_dir)
    print(f"      -> {project_dir}")
    return project_dir

# -- STEP 2: quest-3d-reconstruction -> COLMAP -------------------------------------------------
# Actual command from github.com/samuelm2/quest-3d-reconstruction README:
#
#   python scripts/e2e_quest_to_colmap.py \
#     --project_dir path/to/raw_capture \
#     --output_dir  path/to/colmap_project \
#     --use_colored_pointcloud
#
# For indoor scenes with windows enable tone mapping in config/pipeline_config.yml:
#   yuv_to_rgb:
#     tone_mapping: true
#     tone_mapping_method: "clahe+gamma"
#     clahe_clip_limit: 2.0

def step2_reconstruct(project_dir: Path, work_dir: Path) -> Path:
    print(f"\n[2/4] Running quest-3d-reconstruction (COLMAP export)...")
    colmap_output = work_dir / "colmap_project"
    colmap_output.mkdir(parents=True, exist_ok=True)
    cmd = [
        sys.executable,
        str(QUEST_RECON_DIR / "scripts" / "e2e_quest_to_colmap.py"),
        "--project_dir", str(project_dir),
        "--output_dir",  str(colmap_output),
        "--use_colored_pointcloud",
    ]
    print(f"      CMD: {' '.join(cmd)}")
    subprocess.run(cmd, check=True)
    print(f"      -> {colmap_output}")
    return colmap_output

# -- STEP 3: LichtFeld Studio -> point_cloud.ply ----------------------------------------------
# Actual CLI flags from github.com/MrNeRF/LichtFeld-Studio wiki:
#
#   ./LichtFeld-Studio \
#     -d  <colmap_dir>       path to COLMAP dataset (needs images/ + sparse/0/)
#     -o  <output_dir>       output directory
#     -i  30000              training iterations
#     --strategy mcmc        better convergence for indoor scenes
#     --max-cap 500000       500k Gaussians -- good balance for room-scale
#     --headless             no GUI, terminal only (required for automation)
#     --bilateral-grid       helps with uneven indoor lighting

def step3_train_splat(colmap_dir: Path, work_dir: Path) -> Path:
    print(f"\n[3/4] Training Gaussian Splat with LichtFeld Studio...")
    print(f"      Expected time: 10-30 min depending on GPU")
    splat_output = work_dir / "splat_output"
    splat_output.mkdir(parents=True, exist_ok=True)
    cmd = [
        str(LICHTFELD_BIN),
        "-d", str(colmap_dir),
        "-o", str(splat_output),
        "--strategy", "mcmc",
        "--max-cap", "500000",
        "--headless",
        "-i", "30000",
        "--bilateral-grid",
    ]
    print(f"      CMD: {' '.join(cmd)}")
    subprocess.run(cmd, check=True)
    ply_path = splat_output / "point_cloud" / "iteration_30000" / "point_cloud.ply"
    if not ply_path.exists():
        found = list(splat_output.rglob("*.ply"))
        if found:
            ply_path = found[0]
        else:
            raise FileNotFoundError(f"No .ply found in {splat_output}")
    print(f"      -> {ply_path}")
    return ply_path

# -- STEP 4: Stage .ply for ALCOM / Unity import -----------------------------------------------
# VRChat import uses ALCOM (https://vrc-get.anatawa12.com/alcom/)
# Add VRChatGaussianSplatting package in ALCOM, then import .ply in Unity.
# CRITICAL: Disable MSAA in Unity Project Settings before uploading world.

def step4_stage_for_vrchat(ply_path: Path, name: str) -> Path:
    print(f"\n[4/4] Staging .ply for ALCOM / Unity import...")
    VRCHAT_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    dest = VRCHAT_OUTPUT_DIR / f"{name}.ply"
    shutil.copy2(ply_path, dest)
    instructions = VRCHAT_OUTPUT_DIR / f"{name}_IMPORT_INSTRUCTIONS.txt"
    instructions.write_text(f"""VRChat Import Instructions -- {name}
=====================================================

PACKAGE MANAGER: ALCOM (https://vrc-get.anatawa12.com/alcom/)
  Download ALCOM and use it instead of VRChat Creator Companion.
  ALCOM is faster and open-source. Uses the same project settings as VCC.

UNITY PACKAGE:
  github.com/MichaelMoroz/VRChatGaussianSplatting
  Add via ALCOM: Packages > Add Package from URL

IMPORT STEPS:
  1. Open your VRChat Unity project via ALCOM
  2. Assets > Import Gaussian Splat > select {name}.ply
  3. Optional: enable "presorting" for better compatibility
  4. Add GaussianSplattingRenderer prefab to your scene
  5. Drag the imported splat asset into the renderer component

REQUIRED SETTINGS:
  - DISABLE MSAA: Edit > Project Settings > Quality > Anti Aliasing = Disabled
  - Sorting Distance: min=0.1, max=50 for room-scale scenes
  - Quest standalone: reduce splat count to 100k-200k max

PLY FILE: {dest}
""")
    print(f"      -> {dest}")
    return dest

# -- MAIN --------------------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="VRScan: Quest 3 capture -> Gaussian Splat -> VRChat"
    )
    parser.add_argument("--input",   required=True, help="Path to OpenQuestCapture export .zip")
    parser.add_argument("--output",  default="./vrscan_workspace", help="Working directory")
    parser.add_argument("--name",    default="my_scan", help="Name for this capture")
    parser.add_argument("--skip-to", choices=["reconstruct", "train", "stage"],
                        help="Skip earlier steps when re-running")
    args = parser.parse_args()

    zip_path = Path(args.input)
    work_dir = Path(args.output) / args.name
    work_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'='*58}")
    print(f"  VRScan Pipeline")
    print(f"  Input:  {zip_path}")
    print(f"  Output: {work_dir}")
    print(f"  Name:   {args.name}")
    print(f"{'='*58}")

    try:
        skip = args.skip_to
        project_dir = work_dir / "raw_capture"
        colmap_dir  = work_dir / "colmap_project"

        if skip not in ("reconstruct", "train", "stage"):
            project_dir = step1_unpack(zip_path, work_dir)
        if skip not in ("train", "stage"):
            colmap_dir = step2_reconstruct(project_dir, work_dir)
        if skip != "stage":
            ply_path = step3_train_splat(colmap_dir, work_dir)
        else:
            ply_path = next((work_dir / "splat_output").rglob("*.ply"))

        final_ply = step4_stage_for_vrchat(ply_path, args.name)

        print(f"\n{'='*58}")
        print(f"  DONE -- .ply ready for ALCOM / Unity import:")
        print(f"  {final_ply}")
        print(f"{'='*58}\n")

    except subprocess.CalledProcessError as e:
        print(f"\n[ERROR] Pipeline step failed: {e}")
        sys.exit(1)
    except FileNotFoundError as e:
        print(f"\n[ERROR] {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()