scripts/export_cosmos_input.py

#!/usr/bin/env python
"""Package a rendered conditioning set as Cosmos-Transfer2.5-2B control inputs.

``render_conditioning_set.py`` emits the *source of truth*: 16-bit metric
depth, raw-id segmentation, and a CG render, as PNG sequences. None of those
are directly consumable by Cosmos-Transfer2.5, which wants 1280x720 H.264
control videos at 16 FPS whose length is a multiple of 93 frames. This
converts the former into the latter.

Two decisions here are the difference between a control signal that works and
one that quietly degrades the generation:

**The background is a real photograph, not the point cloud.** The capture
camera in this episode is static to machine precision (verified: the per-frame
world-to-camera matrices are bit-identical), so a per-pixel *temporal median*
over the episode's real video removes the moving arm and leaves a clean,
robot-free photographic plate that is pixel-aligned with the render by
construction. Compositing the CG robot/object over that plate -- instead of
over the unprojected point cloud -- eliminates the point cloud's holes,
dithering and 320x180 blockiness from the appearance control entirely. It also
makes the Canny/edge control usable, since the background then carries real
texture edges rather than point-cloud speckle.

**Implausibly near background depth is rejected, not rendered.** The capture is a
ZED stereo pair, and in this frame its minimum reported depth is 233 mm. Stereo
matching failure saturates disparity, which puts the failed pixels at the sensor's
*near* limit -- so a cluster of background readings pinned just above 233 mm, in
exactly the dark and textureless parts of the frame (the black curtain, the far wall),
is a failure mode and not a measurement. Left in, those readings are the *brightest*
thing in an inverse-depth control: they both draw a hard bogus structure across the
upper frame and drag the normalisation range so everything real loses contrast. They
are dropped and re-filled from valid neighbours. This is applied to the background
only -- the robot genuinely does come within 30 cm of this wide-angle camera, and its
depth comes from the mesh renderer, which is never touched by the filter.

**Depth is normalised once for the whole clip, never per frame.** Cosmos's own
depth pipeline (DepthAnything) produces *relative inverse* depth, so this
converts metric metres to inverse depth and maps it to 8-bit. Doing that
per-frame -- the obvious implementation -- makes the depth video flicker as
the scene's near/far extremes change, and the model faithfully reproduces the
flicker as brightness pumping. A single global range fixed across the clip
costs nothing and removes the failure mode. Robust percentiles rather than
min/max, so one stray pixel cannot crush the usable range.

Usage:
    PYTHONPATH=src python scripts/export_cosmos_input.py \\
        --render-dir outputs/conditioning_cosmos_720p \\
        --episode AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera ext1 \\
        --start-row 20 \\
        --out-dir outputs/cosmos_input
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

import cv2
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("export_cosmos_input")

_FFMPEG = Path.home() / "miniconda3" / "envs" / "ffmpeg_libs" / "bin" / "ffmpeg"

#: Cosmos-Transfer2.5-2B hard requirements, from its model card. Checked, not
#: assumed: a silently-wrong resolution or frame count is the kind of thing
#: that produces a degraded generation rather than an error.
_REQUIRED_WIDTH, _REQUIRED_HEIGHT = 1280, 720
_REQUIRED_FPS = 16
_PREFERRED_FRAME_MULTIPLE = 93

#: Stable across frames *and* across runs. A palette that changes with the set
#: of ids present in a given frame makes the segmentation control flicker
#: identity, which is worse than having no segmentation control at all.
_SEG_PALETTE: dict[int, tuple[int, int, int]] = {
    0: (0, 0, 0),
    1: (0, 200, 255),
    2: (255, 80, 0),
    3: (80, 255, 80),
    4: (255, 0, 200),
    5: (255, 220, 0),
}


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--render-dir", type=Path, required=True,
                   help="output of render_conditioning_set.py (rgb/, depth/, seg/, meta.json)")
    p.add_argument("--out-dir", type=Path, required=True)
    p.add_argument("--episode", required=True, help="DROID episode uuid (for the real video)")
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--start-row", type=int, default=20,
                   help="trajectory row 0 corresponds to this real video frame index")
    p.add_argument("--background", choices=["plate", "cloud"], default="plate",
                   help="plate = real temporal-median photograph (default, recommended); "
                        "cloud = the renderer's own point-cloud background")
    p.add_argument("--fps", type=int, default=_REQUIRED_FPS)
    p.add_argument("--blur-sigma", type=float, default=18.0,
                   help="Gaussian sigma for the 'vis' (blurred RGB) control, in pixels at 720p")
    p.add_argument("--depth-percentile", type=float, default=1.0,
                   help="robust percentile for the global inverse-depth range (each tail)")
    p.add_argument("--min-background-depth", type=float, default=0.40,
                   help="metres. BACKGROUND depth readings nearer than this are treated as "
                        "stereo-matching failures rather than measurements and are re-filled "
                        "from valid neighbours. Robot/object mesh depth is never touched. "
                        "0 disables. See the module docstring for the justification.")
    p.add_argument("--depth-normalize", choices=["global", "perframe"], default="global",
                   help="global is strongly recommended; perframe is provided only for A/B")
    p.add_argument("--weight-vis-fg", type=float, default=0.40,
                   help="control_weight for the vis (blur) branch")
    p.add_argument("--weight-depth-fg", type=float, default=0.90,
                   help="control_weight for the depth branch")
    p.add_argument("--spec-name", default="pick_place_books",
                   help="'name' field written into each spec_*.json")
    p.add_argument("--guidance", type=float, default=3.0,
                   help="classifier-free guidance; 3 is what the shipped examples use")
    p.add_argument("--crf", type=int, default=12, help="x264 CRF for control videos (low = clean)")
    p.add_argument("--prompt-file", type=Path, default=None,
                   help="override the built-in prompt (must be < 300 words)")
    return p.parse_args()


# --------------------------------------------------------------------------- #
# ffmpeg
# --------------------------------------------------------------------------- #


class _H264Writer:
    """Raw RGB24 -> H.264/yuv420p mp4, via an ffmpeg pipe.

    cv2.VideoWriter is avoided deliberately: its codec/pixel-format selection
    depends on how OpenCV was built, and a control video that silently lands in
    a different pixel format or gets chroma-subsampled differently is exactly
    the kind of difference that is invisible in a preview and harmful to the
    model. Here every parameter is explicit.
    """

    def __init__(self, path: Path, width: int, height: int, fps: int, crf: int):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.path = path
        cmd = [
            str(_FFMPEG), "-y", "-loglevel", "error",
            "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{width}x{height}", "-r", str(fps), "-i", "pipe:0",
            "-c:v", "libx264", "-preset", "slow", "-crf", str(crf),
            "-pix_fmt", "yuv420p",
            # Closed GOP, no scene-cut re-keying: decoders (and dataloaders that
            # seek) then agree exactly on frame indices, which matters when four
            # control videos must stay frame-aligned with each other.
            "-g", "16", "-sc_threshold", "0",
            "-movflags", "+faststart",
            str(path),
        ]
        self._p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

    def write(self, rgb: np.ndarray) -> None:
        self._p.stdin.write(np.ascontiguousarray(rgb, dtype=np.uint8).tobytes())

    def close(self) -> None:
        self._p.stdin.close()
        err = self._p.stderr.read().decode(errors="replace")
        rc = self._p.wait()
        if rc != 0:
            raise RuntimeError(f"ffmpeg failed for {self.path}: {err}")


# --------------------------------------------------------------------------- #
# background plate
# --------------------------------------------------------------------------- #


def _real_video_path(episode: str, camera: str, render_meta_dir: Path) -> tuple[Path, str]:
    """Locate the episode mp4 for `camera`, resolving ext1/ext2 -> serial via the mask meta."""
    ep_dir = REPO_ROOT / "data" / "droid_raw" / episode / "recordings" / "MP4"
    serial = None
    for meta in sorted((REPO_ROOT / "outputs" / episode).glob("objects*/*/objects/1_mask/meta.json")):
        m = json.loads(meta.read_text())
        if m.get("camera") == camera:
            serial = m.get("camera_serial")
            break
    if serial is None:
        raise SystemExit(
            f"could not resolve {camera} to a camera serial from any mask meta.json under "
            f"outputs/{episode}/ -- pass the right --episode/--camera"
        )
    path = ep_dir / f"{serial}.mp4"
    if not path.exists():
        raise SystemExit(f"episode video not found: {path}")
    return path, serial


def _temporal_median_plate(video_path: Path, cache: Path) -> np.ndarray:
    """Per-pixel temporal median of the whole episode -> a robot-free RGB plate.

    Valid only because the capture camera does not move (asserted by the caller
    against the render's own per-frame extrinsics). The arm occupies any given
    background pixel for well under half the episode, so the median lands on
    the background at every pixel it ever vacates.
    """
    if cache.exists():
        logger.info("background plate: reusing %s", cache)
        return cv2.cvtColor(cv2.imread(str(cache)), cv2.COLOR_BGR2RGB)

    cap = cv2.VideoCapture(str(video_path))
    frames = []
    while True:
        ok, bgr = cap.read()
        if not ok:
            break
        frames.append(bgr)
    cap.release()
    if not frames:
        raise SystemExit(f"no frames decoded from {video_path}")

    plate_bgr = np.median(np.stack(frames), axis=0).astype(np.uint8)
    cache.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(cache), plate_bgr)
    logger.info("background plate: temporal median of %d real frames -> %s", len(frames), cache)
    return cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2RGB)


def _read_real_frame(video_path: Path, index: int) -> np.ndarray | None:
    cap = cv2.VideoCapture(str(video_path))
    cap.set(cv2.CAP_PROP_POS_FRAMES, index)
    ok, bgr = cap.read()
    cap.release()
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) if ok else None


# --------------------------------------------------------------------------- #
# channel conversion
# --------------------------------------------------------------------------- #


def _fill_holes_smooth(values: np.ndarray, valid: np.ndarray, levels: int = 7) -> np.ndarray:
    """Fill depth holes by pyramid push-pull, keeping measured pixels exact.

    **Do not replace this with a nearest-valid-pixel fill.** That was the first
    implementation here and it produced a visible, serious artifact in the
    generated video: assigning every hole pixel the value of its nearest valid
    pixel is *precisely* a Voronoi partition, so a filled region breaks into
    flat polygonal cells meeting along straight seams. With ~17% of pixels
    holed and the holes forming large contiguous blobs, that paints a
    tessellation of fake planar facets into the depth control -- and the model
    faithfully rendered them as shattered-glass surfaces hanging in the scene.
    It is nearly invisible in a downscaled preview of the depth video and
    obvious the moment you run an edge detector over it, which is how it was
    finally caught.

    Push-pull instead: decimate value and coverage through a Gaussian pyramid,
    then reconstruct coarse-to-fine, letting each level supply only what the
    finer level lacks. The interpolant is smooth by construction, so it adds no
    edges that were not already in the data. Measured pixels are written back
    verbatim afterwards -- the fill only ever invents values where there were
    none, and it invents them smoothly, which is the honest representation of
    "we don't know" for a depth surface.
    """
    if valid.all():
        return values
    m = valid.astype(np.float32)
    vs, ms = [values.astype(np.float32) * m], [m]
    for _ in range(levels):
        vs.append(cv2.pyrDown(vs[-1]))
        ms.append(cv2.pyrDown(ms[-1]))

    up = vs[-1] / np.maximum(ms[-1], 1e-6)
    for k in range(len(vs) - 2, -1, -1):
        up = cv2.resize(up, (vs[k].shape[1], vs[k].shape[0]), interpolation=cv2.INTER_LINEAR)
        w = np.clip(ms[k], 0.0, 1.0)
        up = vs[k] / np.maximum(ms[k], 1e-6) * w + up * (1.0 - w)

    out = up
    out[valid] = values[valid]
    return out


def _valid_depth(depth_m: np.ndarray, seg: np.ndarray, min_background_m: float) -> np.ndarray:
    """Which depth pixels are measurements.

    Zero means "no data" by the renderer's convention. On top of that, a
    *background* reading nearer than ``min_background_m`` is rejected: see the
    module docstring -- those are the stereo's disparity saturating on dark or
    textureless regions, and they land at the sensor's near limit rather than
    anywhere near the truth. The foreground is exempt on purpose, because the
    arm really does come that close to this camera and its depth is rendered
    from the mesh, not measured.
    """
    valid = depth_m > 0
    if min_background_m > 0:
        valid &= ~((seg == 0) & (depth_m < min_background_m))
    return valid


def _depth_to_inverse_8bit(depth_m: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """Metric metres -> 8-bit relative inverse depth (near = bright), Cosmos convention."""
    inv = 1.0 / np.maximum(depth_m, 1e-6)
    return np.clip((inv - lo) / max(hi - lo, 1e-9), 0.0, 1.0).astype(np.float32)


def _colorize_seg(seg: np.ndarray) -> np.ndarray:
    out = np.zeros((*seg.shape, 3), dtype=np.uint8)
    for sid in np.unique(seg):
        out[seg == sid] = _SEG_PALETTE.get(int(sid), (128, 128, 128))
    return out


def _feathered_fg(seg: np.ndarray) -> np.ndarray:
    """Foreground coverage in [0,1] with a 1px feather, to kill compositing jaggies."""
    fg = (seg != 0).astype(np.float32)
    return cv2.GaussianBlur(fg, (3, 3), 0.8)


_DEFAULT_PROMPT = (
    "A Franka Panda robot arm with a black two-finger gripper works at a cluttered wooden "
    "laboratory workbench. The arm picks up a small blue plastic brick from the bench top, "
    "lifts it, carries it smoothly to the right, and sets it down on top of a tall stack of "
    "hardcover and paperback books sitting on the bench, then releases the brick and draws "
    "back. Behind the bench is a blue-painted pegboard wall with an aluminium extrusion frame "
    "and a yellow bar. On the bench are stacked books, a white plastic bin holding black "
    "binders, a white mug of pens, a cardboard box, and a black tool case. Below the bench top "
    "is a dark steel tool chest with one drawer pulled open. The room is lit by even overhead "
    "fluorescent light. The camera is fixed on a tripod and does not move. Photorealistic, "
    "sharp focus, natural indoor lighting, real laboratory footage."
)


def _settle_result(render_dir: Path) -> dict | None:
    """Whether this trajectory's release was settled under gravity, and to what.

    Returns None if the action sidecar shows no settle step, so the caller keeps
    the renderer's (then correct) "no settling" limitation.
    """
    path = render_dir / "action.json"
    if not path.exists():
        return None
    try:
        action = json.loads(path.read_text())
    except json.JSONDecodeError:
        return None
    sim, final = action.get("simulation"), action.get("final_pose")
    if not (isinstance(sim, dict) and sim.get("settled") and isinstance(final, dict)):
        return None
    return {
        "engine": f"MuJoCo {sim.get('mujoco_version', '?')}",
        "tilt_release_deg": float(final.get("tilt_from_flat_deg_at_release", float("nan"))),
        "tilt_final_deg": float(final.get("tilt_from_flat_deg", float("nan"))),
        "settle_time_s": float(sim.get("settle_time_s", float("nan"))),
    }


def _readme(args, meta: dict, n: int, w: int, h: int) -> str:
    d = meta["depth"]
    return f"""\
# Cosmos-Transfer2.5-2B input bundle

{n} frames, {w}x{h}, {meta['fps']} FPS = {meta['duration_s']:.3f} s.
{n} = {n // 93} x 93, the frame multiple the model card reports as best.

## What to feed the model

| File | Role |
|---|---|
| `prompt.txt` | text prompt ({len(meta['prompt'].split())} words; the card requires < 300) |
| `input_video.mp4` | the video being transferred: real background plate + CG robot/brick |
| `depth.mp4` | **depth control** -- relative *inverse* depth, 8-bit, near = bright |
| `fg_mask.mp4` | binary robot+brick mask, for a control's `mask_path` |
| `seg.mp4` | segmentation control (stable palette, see `export_meta.json`) |
| `edge.mp4` | Canny edge control |
| `vis.mp4` | a pre-blurred copy -- **normally unused**, see below |
| `first_frame.png` | **real photograph** of this scene at the trajectory's t=0 |
| `spec_*.json` | ready-to-run parameter files (paths relative to the spec) |

## Run it

    python examples/inference.py -i <this dir>/spec_depth_vis.json -o outputs/run1
    # multi-GPU:
    torchrun --nproc_per_node=4 examples/inference.py -i ... -o ...

| Spec | What it does |
|---|---|
| `spec_depth_vis.json` | **start here.** depth {args.weight_depth_fg} + vis {args.weight_vis_fg} |
| `spec_depth_masked_vis.json` | same, but depth applies only to the robot/brick |
| `spec_depth_only.json` | depth alone at 1.0, matching the shipped example |
| `spec_multicontrol.json` | all four branches |

## Why `vis` has no `control_path`

The repo computes the blur itself from `video_path` when you give `vis` a
`preset_blur_strength` instead of a control video -- and that preset is, by
construction, the blur the branch was trained on. `vis.mp4` here was made with a
hand-picked Gaussian sigma ({args.blur_sigma:g} px), which is a *guess* at that
distribution. The specs therefore use the preset. `vis.mp4` is kept only so you can
A/B it.

Do not raise the `vis` weight much: it is the strongest appearance control, and
pushing it up makes the output a de-blurred CG render rather than a photorealistic
video. `depth` is where our real information is.

## On masks

`mask_path` in this repo is **binary** -- white means "use this control here", black
means "don't". It is not a per-pixel weight map; per-control strength is the scalar
`control_weight`. `spec_depth_masked_vis.json` uses `fg_mask.mp4` to confine depth to
the robot and brick, which is worth trying because foreground depth is exact mesh
z-buffer at full resolution while background depth is upsampled from a 320x180
capture and partly interpolated.

## Two things that are easy to get wrong

**Depth is not metric here.** It is relative inverse depth, matching what
DepthAnything (Cosmos's own extractor) produces, normalised **once for the whole clip**
over `{d['global_range_m'][0]:.3f} .. {d['global_range_m'][1]:.3f} m`. Per-frame
normalisation makes the control flicker and the model turns that into brightness
pumping. The metric 16-bit millimetre PNGs remain in `{meta['source_render_dir']}`
and are the source of truth; regenerate from those, never from `depth.mp4`.

**The background is a real photograph.** The capture camera is static to machine
precision, so a per-pixel temporal median over the real episode footage removes the
moving arm and leaves `background_plate.png`, pixel-aligned with the render by
construction. Only the robot and the brick are CG. This is why `edge.mp4` is usable at
all -- the background carries real texture edges rather than point-cloud speckle.

## Spec schema

Taken from the repo's own shipped examples
(`assets/robot_example/*/*_spec.json`), not guessed: `name`, `prompt_path`,
`video_path`, `guidance`, and per-control `{{control_path, control_weight, mask_path}}`.
Paths are relative to the spec file.

## Known limitations inherited from the render

{chr(10).join(f'- **{k}**: {v}' for k, v in meta['source_render_meta'].items())}
- **background depth resolution**: {d['background_depth_native_resolution_note']}
  Holes were filled by nearest valid pixel; mean {d['hole_fraction_mean'] * 100:.2f}%,
  max {d['hole_fraction_max'] * 100:.2f}% of pixels per frame.
"""


def main() -> int:
    args = parse_args()
    setup_logging()

    rd = args.render_dir
    render_meta = json.loads((rd / "meta.json").read_text())
    camera_json = json.loads((rd / "camera.json").read_text())

    # render_conditioning_set.py states "no_settling_at_release" unconditionally,
    # because it cannot know whether the trajectory it was handed had already
    # been through the settle step. If the action sidecar shows it has, that
    # limitation is simply false for this bundle, and shipping it would tell a
    # reader the opposite of what the data does.
    limitations = dict(render_meta.get("limitations", {}))
    settle = _settle_result(rd)
    if settle is not None:
        limitations["no_settling_at_release"] = (
            "SUPERSEDED for this bundle: the released object WAS settled under gravity "
            f"({settle['engine']}). Its tilt at release was {settle['tilt_release_deg']:.2f} deg "
            f"and {settle['tilt_final_deg']:.2f} deg after settling, reached in "
            f"{settle['settle_time_s']:.3f} s of simulated time. Mass is not load-bearing: "
            "a rigid body's fall-and-settle trajectory is mass-independent, which was verified "
            "by re-running at 10x mass. Everything else in no_dynamics still holds -- the "
            "transport motion itself is kinematic."
        )

    rgb_paths = sorted((rd / "rgb").glob("*.png"))
    depth_paths = sorted((rd / "depth").glob("*.png"))
    seg_paths = sorted((rd / "seg").glob("*.png"))
    n = len(rgb_paths)
    if not (n == len(depth_paths) == len(seg_paths)):
        raise SystemExit(f"channel length mismatch: rgb={n} depth={len(depth_paths)} "
                         f"seg={len(seg_paths)}")

    h, w = cv2.imread(str(rgb_paths[0])).shape[:2]

    # --- hard requirement checks, up front ---------------------------------
    problems = []
    if (w, h) != (_REQUIRED_WIDTH, _REQUIRED_HEIGHT):
        problems.append(f"resolution is {w}x{h}, Cosmos-Transfer2.5-2B requires "
                        f"{_REQUIRED_WIDTH}x{_REQUIRED_HEIGHT}")
    if args.fps != _REQUIRED_FPS:
        problems.append(f"--fps {args.fps}, the model produces {_REQUIRED_FPS} FPS")
    if problems:
        raise SystemExit("refusing to export:\n  - " + "\n  - ".join(problems))
    if n % _PREFERRED_FRAME_MULTIPLE != 0:
        logger.warning(
            "%d frames is not a multiple of %d; the model card reports multiples of %d "
            "perform best (re-run resample_trajectory.py with --frames %d)",
            n, _PREFERRED_FRAME_MULTIPLE, _PREFERRED_FRAME_MULTIPLE,
            _PREFERRED_FRAME_MULTIPLE * max(1, round(n / _PREFERRED_FRAME_MULTIPLE)),
        )

    # The plate is only valid if the camera really is static. Check it rather
    # than trusting the earlier observation.
    W = np.asarray(camera_json["world_to_cam_per_frame"])
    cam_motion = float(np.abs(W - W[0]).max())
    if args.background == "plate" and cam_motion > 1e-9:
        raise SystemExit(
            f"--background plate requires a static camera, but the per-frame extrinsics vary "
            f"by {cam_motion:.3e}. Use --background cloud, which re-renders per frame."
        )

    video_path, serial = _real_video_path(args.episode, args.camera, rd)
    args.out_dir.mkdir(parents=True, exist_ok=True)

    plate = None
    if args.background == "plate":
        plate = _temporal_median_plate(video_path, args.out_dir / "background_plate.png")
        # The renderer's wording covers appearance and geometry together; with a
        # photographic plate only the geometry half still applies.
        limitations["single_viewpoint_capture"] = (
            "Applies to the DEPTH channel only in this bundle. Background depth is still a "
            "single-viewpoint 2.5D capture, so surfaces the camera never saw have no points "
            "and any view far from the capture pose exposes real holes. Background "
            "*appearance* no longer has this limitation: rgb/vis/edge take their background "
            "from a real photograph (temporal median of the episode), not from the cloud."
        )
        if plate.shape[:2] != (h, w):
            plate = cv2.resize(plate, (w, h), interpolation=cv2.INTER_AREA)

    # --- pass 1: global inverse-depth range --------------------------------
    logger.info("pass 1/2: scanning %d depth frames for the global inverse-depth range", n)
    samples, rejected = [], 0
    total_bg = 0
    for i, dp in enumerate(depth_paths):
        d_m = cv2.imread(str(dp), cv2.IMREAD_UNCHANGED).astype(np.float32) / 1000.0
        seg = cv2.imread(str(seg_paths[i]), cv2.IMREAD_UNCHANGED)
        valid = _valid_depth(d_m, seg, args.min_background_depth)
        bg = seg == 0
        total_bg += int(bg.sum())
        rejected += int((bg & (d_m > 0) & ~valid).sum())
        if valid.any():
            samples.append(1.0 / d_m[valid])
    allinv = np.concatenate(samples)
    if args.min_background_depth > 0:
        logger.info(
            "rejected %d background depth readings nearer than %.2f m as stereo-matching "
            "failures (%.2f%% of all background pixels)",
            rejected, args.min_background_depth, 100.0 * rejected / max(total_bg, 1),
        )
    lo = float(np.percentile(allinv, args.depth_percentile))
    hi = float(np.percentile(allinv, 100.0 - args.depth_percentile))
    logger.info(
        "global inverse-depth range: [%.4f, %.4f] 1/m  (= %.3f m far .. %.3f m near) "
        "from %.1f/%.1f percentiles",
        lo, hi, 1.0 / max(lo, 1e-9), 1.0 / max(hi, 1e-9),
        args.depth_percentile, 100 - args.depth_percentile,
    )

    # --- pass 2: write every video -----------------------------------------
    fps, crf = args.fps, args.crf
    W_ = dict(width=w, height=h, fps=fps, crf=crf)
    writers = {
        "input_video": _H264Writer(args.out_dir / "input_video.mp4", **W_),
        "vis": _H264Writer(args.out_dir / "vis.mp4", **W_),
        "depth": _H264Writer(args.out_dir / "depth.mp4", **W_),
        "seg": _H264Writer(args.out_dir / "seg.mp4", **W_),
        "edge": _H264Writer(args.out_dir / "edge.mp4", **W_),
        # Binary, because that is what cosmos-transfer2.5's `mask_path` actually
        # takes: white = apply this control here, black = do not. It is NOT a
        # continuous weight map -- per-control strength is the scalar
        # `control_weight`, and there is no per-pixel weighting to export.
        "fg_mask": _H264Writer(args.out_dir / "fg_mask.mp4", **W_),
        # The complement, so a control can be confined to the *background*.
        # This is what makes the split configuration possible: our background is
        # a real photograph and our foreground is CG with exact geometry, so the
        # two regions want opposite treatment -- preserve the background's real
        # texture, but let the model re-imagine the robot's appearance under a
        # hard geometric constraint.
        "bg_mask": _H264Writer(args.out_dir / "bg_mask.mp4", **W_),
    }

    ksize = int(2 * round(3 * args.blur_sigma) + 1)
    hole_fracs, first_composite = [], None
    logger.info("pass 2/2: writing %d frames of 8 control videos", n)
    try:
        for i in range(n):
            render_rgb = cv2.cvtColor(cv2.imread(str(rgb_paths[i])), cv2.COLOR_BGR2RGB)
            seg = cv2.imread(str(seg_paths[i]), cv2.IMREAD_UNCHANGED)
            d_mm = cv2.imread(str(depth_paths[i]), cv2.IMREAD_UNCHANGED).astype(np.float32)

            fg = _feathered_fg(seg)[:, :, None]
            base = plate if plate is not None else render_rgb
            composite = (fg * render_rgb + (1.0 - fg) * base).astype(np.uint8)
            if i == 0:
                first_composite = composite.copy()

            # depth
            d_m_raw = d_mm / 1000.0
            valid = _valid_depth(d_m_raw, seg, args.min_background_depth)
            hole_fracs.append(float(1.0 - valid.mean()))
            d_m = _fill_holes_smooth(d_m_raw, valid)
            if args.depth_normalize == "perframe":
                inv = 1.0 / np.maximum(d_m, 1e-6)
                f_lo = float(np.percentile(inv, args.depth_percentile))
                f_hi = float(np.percentile(inv, 100 - args.depth_percentile))
                dnorm = _depth_to_inverse_8bit(d_m, f_lo, f_hi)
            else:
                dnorm = _depth_to_inverse_8bit(d_m, lo, hi)
            depth_u8 = np.repeat((dnorm * 255).astype(np.uint8)[:, :, None], 3, axis=2)

            blurred = cv2.GaussianBlur(composite, (ksize, ksize), args.blur_sigma)
            gray = cv2.cvtColor(composite, cv2.COLOR_RGB2GRAY)
            edge = np.repeat(cv2.Canny(gray, 100, 200)[:, :, None], 3, axis=2)
            hard_fg = (seg != 0)
            mask_u8 = np.repeat((hard_fg * 255).astype(np.uint8)[:, :, None], 3, axis=2)

            writers["input_video"].write(composite)
            writers["vis"].write(blurred)
            writers["depth"].write(depth_u8)
            writers["seg"].write(_colorize_seg(seg))
            writers["edge"].write(edge)
            writers["fg_mask"].write(mask_u8)
            writers["bg_mask"].write(255 - mask_u8)

            if i % 40 == 0:
                logger.info("  frame %d/%d", i, n)
    finally:
        for wr in writers.values():
            wr.close()

    # --- first frame -------------------------------------------------------
    real_first = _read_real_frame(video_path, args.start_row)
    if real_first is not None:
        if real_first.shape[:2] != (h, w):
            real_first = cv2.resize(real_first, (w, h), interpolation=cv2.INTER_AREA)
        cv2.imwrite(str(args.out_dir / "first_frame.png"),
                    cv2.cvtColor(real_first, cv2.COLOR_RGB2BGR))
        # Alignment is the whole reason a real photo can be used here: the
        # trajectory's frame 0 replays the real recorded joints of this very
        # video frame, so the CG arm should land on the real arm. Emit the
        # blend so that claim is checkable by eye rather than asserted.
        blend = (0.5 * real_first + 0.5 * first_composite).astype(np.uint8)
        cv2.imwrite(str(args.out_dir / "first_frame_alignment_check.png"),
                    cv2.cvtColor(blend, cv2.COLOR_RGB2BGR))
    cv2.imwrite(str(args.out_dir / "first_frame_composite.png"),
                cv2.cvtColor(first_composite, cv2.COLOR_RGB2BGR))

    prompt = args.prompt_file.read_text().strip() if args.prompt_file else _DEFAULT_PROMPT
    n_words = len(prompt.split())
    if n_words >= 300:
        raise SystemExit(f"prompt is {n_words} words; the model card requires fewer than 300")
    (args.out_dir / "prompt.txt").write_text(prompt + "\n")

    # Spec schema taken from the repo's own examples
    # (third_party/cosmos-transfer2.5/assets/robot_example/*/*.json), not guessed.
    # Paths are relative to the spec file, which is how the shipped examples do it.
    #
    # `vis` deliberately carries NO control_path: the repo computes the blur on the
    # fly from `video_path` using its own `preset_blur_strength`, which is by
    # construction the blur the branch was trained on. Our hand-rolled vis.mp4 is
    # still exported, but a hand-chosen Gaussian sigma is a guess at that
    # distribution and the preset is not, so the preset is what the specs use.
    def _write_spec(name: str, body: dict) -> None:
        (args.out_dir / f"spec_{name}.json").write_text(
            json.dumps({"name": f"{args.spec_name}_{name}", "prompt_path": "prompt.txt",
                        "video_path": "input_video.mp4", "guidance": args.guidance, **body},
                       indent=2) + "\n"
        )

    _write_spec("depth_vis", {
        "depth": {"control_path": "depth.mp4", "control_weight": args.weight_depth_fg},
        "vis": {"control_weight": args.weight_vis_fg, "preset_blur_strength": "high"},
    })
    # Depth restricted to the robot + object. `mask_path` is binary (white = use the
    # control here), NOT a weight map. Worth trying because our background depth is
    # the weakest channel -- upsampled from a 320x180 capture and part-interpolated --
    # while the foreground depth is exact mesh z-buffer at full resolution.
    _write_spec("depth_masked_vis", {
        "depth": {"control_path": "depth.mp4", "control_weight": 1.0,
                  "mask_path": "fg_mask.mp4"},
        "vis": {"control_weight": args.weight_vis_fg, "preset_blur_strength": "high"},
    })
    _write_spec("depth_only", {
        "depth": {"control_path": "depth.mp4", "control_weight": 1.0},
    })
    # The one that actually targets photorealism.
    #
    # `image_context_path` is the style reference, and its default is the trap:
    # "If None and context_frame_idx is not provided, use a RANDOM FRAME FROM
    # THE INPUT VIDEO". Our input video is a CG composite, so leaving it unset
    # silently anchors the output's appearance to CG -- which is exactly the
    # failure seen on the first run. first_frame.png is a real photograph of
    # this scene at this trajectory's t=0, so it is the right style anchor, and
    # it is geometrically consistent with every control channel because the
    # capture camera never moves.
    #
    # `vis` is dropped rather than down-weighted: config.py rejects vis and
    # image_context_path together ("both used to transfer style ... please only
    # provide one"). Geometry still comes from depth, which is where our real
    # information is anyway.
    _write_spec("depth_realstyle", {
        "image_context_path": "first_frame.png",
        "depth": {"control_path": "depth.mp4", "control_weight": args.weight_depth_fg},
    })
    _write_spec("multicontrol", {
        "depth": {"control_path": "depth.mp4", "control_weight": args.weight_depth_fg},
        "edge": {"control_path": "edge.mp4", "control_weight": 0.2},
        "seg": {"control_path": "seg.mp4", "control_weight": 1.0},
        "vis": {"control_weight": args.weight_vis_fg, "preset_blur_strength": "high"},
    })

    meta = {
        "target_model": "nvidia/Cosmos-Transfer2.5-2B",
        "prompt": prompt,
        "prompt_words": n_words,
        "resolution": [w, h],
        "frames": n,
        "fps": fps,
        "duration_s": n / fps,
        "frames_is_multiple_of_93": n % 93 == 0,
        "source_render_dir": str(rd),
        "source_render_meta": limitations,
        "episode": args.episode,
        "camera": args.camera,
        "camera_serial": serial,
        "camera_static_max_extrinsic_deviation": cam_motion,
        "background": {
            "mode": args.background,
            "method": (
                "per-pixel temporal median of all real frames of the episode video, which "
                "removes the moving arm and leaves a robot-free photographic plate. Valid "
                "only because the capture camera is static (checked above)."
                if args.background == "plate" else
                "the renderer's unprojected point-cloud background"
            ),
        },
        "depth": {
            "convention": "relative INVERSE depth, near = bright, 8-bit -- matching what "
                          "DepthAnything (Cosmos's own depth extractor) produces, NOT metric "
                          "depth. The metric 16-bit millimetre PNGs in the render dir remain "
                          "the source of truth.",
            "normalization": args.depth_normalize,
            "global_inverse_range_per_m": [lo, hi],
            "global_range_m": [1.0 / hi, 1.0 / lo],
            "why_global": "Per-frame normalization makes the control video flicker as the "
                          "scene's near/far extremes change, and the model reproduces that as "
                          "brightness pumping in the output.",
            "hole_fill": "pyramid push-pull (smooth interpolation), measured pixels kept exact, applied before inversion. NOT nearest-neighbour: that is a Voronoi partition and tessellates the filled area into flat polygonal facets, which the model renders as fake glass-like planes.",
            "hole_fraction_mean": float(np.mean(hole_fracs)),
            "hole_fraction_max": float(np.max(hole_fracs)),
            "min_background_depth_m": args.min_background_depth,
            "background_readings_rejected": rejected,
            "background_readings_rejected_fraction": rejected / max(total_bg, 1),
            "why_reject_near_background":
                "The ZED's minimum reported depth in this capture is 233 mm. Stereo matching "
                "failure saturates disparity and therefore lands at the sensor's NEAR limit, so "
                "background readings pinned just above that minimum -- concentrated in the dark "
                "curtain and the far wall -- are failures, not measurements. In an inverse-depth "
                "control they are the brightest pixels in the frame: they draw a hard bogus "
                "structure across the upper frame and compress the range available to everything "
                "real. Foreground is exempt: the arm genuinely reaches within 30 cm of this "
                "wide-angle camera, and its depth is rendered from the mesh, not measured.",
            "background_depth_native_resolution_note":
                "Background depth is unprojected from a single ~320x180 capture, so it is "
                "smooth/blocky relative to the 1280x720 output. Robot and object depth comes "
                "from the mesh renderer at full output resolution and is sharp.",
        },
        "vis": {
            "content": "the composited RGB (real plate + CG robot/object), Gaussian blurred",
            "blur_sigma_px": args.blur_sigma,
            "kernel_px": ksize,
            "not_used_by_default": "The generated specs give the vis branch a "
                                   "preset_blur_strength instead of this file, so the repo blurs "
                                   "video_path itself with the strength the branch was trained "
                                   "on. This file's sigma is a hand-picked guess at that "
                                   "distribution; it is exported only for A/B.",
        },
        "segmentation": {"id_map": render_meta.get("seg_id_map", {}), "palette": _SEG_PALETTE},
        "control_weights": {
            "vis": args.weight_vis_fg,
            "depth": args.weight_depth_fg,
            "guidance": args.guidance,
            "rationale": "Depth is where our real information is -- foreground depth is exact "
                         "mesh z-buffer. vis is the strongest appearance control, so it starts "
                         "low; raising it turns the output into a de-blurred CG render.",
            "mask_semantics": "cosmos-transfer2.5's mask_path is BINARY (white = apply this "
                              "control here), not a per-pixel weight map. Per-control strength "
                              "is the scalar control_weight only.",
        },
        "video_encoding": {"codec": "libx264", "pix_fmt": "yuv420p", "crf": crf,
                           "gop": 16, "scene_cut_detection": False},
        "spec_schema_source":
            "third_party/cosmos-transfer2.5/assets/robot_example/*/*_spec.json -- the repo's own "
            "shipped examples, read directly rather than inferred from documentation.",
    }
    (args.out_dir / "export_meta.json").write_text(json.dumps(meta, indent=2, default=str))
    (args.out_dir / "README.md").write_text(_readme(args, meta, n, w, h))

    print("=" * 78)
    print(f"{args.out_dir}   {n} frames  {w}x{h} @ {fps} fps  = {n/fps:.3f} s")
    print(f"  multiple of 93:      {'yes' if n % 93 == 0 else 'NO'}  ({n/93:.3f} x 93)")
    print(f"  depth holes filled:  mean {np.mean(hole_fracs)*100:.2f}%  "
          f"max {np.max(hole_fracs)*100:.2f}%")
    print(f"  background:          {args.background}")
    for name in sorted(writers):
        f = args.out_dir / f"{name}.mp4"
        print(f"  {name+'.mp4':22s} {f.stat().st_size/1e6:7.2f} MB")
    print("=" * 78)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())