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 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 / "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 / "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 _reconstruct_cmd(job: Path) -> list[str]:
    ck = config.GENRECON_CKPTS
    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", str(config.RADIUS_M),
        "--proj_batch_voxels", str(config.PROJ_BATCH_VOXELS),
    ])


def run_reconstruct(job: Path, log: Log) -> None:
    (job / "out").mkdir(exist_ok=True)
    env = dict(os.environ, CUDA_VISIBLE_DEVICES=config.CUDA_DEVICE)
    _run(_reconstruct_cmd(job), cwd=config.GENRECON_SRC, log=log, env=env)


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)
    if not (out / "scene.glb").exists():
        raise RuntimeError("chunked_to_glb.py did not produce scene.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")