pipeline/gpu.py

"""
The ONLY module that touches the GPU. Models are lazily loaded so each stage
pays only for what it needs, and only the detector you actually pick is loaded:

  detect_frame(..., backend) -> ViTDet-H Cascade Mask R-CNN  OR  RF-DETR-Seg.
  reconstruct_selected()     -> SAM 3D Body: meshes for ONLY the selected boxes.

Both detectors return boxes + per-player masks, so the UI can show/select by box
or by segment with either one. Build still reconstructs only the selected players.
"""

import os
import sys
import functools

import numpy as np

# The sam-3d-body repo is cloned here by the Dockerfile and added to sys.path.
SAM3D_DIR = os.environ.get("SAM3D_DIR", "/app/sam-3d-body")
if SAM3D_DIR not in sys.path:
    sys.path.insert(0, SAM3D_DIR)

HF_REPO_ID = os.environ.get("SAM3D_REPO_ID", "facebook/sam-3d-body-dinov3")
RFDETR_SIZE = os.environ.get("RFDETR_SIZE", "large").lower()

VITDET_PERSON = 0   # detectron2 COCO is 0-indexed
RFDETR_PERSON = 1   # rfdetr/COCO is 1-indexed

_DETECTORS = {}     # backend -> loaded model
_ESTIMATOR = None
_FACES = None


# ----------------------------------------------------------------------------
# Detectors (loaded on first use of that backend).
# ----------------------------------------------------------------------------
def _get_vitdet():
    if "vitdet" not in _DETECTORS:
        import torch
        from tools.build_detector import load_detectron2_vitdet
        d = load_detectron2_vitdet()
        _DETECTORS["vitdet"] = d.to("cuda").eval() if torch.cuda.is_available() else d.eval()
    return _DETECTORS["vitdet"]


def _get_rfdetr():
    if "rfdetr" not in _DETECTORS:
        from rfdetr import (RFDETRSegNano, RFDETRSegSmall,        # type: ignore
                            RFDETRSegMedium, RFDETRSegLarge)
        sizes = {"nano": RFDETRSegNano, "small": RFDETRSegSmall,
                 "medium": RFDETRSegMedium, "large": RFDETRSegLarge}
        m = sizes.get(RFDETR_SIZE, RFDETRSegLarge)()
        try:
            m.optimize_for_inference()
        except Exception:
            pass
        _DETECTORS["rfdetr"] = m
    return _DETECTORS["rfdetr"]


def _detect_vitdet(frame_rgb, conf):
    import torch
    import detectron2.data.transforms as T
    det = _get_vitdet()
    img_bgr = np.ascontiguousarray(frame_rgb[:, :, ::-1])
    h, w = img_bgr.shape[:2]
    aug = T.ResizeShortestEdge(short_edge_length=1024, max_size=1024)
    img_t = aug(T.AugInput(img_bgr)).apply_image(img_bgr)
    img_t = torch.as_tensor(img_t.astype("float32").transpose(2, 0, 1))
    with torch.no_grad():
        out = det([{"image": img_t, "height": h, "width": w}])
    inst = out[0]["instances"].to("cpu")
    boxes = inst.pred_boxes.tensor.numpy()
    classes = inst.pred_classes.numpy()
    scores = inst.scores.numpy()
    masks = inst.pred_masks.numpy() if inst.has("pred_masks") else None
    people = []
    for k in range(len(boxes)):
        if int(classes[k]) != VITDET_PERSON or scores[k] < conf:
            continue
        m = None
        if masks is not None:
            mk = np.asarray(masks[k]).astype(bool)
            m = mk[0] if mk.ndim == 3 else mk
        people.append({"bbox": boxes[k].astype(float),
                       "score": float(scores[k]), "mask": m})
    return people


def _detect_rfdetr(frame_rgb, conf):
    det = _get_rfdetr().predict(frame_rgb, threshold=conf)
    masks = getattr(det, "mask", None)
    people = []
    for k in range(len(det.xyxy)):
        if int(det.class_id[k]) != RFDETR_PERSON:
            continue
        x1, y1, x2, y2 = [float(v) for v in det.xyxy[k]]
        people.append({"bbox": np.array([x1, y1, x2, y2], dtype=float),
                       "score": float(det.confidence[k]),
                       "mask": (np.asarray(masks[k], dtype=bool) if masks is not None else None)})
    return people


@functools.lru_cache(maxsize=16)
def _detect_cached(video_path, idx, conf, backend):
    """Run the chosen detector on one frame; person boxes + masks. Cached per key."""
    from .video import grab_frame
    frame_rgb = grab_frame(video_path, idx)
    if frame_rgb is None:
        return None
    people = (_detect_rfdetr if backend == "rfdetr" else _detect_vitdet)(frame_rgb, conf)
    people.sort(key=lambda p: (p["bbox"][0], p["bbox"][1]))   # stable L->R order
    return people


def detect_frame(video_path, idx, conf=0.3, backend="vitdet"):
    """CPU-cheap wrapper around the cached detection for the chosen backend."""
    return _detect_cached(str(video_path), int(idx), round(float(conf), 3), backend)


# ----------------------------------------------------------------------------
# Reconstructor — SAM 3D Body. Loaded on the first Build.
# ----------------------------------------------------------------------------
def get_estimator():
    """Lazy-load the SAM 3D Body estimator once; returns (estimator, faces)."""
    global _ESTIMATOR, _FACES
    if _ESTIMATOR is None:
        from huggingface_hub import login
        token = os.environ.get("HF_TOKEN")
        if token:
            login(token=token)
        from notebook.utils import setup_sam_3d_body
        _ESTIMATOR = setup_sam_3d_body(hf_repo_id=HF_REPO_ID)
        _FACES = np.asarray(_ESTIMATOR.faces)
    return _ESTIMATOR, _FACES


def get_faces():
    return get_estimator()[1]


@functools.lru_cache(maxsize=16)
def _reconstruct_cached(video_path, idx, boxes_key):
    """Reconstruct ONLY the given boxes. boxes_key is a hashable tuple of int xyxy."""
    from .video import grab_frame
    est, _ = get_estimator()
    frame_rgb = grab_frame(video_path, idx)
    if frame_rgb is None:
        return None
    boxes = np.array(boxes_key, dtype=np.float32).reshape(-1, 4)
    # Providing bboxes reconstructs exactly these people (in order); the FOV
    # estimator still runs for focal_length.
    people = est.process_one_image(frame_rgb, bboxes=boxes)
    slim = []
    for p in people:
        kp = p.get("pred_keypoints_3d")
        slim.append({
            "bbox": np.asarray(p["bbox"]).reshape(-1)[:4].astype(float),
            "pred_vertices": np.asarray(p["pred_vertices"], dtype=np.float32),
            "pred_cam_t": np.asarray(p["pred_cam_t"], dtype=np.float32).reshape(3),
            "focal_length": float(np.asarray(p["focal_length"]).reshape(-1)[0]),
            "pred_keypoints_3d": (None if kp is None
                                  else np.asarray(kp, dtype=np.float32)),
        })
    return slim


def reconstruct_selected(video_path, idx, boxes):
    """Reconstruct meshes for the selected boxes (list of [x1,y1,x2,y2])."""
    boxes_key = tuple(int(round(v)) for b in boxes for v in b[:4])
    return _reconstruct_cached(str(video_path), int(idx), boxes_key)