scripts/mask_object_manual.py

#!/usr/bin/env python
"""Mask stage variant for a SECOND object per episode, with an explicit prompt.

`run_object_pipeline.py`'s own `mask` stage always derives its SAM 3.1 prompt
from the episode's `current_task` sentence (`TaskPromptDeriver`), which only
ever targets the sentence's primary/direct-object noun phrase (e.g. "brick"
for "Put brick in drawer shelf and close drawer") -- there is no CLI override
for a *second* object like the drawer/shelf mentioned later in the same
sentence. This script is exactly `run_mask()`'s body, with the prompt
candidate list passed in explicitly instead of derived, and writing the same
on-disk layout (`1_mask/{meta.json,mask_frame0.npz,frame0_bgr.png,crop_*}`) so
`run_object_pipeline.py --stages mesh,align` can be pointed at this script's
`--outputs` directory afterwards and just work, unmodified.

Usage:
    PYTHONPATH=src python scripts/mask_object_manual.py \\
        --episode <uuid> --camera ext1 --clip 0:11 \\
        --prompts "drawer shelf,drawer,cabinet drawer,shelf" \\
        --outputs outputs/<uuid>/objects_shelf/0_11
"""

from __future__ import annotations

import argparse
import json
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.objects import crop  # noqa: E402
from fpgm.pipeline.frames import ClipFrameSource  # noqa: E402
from fpgm.types import DataError  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

# Reuse run_object_pipeline.py's private clip-context/metadata loaders instead
# of duplicating them -- same trick as scripts/render_fixed_initial_view.py.
import importlib.util

_spec = importlib.util.spec_from_file_location(
    "run_object_pipeline", REPO_ROOT / "scripts" / "run_object_pipeline.py"
)
rop = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = rop
_spec.loader.exec_module(rop)

logger = get_logger("mask_object_manual")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--episode", required=True)
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--clip", required=True)
    p.add_argument("--prompts", required=True, help="comma-separated SAM 3.1 prompt candidates, tried in order")
    p.add_argument("--outputs", type=Path, required=True)
    p.add_argument("--device", default="cuda")
    return p.parse_args()


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

    from fpgm.config import SegmentationConfig
    from fpgm.segmentation.sam3 import Sam3VideoSegmenter

    candidates = [c.strip() for c in args.prompts.split(",") if c.strip()]
    out_dir = rop.ensure_dir(rop.stage_dir(args.outputs, "mask"))
    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    logger.info("clip %s: manual prompt candidates %s", args.clip, candidates)

    seg_cfg = SegmentationConfig()
    segmenter = Sam3VideoSegmenter(seg_cfg, device=args.device)

    video_start = int(ctx.timing.clip_frame_to_video_frame(0))
    video_end = int(ctx.timing.clip_frame_to_video_frame(ctx.timing.n_frames))
    with ClipFrameSource(
        str(ctx.mp4_path), video_start, video_end, stride=ctx.timing.annotation_stride
    ) as frames:
        width, height = frames.resolution
        frame0_bgr = frames.read(0)
        cv2.imwrite(str(out_dir / "frame0_bgr.png"), frame0_bgr)

        masklets = None
        obj_id: int | None = None
        used_prompt = ""
        for prompt in candidates:
            segmenter.start_session(str(frames.frame_dir))
            try:
                segmenter.add_text_prompt(seg_cfg.prompt_frame_idx, prompt)
                found = segmenter.collect_masklets()
                if not found:
                    logger.info("prompt %r matched nothing; trying next candidate", prompt)
                    continue
                obj_id = segmenter.select_object(found)
                masklets, used_prompt = found, prompt
                break
            finally:
                segmenter.close_session()

        if masklets is None or obj_id is None:
            raise DataError(f"mask stage: SAM 3.1 found nothing for any of {candidates}")

        masklet = masklets[obj_id]
        if 0 not in masklet.frames or not masklet.frames[0].any():
            raise DataError(
                f"mask stage: object mask is empty on frame 0 of clip {args.clip}; frame 0 "
                "is the only frame with dense depth, so reconstruction cannot proceed -- "
                "try a different --clip"
            )
        mask_frame0 = masklet.frames[0]

        best_frame = crop.select_best_frame(masklet, prefer="largest")
        best_frame_bgr = frames.read(best_frame)
        obj_crop = crop.crop_from_mask(
            best_frame_bgr, masklet.frames[best_frame], frame_idx=best_frame, prompt=used_prompt
        )
        crop.save_debug(obj_crop, out_dir, frame_bgr=best_frame_bgr)

    np.savez_compressed(out_dir / "mask_frame0.npz", mask=mask_frame0)
    meta = {
        "episode": args.episode,
        "camera": args.camera,
        "camera_serial": ctx.serial,
        "clip": args.clip,
        "prompt": used_prompt,
        "obj_id": int(obj_id),
        "video_resolution": [width, height],
        "best_frame": int(best_frame),
        "n_masklets": len(masklets),
    }
    (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))
    print(f"used_prompt={used_prompt!r} -> {out_dir}")
    return 0


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