worker/sam2_split.py

"""SAM 2 dynamic/static split: a coherent mask instead of a saturated field.

Why this exists
---------------
MegaSaM's motion_prob for real handheld footage is saturated, not shifted:
run 26 measured 60% of pixels above 0.99 and 26.9% pinned at exactly 1.0, so
every threshold family (fixed, Otsu, percentile) either called the whole scene
dynamic or selected nothing. The ceiling-detection hack in prune.py works but
is a heuristic standing on a broken distribution.

SAM 2 sidesteps the distribution entirely: prompt it once on the frame where
motion evidence is strongest, let its video tracker propagate a temporally
consistent object mask through the clip, and use THAT as the dynamic region.
motion_prob is demoted from decision-maker to prompt-finder.

Strategy (fully automatic, no user click needed):
  1. pick the prompt frame + point from the flow-residual peak (or the
     motion_prob ceiling blob when available - it is spatially coherent even
     though its values are useless as probabilities),
  2. run SAM2 video propagation from that prompt,
  3. per frame: dynamic = SAM2 mask; fall back to the old ceiling/Otsu split
     only if SAM2 tracked nothing (mask area < 0.5% everywhere).

Output matches what prune.py consumes: one bool mask per subsampled frame,
flattened alongside the point cloud.
"""
from __future__ import annotations

import os

import numpy as np


def pick_prompt(motion_prob: np.ndarray) -> tuple[int, tuple[int, int]]:
    """Choose (frame_index, (x, y)) from the per-frame motion field (F, H, W).

    The ceiling blob's centroid on the frame where it is largest: saturation
    made the VALUES useless but the SUPPORT is still the moving subject.
    """
    mp = np.asarray(motion_prob, dtype=np.float32)
    ceil = mp >= mp.max() - 1e-6
    areas = ceil.reshape(len(mp), -1).mean(axis=1)
    f = int(np.argmax(areas))
    ys, xs = np.nonzero(ceil[f])
    if xs.size == 0:  # nothing at ceiling: fall back to the global argmax
        f, y, x = np.unravel_index(int(np.argmax(mp)), mp.shape)
        return int(f), (int(x), int(y))
    return f, (int(np.median(xs)), int(np.median(ys)))


def sam2_video_masks(frame_paths: list[str], prompt_frame: int,
                     prompt_xy: tuple[int, int], log=print) -> np.ndarray | None:
    """Propagate one point prompt through the clip; returns (F, H, W) bool.

    Returns None when SAM2 is unavailable or tracks nothing, so callers can
    fall back to the legacy split.

    SAM2's ``init_state`` does not accept a list of frame paths - its real
    signature is ``init_state(video_path, ...)`` where ``video_path`` is a
    directory of ``<index>.jpg`` files (``sam2/utils/misc.py:
    load_video_frames_from_jpg_images`` hard-filters on .jpg/.jpeg and sorts
    by ``int(stem)``). Instant4D's extracted frames are ``NNNNN.png``, so they
    are re-encoded into a throwaway numeric-jpg directory first.
    """
    try:
        import torch
        from sam2.build_sam import build_sam2_video_predictor
    except ImportError as e:
        log(f'[sam2] unavailable ({e}); falling back to legacy split')
        return None

    import shutil
    import tempfile
    import cv2

    jpg_dir = tempfile.mkdtemp(prefix='sam2_frames_')
    try:
        for i, p in enumerate(frame_paths):
            img = cv2.imread(p)
            if img is None:
                raise RuntimeError(f'could not read frame {p}')
            cv2.imwrite(os.path.join(jpg_dir, f'{i:05d}.jpg'), img,
                       [cv2.IMWRITE_JPEG_QUALITY, 95])

        ckpt = build_sam2_video_predictor('configs/sam2.1/sam2.1_hiera_s.yaml',
                                          'checkpoints/sam2.1_hiera_small.pt')
        with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16):
            state = ckpt.init_state(video_path=jpg_dir)
            ckpt.add_new_points_or_box(state, frame_idx=prompt_frame, obj_id=1,
                                       points=np.array([prompt_xy], dtype=np.float32),
                                       labels=np.array([1], dtype=np.int32))
            masks = {}
            for f_idx, _ids, logits in ckpt.propagate_in_video(state):
                masks[f_idx] = (logits[0, 0] > 0).cpu().numpy()
            # propagate_in_video only walks forward; run the reverse leg too.
            for f_idx, _ids, logits in ckpt.propagate_in_video(state, reverse=True):
                masks[f_idx] = (logits[0, 0] > 0).cpu().numpy()
    finally:
        shutil.rmtree(jpg_dir, ignore_errors=True)

    out = np.stack([masks[i] for i in range(len(frame_paths))])
    frac = out.reshape(len(out), -1).mean(axis=1)
    log(f'[sam2] mask area per frame: min {frac.min():.1%} p50 {np.median(frac):.1%} '
        f'max {frac.max():.1%}')
    if frac.max() < 0.005:
        log('[sam2] tracked nothing; falling back to legacy split')
        return None
    return out