diff --git a/data_only_viz/action_head_pub.py b/data_only_viz/action_head_pub.py index ca36bb2..8c8f0d6 100644 --- a/data_only_viz/action_head_pub.py +++ b/data_only_viz/action_head_pub.py @@ -129,6 +129,10 @@ class ActionHeadPublisher(threading.Thread): v3d = p.get("v3d") if v3d is None: continue + # CoreMLArray wraps a numpy array but has no __array__ + # protocol; unwrap via .numpy() before np.asarray. + if hasattr(v3d, "numpy") and not isinstance(v3d, np.ndarray): + v3d = v3d.numpy() v3d_np = np.asarray(v3d, dtype=np.float32) if v3d_np.shape[0] < max(SMPLX_JOINT_ANCHOR_VERTS) + 1: continue diff --git a/data_only_viz/scripts/extract_j3d_offline.py b/data_only_viz/scripts/extract_j3d_offline.py new file mode 100644 index 0000000..27ca9d4 --- /dev/null +++ b/data_only_viz/scripts/extract_j3d_offline.py @@ -0,0 +1,127 @@ +"""Extract j3d (22 SMPL-X joint anchors) from a recorded MP4 using the +Multi-HMR CoreML backend, write per-frame per-person jsonl rows. + +Usage: + uv run python -m data_only_viz.scripts.extract_j3d_offline \ + --session sess03 \ + --video ~/.cache/av-live-action/raw/sess03.mp4 \ + --out ~/.cache/av-live-action/raw/sess03.jsonl +""" +from __future__ import annotations + +import argparse +import json +import logging +from pathlib import Path + +import cv2 +import numpy as np + +from data_only_viz.action_head_pub import SMPLX_JOINT_ANCHOR_VERTS +from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend + +LOG = logging.getLogger("extract_j3d_offline") +IMG_SIZE = 672 +DEFAULT_OUT_DIR = Path("~/.cache/av-live-action/raw").expanduser() + + +def _default_K(size: int = IMG_SIZE) -> np.ndarray: + """Synthetic camera intrinsics, focal ~ image size, principal point centred.""" + f = float(size) + cx = cy = f * 0.5 + return np.array( + [[f, 0.0, cx], [0.0, f, cy], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + + +def _frame_to_chw(frame_bgr: np.ndarray, size: int = IMG_SIZE) -> np.ndarray: + """BGR uint8 (H, W, 3) -> float32 CHW (3, size, size) in [0, 1].""" + h, w = frame_bgr.shape[:2] + side = min(h, w) + y0 = (h - side) // 2 + x0 = (w - side) // 2 + crop = frame_bgr[y0:y0 + side, x0:x0 + side] + resized = cv2.resize(crop, (size, size)) + rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + return rgb.transpose(2, 0, 1) # CHW + + +def _person_to_j3d22(person: dict, anchors: tuple[int, ...]) -> np.ndarray | None: + v3d = person.get("v3d") + if v3d is None: + return None + # CoreMLArray wraps numpy but lacks __array__; unwrap before asarray. + if hasattr(v3d, "numpy") and not isinstance(v3d, np.ndarray): + v3d = v3d.numpy() + v3d_np = np.asarray(v3d, dtype=np.float32) + if v3d_np.shape[0] < max(anchors) + 1: + return None + return v3d_np[list(anchors)].astype(np.float32) + + +def extract(session: str, video: Path, out: Path, + det_thresh: float = 0.3, + mlpackage_path: Path | None = None, + anchors: tuple[int, ...] = SMPLX_JOINT_ANCHOR_VERTS) -> int: + """Returns the number of (frame, person) rows written.""" + out.parent.mkdir(parents=True, exist_ok=True) + cap = cv2.VideoCapture(str(video)) + if not cap.isOpened(): + raise RuntimeError(f"cannot open {video}") + fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + backend = MultiHMRCoreMLBackend(mlpackage_path) if mlpackage_path \ + else MultiHMRCoreMLBackend() + K = _default_K(IMG_SIZE) + n_frames = 0 + n_rows = 0 + with out.open("w") as f: + while True: + ok, frame = cap.read() + if not ok: + break + chw = _frame_to_chw(frame) + try: + persons = backend.infer(chw, K, det_thresh=det_thresh) + except Exception: + LOG.exception("infer failed at frame=%d", n_frames) + n_frames += 1 + continue + ts = n_frames / fps + for i, person in enumerate(persons): + j3d = _person_to_j3d22(person, anchors) + if j3d is None: + continue + f.write(json.dumps({ + "ts": ts, + "session": session, + "pid": int(person.get("pid", i)), + "j3d": j3d.tolist(), + }) + "\n") + n_rows += 1 + n_frames += 1 + if n_frames % 100 == 0: + LOG.info("frame=%d rows=%d", n_frames, n_rows) + cap.release() + LOG.info("done: %d frames, %d rows -> %s", n_frames, n_rows, out) + return n_rows + + +def _cli() -> None: + p = argparse.ArgumentParser() + p.add_argument("--session", required=True) + p.add_argument("--video", required=True, type=Path) + p.add_argument("--out", type=Path) + p.add_argument("--det-thresh", type=float, default=0.3) + p.add_argument("--mlpackage", type=Path, default=None) + args = p.parse_args() + logging.basicConfig(level=logging.INFO, + format="%(asctime)s [%(name)s] %(message)s") + DEFAULT_OUT_DIR.mkdir(parents=True, exist_ok=True) + out = args.out or (DEFAULT_OUT_DIR / f"{args.session}.jsonl") + extract(args.session, args.video, out, + det_thresh=args.det_thresh, mlpackage_path=args.mlpackage) + + +if __name__ == "__main__": + _cli() diff --git a/data_only_viz/scripts/train_on_studio.sh b/data_only_viz/scripts/train_on_studio.sh index def7bcd..e4ec3fe 100755 --- a/data_only_viz/scripts/train_on_studio.sh +++ b/data_only_viz/scripts/train_on_studio.sh @@ -36,7 +36,17 @@ REMOTE_CKPT="$REMOTE_ROOT/checkpoints" DATASET_FILE="${DATASET_FILE:-$LOCAL_DATASET/dataset.jsonl}" CKPT_NAME="${CKPT_NAME:-action_head.pt}" -TRAIN_ARGS="$*" + +# Quote train args defensively before forwarding through bastion ssh + +# studio ssh (each layer reparses). Reject single quotes — they break +# the single-quoted payload in bastion_ssh and could allow injection. +for a in "$@"; do + if [[ "$a" == *"'"* ]]; then + printf '[train_on_studio] forbidden single quote in arg: %s\n' "$a" >&2 + exit 3 + fi +done +TRAIN_ARGS="$(printf '%q ' "$@")" log() { printf '[train_on_studio] %s\n' "$*" >&2; }