genrecon/backend/pipeline.py

"""The reconstruction pipeline: phone video -> walkable scene.glb.

  frames   ffmpeg extracts frames from the uploaded video
  colmap   COLMAP SfM -> <job>/colmap/{cameras,images,points3D}.txt
  reconstruct  GenRecon reconstruct_scene.py --mode Iphone
  glb      chunked_to_glb.py -> <job>/out/scene.glb

MOCK mode short-circuits every stage and copies the bundled sample scene so
the API + PWA can be exercised end-to-end without a GPU.

NOTE for the executor (issue i-opavuyaq): the exact on-disk layout GenRecon's
`Iphone` mode wants is confirmed during the Step-2 smoke test. This module
produces the layout the README implies (`<job>/images/` + `<job>/colmap/cameras.txt`);
adjust `_reconstruct_cmd` / COLMAP output paths if the smoke test shows otherwise.
"""
from __future__ import annotations

import os
import shutil
import subprocess
import time
from collections import deque
from pathlib import Path
from typing import Callable

import config

Log = Callable[[str], None]


def _run(cmd: list[str], cwd: Path | None, log: Log, env: dict | None = None) -> None:
    """Run a command, streaming combined output into the job log. Raise on non-zero."""
    log(f"$ {' '.join(str(c) for c in cmd)}")
    proc = subprocess.Popen(
        cmd, cwd=str(cwd) if cwd else None,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, bufsize=1, env=env,
    )
    assert proc.stdout is not None
    for line in proc.stdout:
        log(line.rstrip())
    code = proc.wait()
    if code != 0:
        raise RuntimeError(f"command failed (exit {code}): {' '.join(str(c) for c in cmd)}")


def _conda(args: list[str]) -> list[str]:
    """Wrap a python invocation so it runs inside the genrecon conda env."""
    inner = "conda run -n %s --no-capture-output %s" % (config.CONDA_ENV, " ".join(args))
    return ["bash", "-lc", inner]


# --- stages -----------------------------------------------------------------

def extract_frames(job: Path, log: Log) -> None:
    images = job / "rgb"   # GenRecon Iphone mode reads frames from <job>/rgb/, not images/
    images.mkdir(exist_ok=True)
    src = next(job.glob("input.*"))
    _run([config.FFMPEG_BIN, "-y", "-i", str(src),
          "-vf", f"fps={config.FRAMES_FPS}", "-qscale:v", "2",
          str(images / "%06d.jpg")], cwd=job, log=log)
    frames = sorted(images.glob("*.jpg"))
    if not frames:
        raise RuntimeError("ffmpeg produced no frames — is the upload a valid video?")
    # Evenly subsample if we overshot the cap (keeps COLMAP tractable).
    if len(frames) > config.FRAMES_MAX:
        step = len(frames) / config.FRAMES_MAX
        keep = {frames[int(i * step)] for i in range(config.FRAMES_MAX)}
        for f in frames:
            if f not in keep:
                f.unlink()
    log(f"frames: {len(list(images.glob('*.jpg')))}")


def run_colmap(job: Path, log: Log) -> None:
    images = job / "rgb"   # GenRecon Iphone mode reads frames from <job>/rgb/, not images/
    db = job / "colmap.db"
    sparse = job / "sparse"
    out = job / "colmap"
    sparse.mkdir(exist_ok=True)
    out.mkdir(exist_ok=True)
    C = config.COLMAP_BIN
    # COLMAP links Qt and aborts on QApplication init in a headless container
    # (no display). QT_QPA_PLATFORM=offscreen makes every CLI subcommand run
    # without a display. SIFT extract/match are forced to CPU because the apt
    # COLMAP build's GPU SIFT needs an OpenGL/EGL context we don't have headless.
    cenv = dict(os.environ, QT_QPA_PLATFORM="offscreen")
    _run([C, "feature_extractor", "--database_path", str(db), "--image_path", str(images),
          "--ImageReader.single_camera", "1", "--SiftExtraction.use_gpu", "0"], cwd=job, log=log, env=cenv)
    _run([C, "exhaustive_matcher", "--database_path", str(db),
          "--SiftMatching.use_gpu", "0"], cwd=job, log=log, env=cenv)
    _run([C, "mapper", "--database_path", str(db), "--image_path", str(images),
          "--output_path", str(sparse)], cwd=job, log=log, env=cenv)
    model = sparse / "0"
    if not model.exists():
        raise RuntimeError("COLMAP mapper produced no model — not enough overlapping views?")
    _run([C, "model_converter", "--input_path", str(model),
          "--output_path", str(out), "--output_type", "TXT"], cwd=job, log=log, env=cenv)
    if not (out / "cameras.txt").exists():
        raise RuntimeError("COLMAP did not write colmap/cameras.txt")


def _resolve_radius_m(job: Path, log: Log) -> str:
    """radius_m is a metric radius; a plain-video COLMAP is only up-to-scale, so a
    fixed value either wipes the cloud or keeps outliers. When config.RADIUS_M is
    'auto', derive it from the actual median nearest-neighbour spacing of the COLMAP
    points (× config.RADIUS_K) so the outlier-clean is scale-invariant."""
    rm = str(config.RADIUS_M).strip().lower()
    if rm != "auto":
        return str(config.RADIUS_M)
    pts = job / "colmap" / "points3D.txt"
    if not pts.exists():
        log("radius_m=auto: no points3D.txt, falling back to 0.2"); return "0.2"
    coords = []
    for line in pts.read_text().splitlines():
        if not line or line[0] == "#":
            continue
        f = line.split()
        try:
            coords.append((float(f[1]), float(f[2]), float(f[3])))
        except (IndexError, ValueError):
            continue
    if len(coords) < 50:
        log(f"radius_m=auto: only {len(coords)} points, falling back to 0.2"); return "0.2"
    try:
        import numpy as np
        from scipy.spatial import cKDTree
        pc = np.asarray(coords)
        idx = np.linspace(0, len(pc) - 1, min(3000, len(pc))).astype(int)
        dist, _ = cKDTree(pc).query(pc[idx], k=2)
        median_nn = float(np.median(dist[:, 1]))
        radius = round(median_nn * config.RADIUS_K, 5)
        log(f"radius_m=auto → {radius} (median_nn={median_nn:.5f} × k={config.RADIUS_K}, n={len(pc)})")
        return str(radius) if radius > 0 else "0.2"
    except Exception as e:  # noqa: BLE001 — never let auto-scale break the run
        log(f"radius_m=auto failed ({e}); falling back to 0.2"); return "0.2"


def _reconstruct_cmd(job: Path, log: Log) -> list[str]:
    ck = config.GENRECON_CKPTS
    radius_m = _resolve_radius_m(job, log)
    return _conda([
        "python", "reconstruct_scene.py",
        "--mode", "Iphone",
        "--path", str(job),
        "--output_path", str(job / "out"),
        # GenRecon derives each stage's training config as <parent-of-ckpt-dir>/config.json
        # (setup_utils.load_train_config), so each checkpoint lives in its own
        # run dir: <ck>/<stage>/checkpoints/<name>.pt + <ck>/<stage>/config.json
        # (the matching configs/gen/*/genrecon*.json). Weights are symlinked so
        # the flat 13.7GB download is not duplicated.
        "--ss_ckpt", str(ck / "ss" / "checkpoints" / "sparse_structure.pt"),
        "--shape_ckpt", str(ck / "shape" / "checkpoints" / "shape_slat.pt"),
        "--tex_ckpt", str(ck / "tex" / "checkpoints" / "texture_slat.pt"),
        "--pipeline_config", config.PIPELINE_CONFIG,
        "--num_imgs_per_scene", str(config.NUM_IMGS_PER_SCENE),
        "--chunk_size_factor", str(config.CHUNK_SIZE_FACTOR),
        "--stat_std_ratio", str(config.STAT_STD_RATIO),
        "--radius_nb_points", str(config.RADIUS_NB_POINTS),
        "--radius_m", radius_m,
        "--proj_batch_voxels", str(config.PROJ_BATCH_VOXELS),
    ])


def _friendly_recon_error(recent: "deque[str]") -> str:
    """Turn a raw reconstruct crash into an actionable message for the uploader.
    GenRecon's research code dies with cryptic tracebacks when the input is a poor
    fit (single object, no depth, too little overlap); map the known signatures to
    plain guidance so a user knows how to re-shoot rather than seeing a stack trace."""
    blob = "\n".join(recent)
    if "Number points after cleaning: 0" in blob or "shape (0," in blob:
        return ("Reconstruction found no stable geometry. Film a well-lit, textured "
                "space and move the camera slowly with lots of overlap between frames "
                "— avoid blank walls, reflections and fast motion.")
    if "chunk_centers" in blob or ("_get_transforms" in blob and "index out of range" in blob):
        return ("The clip has no scene depth to reconstruct. Walk through a room / space "
                "so the camera sees both near and far surfaces (not just one object).")
    if "not enough overlapping views" in blob or "mapper produced no model" in blob:
        return ("Camera tracking failed — not enough overlapping views. Re-record moving "
                "slowly and steadily so consecutive frames share plenty of the scene.")
    if "CUDA out of memory" in blob or "OutOfMemoryError" in blob:
        return ("Ran out of GPU memory on this scene. Try a shorter clip or lower the "
                "per-scene image count.")
    return "Reconstruction failed. Try a slower, more thorough capture of a textured space."


def run_reconstruct(job: Path, log: Log) -> None:
    (job / "out").mkdir(exist_ok=True)
    env = dict(os.environ, CUDA_VISIBLE_DEVICES=config.CUDA_DEVICE)
    recent: "deque[str]" = deque(maxlen=80)

    def tee(line: str) -> None:
        recent.append(line)
        log(line)

    try:
        _run(_reconstruct_cmd(job, log), cwd=config.GENRECON_SRC, log=tee, env=env)
    except RuntimeError:
        raise RuntimeError(_friendly_recon_error(recent)) from None


def run_glb(job: Path, log: Log) -> None:
    out = job / "out"
    env = dict(os.environ, CUDA_VISIBLE_DEVICES=config.CUDA_DEVICE)
    _run(_conda([
        "python", "chunked_to_glb.py",
        "--inputs", str(out / "to_glb_inputs.pt"),
        "--chunk_inputs", str(out / "chunk_inputs.pt"),
        "--output_dir", str(out),
    ]), cwd=config.GENRECON_SRC, log=log, env=env)
    glb = out / "scene.glb"
    if not glb.exists():
        raise RuntimeError("chunked_to_glb.py did not produce scene.glb")
    _web_optimize_glb(glb, log)


def _web_optimize_glb(glb: Path, log: Log) -> None:
    """Shrink the raw multi-hundred-MB mesh into a browser-loadable walkthrough glb
    (simplify + webp, no geometry codec). Best-effort: the raw mesh is preserved as
    scene_raw.glb, and any failure leaves the original scene.glb in place — a heavy
    but correct asset always beats a broken job."""
    if not config.WEB_OPTIMIZE:
        return
    script = Path(config.OPTIMIZE_GLB_SCRIPT)
    if not script.exists():
        log(f"web-optimize skipped (script not found: {script})")
        return
    raw = glb.with_name("scene_raw.glb")
    try:
        shutil.copy2(glb, raw)
        # optimize the raw copy back into scene.glb (script never hard-fails)
        _run(["bash", str(script), str(raw), str(glb)], cwd=glb.parent, log=log)
        if not glb.exists() or glb.stat().st_size == 0:
            shutil.copy2(raw, glb)
            log("web-optimize produced no output; kept full-res glb")
    except Exception as e:  # noqa: BLE001 — optimization must never fail the job
        log(f"web-optimize error ({e}); kept full-res glb")
        if raw.exists() and (not glb.exists() or glb.stat().st_size == 0):
            shutil.copy2(raw, glb)


_REAL = {"frames": extract_frames, "colmap": run_colmap,
         "reconstruct": run_reconstruct, "glb": run_glb}


def run_pipeline(job: Path, log: Log, set_progress: Callable[[str, int, str], None]) -> None:
    """Drive all stages, reporting (stage, percent, message). Raises on failure."""
    for stage, lo, hi in config.STAGES:
        set_progress(stage, lo, f"{stage}…")
        if config.MOCK:
            for p in range(lo, hi + 1, max(1, (hi - lo) // 5)):
                set_progress(stage, p, f"{stage} (mock)…")
                time.sleep(0.4)
        else:
            _REAL[stage](job, log)
        set_progress(stage, hi, f"{stage} done")
    if config.MOCK:
        (job / "out").mkdir(exist_ok=True)
        shutil.copyfile(config.SAMPLE_GLB, job / "out" / "scene.glb")
        log("MOCK: copied sample scene.glb")