14_07_expt/gpu_pose.py

"""GPU pose extractors that match pipeline.extract_skeleton_from_video's output schema.

Both models below output standard COCO-17 keypoints in the same joint order that
pipeline.py's MediaPipe extractor already remaps to (nose, l/r eye, l/r ear,
l/r shoulder, l/r elbow, l/r wrist, l/r hip, l/r knee, l/r ankle), so no joint
reordering is needed -- only pixel -> [0, 1] normalization (x / width, y / height)
to match what preprocessing.correct_aspect_ratio expects.
"""

from __future__ import annotations

import os
import pickle
import sys
from pathlib import Path
from typing import Iterable

import numpy as np


def _ensure_onnxruntime_cuda_libs() -> None:
    """Re-exec the process with cuDNN/cuBLAS dirs on LD_LIBRARY_PATH.

    onnxruntime-gpu's CUDA execution provider dlopen()s libcudnn/libcublasLt at
    runtime. torch's pip-installed nvidia/*/lib wheels already contain these
    .so files but aren't on the loader path. The dynamic linker only reads
    LD_LIBRARY_PATH at process start, so mutating os.environ after the fact
    (e.g. in Python before importing onnxruntime) has no effect -- the process
    must be restarted with the variable already set.
    """
    if os.environ.get("_GPU_POSE_LIBS_PATCHED"):
        return
    if not (sys.argv and os.path.isfile(sys.argv[0])):
        return  # not a real script (e.g. -c/-m); can't safely re-exec

    try:
        import nvidia  # noqa: F401
    except ImportError:
        return

    nvidia_root = Path(nvidia.__file__).resolve().parent
    lib_dirs = [
        str(nvidia_root / name / "lib")
        for name in ("cudnn", "cublas", "cuda_runtime", "cuda_nvrtc")
    ]
    lib_dirs = [p for p in lib_dirs if os.path.isdir(p)]
    if not lib_dirs:
        return

    existing = os.environ.get("LD_LIBRARY_PATH", "")
    if all(d in existing for d in lib_dirs):
        return

    os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([existing] if existing else []))
    os.environ["_GPU_POSE_LIBS_PATCHED"] = "1"
    os.execv(sys.executable, [sys.executable] + sys.argv)


_ensure_onnxruntime_cuda_libs()


def _grab_kept_frames(video_path: Path, frame_stride: int, max_video_frames: int | None):
    """Decode a video, skipping non-kept frames cheaply with grab() instead of a full decode."""
    import cv2

    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        return None

    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 16.0
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 9.0

    frames = []
    decoded_frames = 0
    processed_frames = 0
    try:
        while True:
            if max_video_frames is not None and processed_frames >= max_video_frames:
                break
            ok = cap.grab()
            if not ok:
                break
            decoded_frames += 1
            if (decoded_frames - 1) % frame_stride != 0:
                continue
            ok, frame = cap.retrieve()
            if not ok:
                break
            frames.append(frame)
            processed_frames += 1
    finally:
        cap.release()

    return {
        "frames": frames,
        "fps": float(fps),
        "width": float(width),
        "height": float(height),
        "decoded_frames": decoded_frames,
        "processed_frames": processed_frames,
    }


def _yolo_cache_path(
    cache_dir: Path,
    video_path: Path,
    weights: str,
    imgsz: int,
    half: bool,
    frame_stride: int,
    max_video_frames: int | None,
) -> Path:
    weights_tag = Path(weights).stem
    cap_tag = max_video_frames if max_video_frames is not None else "all"
    cache_suffix = f"yolo_{weights_tag}_i{imgsz}_h{int(half)}_s{frame_stride}_m{cap_tag}"
    return cache_dir / f"{video_path.stem}_{cache_suffix}.pkl"


def _load_cached_sample(cache_path: Path) -> dict | None:
    if not cache_path.exists():
        return None
    with cache_path.open("rb") as f:
        return pickle.load(f)


def _save_cached_sample(sample: dict, cache_path: Path) -> None:
    """Write via a temp file + atomic rename so a killed process never leaves a torn pickle
    behind -- resumable runs depend on every file at cache_path being either absent or valid."""
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp")
    with tmp_path.open("wb") as f:
        pickle.dump(sample, f)
    os.replace(tmp_path, cache_path)


def load_yolo_model(weights: str = "yolo11n-pose.pt", device: str = "cuda:0"):
    from ultralytics import YOLO

    model = YOLO(weights)
    model.to(device)
    return model


def _predict_with_oom_retry(model, chunk, imgsz: int, half: bool) -> list:
    """Run model.predict on a frame chunk, degrading gracefully on CUDA OOM.

    Even with a capped batch, a worker whose CUDA reservation has ratcheted up (long videos +
    PyTorch's non-releasing caching allocator) can still OOM on a chunk. Rather than dropping the
    whole video, empty the cache and retry the chunk one frame at a time; any single frame that
    still OOMs contributes a None (treated downstream as a no-detection frame) instead of killing
    the sample. This is issue #6 from the review -- production must not silently lose videos to OOM.
    """
    import torch

    def _is_oom(exc: Exception) -> bool:
        return isinstance(exc, torch.cuda.OutOfMemoryError) or (
            isinstance(exc, RuntimeError) and "out of memory" in str(exc).lower()
        )

    try:
        return list(model.predict(chunk, device=0, verbose=False, imgsz=imgsz, half=half, batch=len(chunk)))
    except Exception as exc:  # noqa: BLE001 -- re-raised below unless it's an OOM
        if not _is_oom(exc):
            raise
        torch.cuda.empty_cache()

    results = []
    for frame in chunk:
        try:
            results.extend(model.predict([frame], device=0, verbose=False, imgsz=imgsz, half=half, batch=1))
        except Exception as exc:  # noqa: BLE001
            if not _is_oom(exc):
                raise
            torch.cuda.empty_cache()
            results.append(None)  # no-detection sentinel for this frame
    return results


def _results_to_skeleton(results: list, width: float, height: float) -> np.ndarray:
    """Convert Ultralytics results (best detection per frame) to a normalized (F, 17, 3) array."""
    joints_per_frame = []
    for result in results:
        if result is None or result.keypoints is None or result.boxes is None or len(result.boxes) == 0:
            joints_per_frame.append(np.zeros((17, 3), dtype=np.float64))
            continue
        best_idx = int(result.boxes.conf.argmax().item())
        kpts = result.keypoints.data[best_idx].detach().cpu().numpy()  # (17, 3) pixel x, y, conf
        kpts = kpts.astype(np.float64)
        kpts[:, 0] /= width
        kpts[:, 1] /= height
        joints_per_frame.append(kpts)
    return np.asarray(joints_per_frame, dtype=np.float64)


def _decoded_to_sample_yolo(
    decoded: dict,
    video_path: Path,
    model,
    imgsz: int,
    half: bool,
    max_batch: int,
    frame_stride: int,
    max_video_frames: int | None,
    weights: str,
) -> dict:
    """Run YOLO on already-decoded frames and build a pipeline-compatible sample dict."""
    frames = decoded["frames"]
    results = []
    for i in range(0, len(frames), max_batch):
        results.extend(_predict_with_oom_retry(model, frames[i : i + max_batch], imgsz, half))

    width, height = decoded["width"], decoded["height"]
    return {
        "skeleton": _results_to_skeleton(results, width, height),
        "fps": decoded["fps"],
        "width": width,
        "height": height,
        "source_file": str(video_path),
        "extract_every_n_frames": frame_stride,
        "max_video_frames": max_video_frames,
        "decoded_frames": decoded["decoded_frames"],
        "processed_frames": decoded["processed_frames"],
        "pose_estimator": "yolo_pose_gpu",
        "yolo_weights": str(weights),
    }


def extract_skeleton_from_video_yolo(
    video_path: Path,
    config: dict,
    model,
    imgsz: int = 384,
    max_batch: int = 64,
) -> dict | None:
    """GPU pose extraction with Ultralytics YOLO-pose, batched over one video's frames.

    max_batch caps frames per predict() call. Under multiprocess deployment, PyTorch's CUDA
    caching allocator never releases reserved memory between calls, so per-worker reservations
    ratchet up unevenly as long (near max_video_frames) videos land on some workers and not
    others -- an unbounded batch=len(frames) call (up to max_video_frames, e.g. 160) then
    reliably OOMs at process counts that would otherwise fit in VRAM. Measured: uncapped batches
    with yolo11m-pose at 8 processes / imgsz=640 dropped 92/300 videos to silent CUDA OOM;
    capping at 64 cut that to 2/300. Any residual OOM is absorbed by _predict_with_oom_retry.
    """
    frame_stride = max(1, int(config.get("extract_every_n_frames", 1)))
    max_video_frames = config.get("max_video_frames")
    max_video_frames = None if max_video_frames is None else max(1, int(max_video_frames))
    half = bool(config.get("yolo_half", False))
    weights = config.get("yolo_weights", "")

    cache_dir = config.get("skeleton_cache_dir")
    cache_path = None
    if cache_dir is not None:
        cache_path = _yolo_cache_path(
            Path(cache_dir), video_path, weights, imgsz, half, frame_stride, max_video_frames
        )
        cached = _load_cached_sample(cache_path)
        if cached is not None:
            return cached

    decoded = _grab_kept_frames(video_path, frame_stride, max_video_frames)
    if decoded is None or not decoded["frames"]:
        return None

    sample = _decoded_to_sample_yolo(
        decoded, video_path, model, imgsz, half, max_batch, frame_stride, max_video_frames, weights
    )
    if cache_path is not None:
        _save_cached_sample(sample, cache_path)
    return sample


def extract_all_videos_yolo(video_samples, config: dict):
    """Extract skeletons for many videos: threaded decode feeding one main-thread GPU consumer.

    `video_samples` is an iterable of dicts each with 'video_path' and 'label' (as produced by
    pipeline.discover_video_samples). Yields extracted sample dicts (same schema as
    pipeline.extract_skeleton_from_video) with 'label' attached, ready for process_sample.

    Runs a single CUDA context in this process (safe, simple, one model load). Decode threads
    release the GIL in cv2, so they parallelize; the GPU is fed serially from the calling thread.
    For higher throughput, run several of these processes over disjoint video shards (the
    benchmark's multiprocess pattern) -- the per-process VRAM footprint is ~1.1 GB (nano) /
    ~1.85 GB (small) / ~3.2 GB (medium) at imgsz=640, fp16, max_batch=64.

    Each extracted sample is written to config["skeleton_cache_dir"] immediately (mirroring
    pipeline.extract_skeleton_from_video's on-disk cache) instead of only living in the returned
    generator. A killed/interrupted run can restart and pick up where it left off -- already-cached
    videos are served straight from disk without redoing decode or GPU work.
    """
    from concurrent.futures import ThreadPoolExecutor, as_completed

    weights = config.get("yolo_weights", "yolo11n-pose.pt")
    imgsz = int(config.get("yolo_imgsz", 640))
    half = bool(config.get("yolo_half", True))
    max_batch = int(config.get("yolo_max_batch", 64))
    decode_workers = max(1, int(config.get("yolo_decode_workers", 8)))
    frame_stride = max(1, int(config.get("extract_every_n_frames", 1)))
    max_video_frames = config.get("max_video_frames")
    max_video_frames = None if max_video_frames is None else max(1, int(max_video_frames))
    cache_dir = config.get("skeleton_cache_dir")
    cache_dir = Path(cache_dir) if cache_dir is not None else None

    video_samples = list(video_samples)

    to_extract = []
    for vs in video_samples:
        cache_path = (
            _yolo_cache_path(cache_dir, vs["video_path"], weights, imgsz, half, frame_stride, max_video_frames)
            if cache_dir is not None
            else None
        )
        cached = _load_cached_sample(cache_path) if cache_path is not None else None
        if cached is not None:
            yield {**cached, "label": vs["label"]}
        else:
            to_extract.append((vs, cache_path))

    if not to_extract:
        return

    model = load_yolo_model(weights)

    prefetch = max(decode_workers * 2, 32)
    with ThreadPoolExecutor(max_workers=decode_workers) as ex:
        pending: dict = {}
        cursor = 0

        def _fill():
            nonlocal cursor
            while cursor < len(to_extract) and len(pending) < prefetch:
                vs, cache_path = to_extract[cursor]
                fut = ex.submit(_grab_kept_frames, vs["video_path"], frame_stride, max_video_frames)
                pending[fut] = (vs, cache_path)
                cursor += 1

        _fill()
        while pending:
            done = next(as_completed(pending))
            vs, cache_path = pending.pop(done)
            _fill()
            try:
                decoded = done.result()
            except Exception as exc:  # noqa: BLE001
                print(f"  skipped video {vs.get('source_file', vs['video_path'])}: {exc}")
                continue
            if decoded is None or not decoded["frames"]:
                continue
            sample = _decoded_to_sample_yolo(
                decoded, vs["video_path"], model, imgsz, half, max_batch,
                frame_stride, max_video_frames, weights,
            )
            if cache_path is not None:
                _save_cached_sample(sample, cache_path)
            yield {**sample, "label": vs["label"]}


def load_rtmpose_model(mode: str = "lightweight", device: str = "cuda"):
    from rtmlib import Body

    return Body(mode=mode, to_openpose=False, backend="onnxruntime", device=device)


def extract_skeleton_from_video_rtmpose(
    video_path: Path,
    config: dict,
    body_model,
) -> dict | None:
    """GPU pose extraction with RTMPose (rtmlib), per-frame (detector + pose each call)."""
    frame_stride = max(1, int(config.get("extract_every_n_frames", 1)))
    max_video_frames = config.get("max_video_frames")
    max_video_frames = None if max_video_frames is None else max(1, int(max_video_frames))

    decoded = _grab_kept_frames(video_path, frame_stride, max_video_frames)
    if decoded is None or not decoded["frames"]:
        return None

    width, height = decoded["width"], decoded["height"]
    joints_per_frame = []
    for frame in decoded["frames"]:
        keypoints, scores = body_model(frame)
        if keypoints is None or len(keypoints) == 0:
            joints_per_frame.append(np.zeros((17, 3), dtype=np.float64))
            continue
        best_idx = int(np.argmax(scores.mean(axis=1)))
        kpts = keypoints[best_idx].astype(np.float64)  # (17, 2) pixel x, y
        conf = scores[best_idx].astype(np.float64)  # (17,)
        joint = np.concatenate([kpts, conf[:, None]], axis=1)
        joint[:, 0] /= width
        joint[:, 1] /= height
        joints_per_frame.append(joint)

    return {
        "skeleton": np.asarray(joints_per_frame, dtype=np.float64),
        "fps": decoded["fps"],
        "width": width,
        "height": height,
        "source_file": str(video_path),
        "extract_every_n_frames": frame_stride,
        "max_video_frames": max_video_frames,
        "decoded_frames": decoded["decoded_frames"],
        "processed_frames": decoded["processed_frames"],
        "pose_estimator": "rtmpose_gpu",
    }