cozyclay 1.9.0 → 2.0.0

package/tools/track/py/masks.py

@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""SAM2 video masks of the one person in a bench clip (tracker silhouette evidence).
+
+Inputs are what a real user has: the video and GVHMR's obs NPZ (``bbx_xys`` =
+centre x, centre y, square size in image pixels; ``kp2d`` ViTPose confidences).
+No plate, truth mask or colour cue is read here.
+
+Prompting:
+* anchor = first frame whose obs box is valid and trusted (mean ViTPose
+  confidence >= TRUST_CONF; frame 0 on every clean clip). If no frame is
+  trusted the most confident valid frame is used; no valid box at all raises
+  ``no-person-prompt``. The anchor gets the obs box plus positive points at the
+  confident ViTPose shoulder/hip midpoints, nose, wrists and (at >= ANKLE_TRUST)
+  ankles inside the box (``--prompt joints``); frames before it are tracked
+  backwards from it. Measured on the 11 shaded truth clips (median of per-item
+  median IoU): box + centre point 0.000 (the centre of GVHMR's loose square box
+  lands on the floor between the legs), box only 0.000, torso midpoints 0.889,
+  joints 0.930, all 17 joints 0.899; without ankles (``upper``) the grey skin
+  set falls to 0.886. The part-coloured mannequin needs points on the head and
+  limbs or SAM2 returns a sub-part.
+* re-prompt: when the mask area changes by more than REPROMPT_CHANGE relative to
+  the previous frame in tracking order, that frame is re-prompted from its own
+  obs box (if trusted), at most once per frame, and tracking resumes from it.
+* chunked propagation: at most ``--chunk-frames`` frames are loaded per SAM2
+  state (frames offloaded to CPU); the next chunk is seeded with the previous
+  chunk's boundary mask. This bounds GPU memory for 15 s (362 frame) clips.
+
+Output ``<out>/masks.npz``: ``prob`` float16 T x H/2 x W/2 (sigmoid of the SAM2
+logits, area-downsampled), ``reprompted`` int32 frame list, ``area`` int32 full
+resolution foreground pixels per frame, ``anchor`` int32, ``chunks`` int32
+(direction, start, end) rows with direction 1 forward / -1 backward. A summary
+is written to ``<out>/masks.json``.
+"""
+from __future__ import annotations
+
+import argparse
+import contextlib
+import json
+import os
+import sys
+import tempfile
+import time
+
+import numpy as np
+
+REPROMPT_CHANGE = 0.4
+TRUST_CONF = 0.5
+DEFAULT_CHUNK_FRAMES = 180
+DEFAULT_MAX_RESERVED_MIB = 5632
+CONFIG = "configs/sam2.1/sam2.1_hiera_s.yaml"
+CHECKPOINT = os.path.expanduser("~/cclay-ingest/cozyfit/checkpoints/sam2.1_hiera_small.pt")
+OBJ_ID = 1
+
+
+class NoPersonPrompt(RuntimeError):
+    """The obs carries no usable person box to prompt SAM2 with."""
+
+    def __init__(self, detail: str):
+        super().__init__(f"no-person-prompt: {detail}")
+
+
+def box_from_xys(xys, width: int, height: int):
+    """obs ``bbx_xys`` row (centre x, centre y, size) -> clipped xyxy box, or None."""
+    xys = np.asarray(xys, dtype=np.float64)
+    if xys.shape != (3,) or not np.all(np.isfinite(xys)) or xys[2] <= 1.0:
+        return None
+    cx, cy, size = xys
+    half = size / 2.0
+    x0, x1 = np.clip([cx - half, cx + half], 0.0, width - 1.0)
+    y0, y1 = np.clip([cy - half, cy + half], 0.0, height - 1.0)
+    if x1 - x0 < 2.0 or y1 - y0 < 2.0:
+        return None
+    return np.array([x0, y0, x1, y1], dtype=np.float32)
+
+
+def box_centre(box) -> np.ndarray:
+    return np.array([[(box[0] + box[2]) / 2.0, (box[1] + box[3]) / 2.0]], dtype=np.float32)
+
+
+PROMPT_MODES = ("centre", "box", "torso", "upper", "joints", "all")
+DEFAULT_PROMPT = "joints"
+TORSO_PAIRS = ((5, 6), (11, 12))  # COCO shoulders, hips
+EXTREMITIES = (0, 9, 10, 15, 16)  # COCO nose, wrists, ankles
+UPPER = (0, 9, 10)  # COCO nose, wrists
+ANKLES = (15, 16)
+# Floor-standing props hide the feet first; a hidden ankle still scores 0.46-0.71
+# (fal stepup/handon) against >= 0.76 when visible, and a click on the prop pulls
+# the prop into the person mask.
+ANKLE_TRUST = 0.75
+
+
+def prompt_points(mode: str, box, kp=None, trust: float = TRUST_CONF):
+    """Positive points accompanying the box prompt of one frame (None = box only).
+
+    centre: the box centre; torso: shoulder and hip midpoints whose two ViTPose
+    joints are both confident and inside the box.
+    """
+    if mode not in PROMPT_MODES:
+        raise ValueError(f"unknown prompt mode {mode}")
+    if box is None or mode == "box":
+        return None
+    if mode == "centre":
+        return box_centre(box)
+    if kp is None:
+        return None
+    kp = np.asarray(kp, dtype=np.float64)
+    candidates = []
+    for a, b in TORSO_PAIRS:
+        if kp[a, 2] >= trust and kp[b, 2] >= trust:
+            candidates.append((kp[a, :2] + kp[b, :2]) / 2.0)
+    extra = {"upper": UPPER, "joints": EXTREMITIES, "all": range(len(kp))}.get(mode, ())
+    candidates += [kp[j, :2] for j in extra if kp[j, 2] >= (max(trust, ANKLE_TRUST) if j in ANKLES else trust)]
+    points = [(x, y) for x, y in candidates if box[0] <= x <= box[2] and box[1] <= y <= box[3]]
+    return np.array(points, dtype=np.float32) if points else None
+
+
+def frame_confidence(kp2d, frames: int) -> np.ndarray:
+    """Mean ViTPose confidence per frame; all ones when the obs has no kp2d."""
+    if kp2d is None:
+        return np.ones(frames, dtype=np.float64)
+    kp2d = np.asarray(kp2d, dtype=np.float64)
+    if kp2d.ndim != 3 or kp2d.shape[0] != frames or kp2d.shape[2] < 3:
+        raise ValueError(f"bad-obs: kp2d shape {kp2d.shape} for {frames} frames")
+    conf = kp2d[:, :, 2].mean(axis=1)
+    return np.where(np.isfinite(conf), conf, 0.0)
+
+
+def prompt_boxes(bbx_xys, width: int, height: int):
+    """Validated per-frame prompt boxes; raises no-person-prompt when none exists."""
+    if bbx_xys is None:
+        raise NoPersonPrompt("obs has no bbx_xys")
+    bbx = np.asarray(bbx_xys, dtype=np.float64)
+    if bbx.ndim != 2 or bbx.shape[0] == 0 or bbx.shape[1] != 3:
+        raise NoPersonPrompt(f"obs bbx_xys is empty or malformed (shape {bbx.shape})")
+    boxes = [box_from_xys(row, width, height) for row in bbx]
+    if all(box is None for box in boxes):
+        raise NoPersonPrompt(f"none of {len(boxes)} obs boxes is a valid person box")
+    return boxes
+
+
+def choose_anchor(boxes, conf, trust: float = TRUST_CONF) -> int:
+    valid = np.array([box is not None for box in boxes])
+    if not valid.any():
+        raise NoPersonPrompt("no valid obs box")
+    trusted = np.flatnonzero(valid & (np.asarray(conf) >= trust))
+    if trusted.size:
+        return int(trusted[0])
+    return int(np.flatnonzero(valid)[np.argmax(np.asarray(conf)[valid])])
+
+
+def should_reprompt(prev_area: int, area: int, threshold: float = REPROMPT_CHANGE) -> bool:
+    """Re-prompt when the mask area jumps by more than ``threshold`` of the previous frame's."""
+    return prev_area > 0 and abs(area - prev_area) > threshold * prev_area
+
+
+def reprompt_frames(areas, eligible=None, threshold: float = REPROMPT_CHANGE):
+    """Frames the rule fires on for a fixed area sequence (tracking order = index order)."""
+    areas = [int(a) for a in areas]
+    eligible = [True] * len(areas) if eligible is None else list(eligible)
+    return [t for t in range(1, len(areas)) if eligible[t] and should_reprompt(areas[t - 1], areas[t], threshold)]
+
+
+def plan_chunks(frames: int, anchor: int, chunk_frames: int):
+    """(direction, start, end) inclusive ranges in processing order.
+
+    Forward chunks walk anchor..T-1, backward chunks anchor..0; neighbouring
+    chunks share their boundary frame, which seeds the later chunk.
+    """
+    if chunk_frames < 2:
+        raise ValueError("chunk-frames must be >= 2")
+    if not 0 <= anchor < frames:
+        raise ValueError(f"anchor {anchor} outside 0..{frames - 1}")
+    chunks = []
+    start = anchor
+    while True:
+        end = min(start + chunk_frames - 1, frames - 1)
+        chunks.append((1, start, end))
+        if end == frames - 1:
+            break
+        start = end
+    end = anchor
+    while end > 0:
+        start = max(end - chunk_frames + 1, 0)
+        chunks.append((-1, start, end))
+        end = start
+    return chunks
+
+
+def read_video(path: str):
+    import cv2
+
+    capture = cv2.VideoCapture(path)
+    if not capture.isOpened():
+        raise RuntimeError(f"cannot open video {path}")
+    frames = []
+    while True:
+        ok, frame = capture.read()
+        if not ok:
+            break
+        frames.append(frame)
+    capture.release()
+    if not frames:
+        raise RuntimeError(f"no frames decoded from {path}")
+    return frames
+
+
+def log(message: str) -> None:
+    print(f"[masks] {message}", flush=True)
+
+
+def track_clip(predictor, open_chunk, count: int, prob_hw, boxes, trusted, anchor: int, chunks,
+               threshold: float = REPROMPT_CHANGE, points=None):
+    """Propagate one object over ``chunks`` with the online re-prompt rule.
+
+    ``predictor`` follows the SAM2 video predictor API; ``open_chunk(start, end)``
+    is a context manager yielding an inference state over frames start..end.
+    Returns (prob float16 T x h x w, area int64 T, reprompted frame list).
+    """
+    import torch
+    import torch.nn.functional as F
+
+    ph, pw = prob_hw
+    prob = np.zeros((count, ph, pw), dtype=np.float16)
+    area = np.full(count, -1, dtype=np.int64)
+    reprompted: list[int] = []
+    boundary: dict = {}  # full-res bool masks: chunk boundary frames and the anchor
+    done = 0
+
+    def record(t: int, logits) -> int:
+        nonlocal done
+        mask = logits[0, 0] > 0
+        small = F.interpolate(torch.sigmoid(logits.float()), size=(ph, pw), mode="area")[0, 0]
+        prob[t] = small.float().cpu().numpy().astype(np.float16)
+        if area[t] < 0:
+            done += 1
+            if done % 20 == 0 or done == count:
+                log(f"{done} / {count}")
+        area[t] = int(mask.sum().item())
+        boundary[t] = mask.cpu()
+        return int(area[t])
+
+    def prompt_box(state, local: int, t: int):
+        positive = None if points is None else points[t]
+        if positive is None:
+            predictor.add_new_points_or_box(state, frame_idx=local, obj_id=OBJ_ID, box=boxes[t])
+        else:
+            predictor.add_new_points_or_box(state, frame_idx=local, obj_id=OBJ_ID, box=boxes[t], points=positive,
+                                            labels=np.ones(len(positive), dtype=np.int32))
+
+    for index, (direction, start, end) in enumerate(chunks):
+        reverse = direction < 0
+        seed = end if reverse else start
+        log(f"chunk {index + 1}/{len(chunks)} {'backward' if reverse else 'forward'} {start}-{end} seed={seed}")
+        with open_chunk(start, end) as state:
+            seed_local = seed - start
+            if seed in boundary and bool(boundary[seed].any()):
+                predictor.add_new_mask(state, frame_idx=seed_local, obj_id=OBJ_ID, mask=boundary[seed])
+            elif boxes[seed] is not None:
+                if seed != anchor:
+                    reprompted.append(seed)
+                    log(f"reprompt f{seed}: empty seed mask")
+                prompt_box(state, seed_local, seed)
+            else:
+                predictor.add_new_mask(state, frame_idx=seed_local, obj_id=OBJ_ID, mask=boundary[seed])
+            resume = seed_local
+            while True:
+                restart = None
+                for local, _, logits in predictor.propagate_in_video(state, start_frame_idx=resume, reverse=reverse):
+                    t = start + local
+                    if t == seed and area[t] >= 0:
+                        continue  # boundary frame already recorded by the previous chunk
+                    new_area = record(t, logits)
+                    prev = t + 1 if reverse else t - 1
+                    if (t != seed and 0 <= prev < count and area[prev] >= 0 and t not in reprompted
+                            and trusted[t] and boxes[t] is not None
+                            and should_reprompt(int(area[prev]), new_area, threshold)):
+                        log(f"reprompt f{t}: area {int(area[prev])} -> {new_area}")
+                        reprompted.append(t)
+                        restart = local
+                        break
+                if restart is None:
+                    break
+                # Forget the tracked output so the obs box is an initial conditioning
+                # prompt (fresh segmentation), not a correction of the drifted mask.
+                state["output_dict_per_obj"][0]["non_cond_frame_outputs"].pop(restart, None)
+                state["frames_tracked_per_obj"][0].pop(restart, None)
+                area[start + restart] = -1
+                done -= 1
+                prompt_box(state, restart, start + restart)
+                resume = restart
+        for t in list(boundary):
+            if t not in (start, end, anchor):
+                del boundary[t]
+    if (area < 0).any():
+        raise RuntimeError(f"frames without a mask: {np.flatnonzero(area < 0).tolist()[:10]}")
+    return prob, area, sorted(reprompted)
+
+
+def segment(video: str, obs: str, chunk_frames: int = DEFAULT_CHUNK_FRAMES, checkpoint: str = CHECKPOINT,
+            max_reserved_mib: int = DEFAULT_MAX_RESERVED_MIB, threshold: float = REPROMPT_CHANGE,
+            prompt: str = DEFAULT_PROMPT):
+    """Run SAM2 over the clip; returns (arrays, summary)."""
+    if prompt not in PROMPT_MODES:
+        raise ValueError(f"unknown prompt mode {prompt}")
+    data = np.load(obs)
+    bbx = data["bbx_xys"] if "bbx_xys" in data.files else None
+    kp2d = data["kp2d"] if "kp2d" in data.files else None
+    # Validate the prompt before decoding video or touching the GPU.
+    if bbx is None or np.asarray(bbx).ndim != 2 or np.asarray(bbx).shape[0] == 0:
+        prompt_boxes(bbx, 1, 1)
+    started = time.time()
+    frames = read_video(video)
+    count = len(frames)
+    height, width = frames[0].shape[:2]
+    if len(bbx) != count:
+        raise ValueError(f"frame-count-mismatch: video {count} frames, obs {len(bbx)}")
+    boxes = prompt_boxes(bbx, width, height)
+    conf = frame_confidence(kp2d, count)
+    anchor = choose_anchor(boxes, conf)
+    chunks = plan_chunks(count, anchor, chunk_frames)
+    points = [prompt_points(prompt, boxes[t], None if kp2d is None else kp2d[t]) for t in range(count)]
+    anchor_points = None if points[anchor] is None else points[anchor].round(1).tolist()
+    log(f"video {width}x{height} frames={count} anchor={anchor} conf={conf[anchor]:.2f} box={boxes[anchor].round(1).tolist()} "
+        f"prompt={prompt} points={anchor_points} chunks={len(chunks)}")
+
+    import cv2
+    import torch
+    from sam2.build_sam import build_sam2_video_predictor
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is unavailable")
+    total_mib = torch.cuda.get_device_properties(0).total_memory / 2**20
+    # Hard cap: an allocation past the budget raises OOM instead of silently exceeding it.
+    torch.cuda.set_per_process_memory_fraction(min(1.0, max_reserved_mib / total_mib), 0)
+    torch.cuda.reset_peak_memory_stats()
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    predictor = build_sam2_video_predictor(CONFIG, checkpoint, device="cuda")
+    load_s = time.time() - started
+
+    @contextlib.contextmanager
+    def jpeg_chunk(start: int, end: int):
+        with tempfile.TemporaryDirectory(prefix="cozyfit-masks-") as tmp:
+            for local, t in enumerate(range(start, end + 1)):
+                cv2.imwrite(os.path.join(tmp, f"{local:05d}.jpg"), frames[t], [cv2.IMWRITE_JPEG_QUALITY, 95])
+            state = predictor.init_state(tmp, offload_video_to_cpu=True)
+            try:
+                yield state
+            finally:
+                del state
+                torch.cuda.empty_cache()
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        prob, area, reprompted = track_clip(predictor, jpeg_chunk, count, (height // 2, width // 2),
+                                            boxes, conf >= TRUST_CONF, anchor, chunks, threshold, points)
+    ph, pw = prob.shape[1:]
+    torch.cuda.synchronize()
+    peak = torch.cuda.max_memory_reserved() / 2**20
+    seconds = time.time() - started
+    arrays = {
+        "prob": prob,
+        "reprompted": np.array(sorted(reprompted), dtype=np.int32),
+        "area": area.astype(np.int32),
+        "anchor": np.array(anchor, dtype=np.int32),
+        "chunks": np.array(chunks, dtype=np.int32).reshape(-1, 3),
+    }
+    summary = {
+        "frames": count, "width": width, "height": height, "probShape": [count, ph, pw],
+        "anchor": anchor, "anchorConfidence": float(conf[anchor]), "chunks": [list(c) for c in chunks],
+        "chunkFrames": chunk_frames, "reprompted": sorted(reprompted), "repromptChange": threshold,
+        "trustConfidence": TRUST_CONF, "prompt": prompt, "emptyFrames": int((area == 0).sum()),
+        "seconds": round(seconds, 3), "modelLoadSeconds": round(load_s, 3), "peakReservedMiB": round(peak, 1),
+        "maxReservedMiB": max_reserved_mib, "config": CONFIG, "checkpoint": os.path.basename(checkpoint),
+        "torch": torch.__version__,
+    }
+    return arrays, summary
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--video", required=True)
+    parser.add_argument("--obs", required=True)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--chunk-frames", type=int, default=DEFAULT_CHUNK_FRAMES)
+    parser.add_argument("--checkpoint", default=CHECKPOINT)
+    parser.add_argument("--max-reserved-mib", type=int, default=DEFAULT_MAX_RESERVED_MIB)
+    parser.add_argument("--prompt", choices=PROMPT_MODES, default=DEFAULT_PROMPT)
+    args = parser.parse_args(argv)
+    try:
+        arrays, summary = segment(args.video, args.obs, args.chunk_frames, args.checkpoint, args.max_reserved_mib,
+                                  prompt=args.prompt)
+    except NoPersonPrompt as error:
+        print(str(error), file=sys.stderr, flush=True)
+        return 3
+    os.makedirs(args.out, exist_ok=True)
+    np.savez_compressed(os.path.join(args.out, "masks.npz"), **arrays)
+    with open(os.path.join(args.out, "masks.json"), "w") as handle:
+        json.dump(summary, handle, indent=1)
+    log(f"done frames={summary['frames']} seconds={summary['seconds']} peakReservedMiB={summary['peakReservedMiB']} "
+        f"reprompted={len(summary['reprompted'])} emptyFrames={summary['emptyFrames']}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())