Pipeline API¶

def boxes_from_tracks(tracks: np.ndarray) -> list[BoundingBox]:
    """Drawable boxes for one frame's BoxMOT ``tracks`` array."""
    return [
        BoundingBox(
            float(row[0]),
            float(row[1]),
            float(row[2]),
            float(row[3]),
            int(row[_ID]),
        )
        for row in np.asarray(tracks)
    ]

def default_device() -> str:
    """Pick the fastest available device as a BoxMOT/Ultralytics device string."""
    try:
        import torch
    except ImportError:
        return "cpu"
    if torch.cuda.is_available():
        return "0"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return "mps"
    return "cpu"

def detect_tracklets(
    video_path: str | Path,
    detector: str = "yolov8n",
    reid: str = "osnet_x0_25_msmt17",
    tracker: str = "botsort",
    device: str | None = None,
    object_classes: Sequence[int] = (0,),
) -> Iterator[FrameResult]:
    """Track objects in a video, yielding BoxMOT's per-frame result.

    Runs BoxMOT (object detection + ReID tracking) and yields each frame's
    :class:`~boxmot.engine.tracking.results.FrameResult` as soon as it is
    computed, so callers can display or accumulate it live. Each result exposes
    ``frame_idx`` (1-based) and ``tracks`` -- an ``(n, 8)`` array of
    ``x1, y1, x2, y2, id, conf, cls, det_ind``. Every frame is yielded, including
    frames with no detections (their ``tracks`` is empty). Stop iterating to
    cancel tracking early.

    ``object_classes`` are the COCO class ids to detect; the default ``(0,)`` is
    people only. ``device`` is a BoxMOT/Ultralytics device string (e.g. ``"0"``,
    ``"mps"``, ``"cpu"``); when ``None`` the fastest available device is chosen
    automatically. The detections sharing a track id form a tracklet, and
    ``conf`` is the per-frame confidence of the object detector.
    """
    # To get lazy, per-frame results so the GUI can display tracking live,
    # we can't use ``boxmot.Boxmot.track`` as this directly writes to a ``runs/`` directory.
    # So we use these internal ``build_*_from_spec`` imports and pin the boxmot version.
    from boxmot import track
    from boxmot.engine.workflows.support import (
        build_detector_from_spec,
        build_tracker_from_spec,
        build_tracker_with_reid_spec,
    )

    if device is None:
        device = default_device()

    # Pass absolute cache paths so both the Ultralytics detector weights and the
    # BoxMOT ReID weights download into the shared cache, rather than the current
    # working directory (detector) or BoxMOT's own in-package models dir (ReID).
    detector_runtime = build_detector_from_spec(
        cached_model_path(detector), classes=list(object_classes), device=device
    )
    tracker_runtime = build_tracker_from_spec(tracker, device=device)
    reid_runtime = build_tracker_with_reid_spec(
        tracker, tracker_runtime, cached_model_path(reid), device=device
    )

    yield from track(
        str(video_path), detector_runtime, reid_runtime, tracker_runtime, verbose=False
    )

def tracks_to_dataframe(frames: Iterable[tuple[int, np.ndarray]]) -> pd.DataFrame:
    """Stack ``(frame_idx, tracks)`` pairs into the stored tracklets DataFrame.

    Returns a :class:`pandas.DataFrame` with columns
    ``frame, track_id, x1, y1, x2, y2, conf``, ``frame`` and ``track_id`` as
    integers -- a compact numeric form to store and analyse. Frames with no
    detections contribute no rows.
    """
    blocks = []
    for frame_idx, tracks in frames:
        tracks = np.asarray(tracks)
        if len(tracks) == 0:
            continue
        block = np.empty((len(tracks), len(_COLUMNS)))
        block[:, 0] = frame_idx
        block[:, 1] = tracks[:, _ID]
        block[:, 2:6] = tracks[:, _XYXY]
        block[:, 6] = tracks[:, _CONF]
        blocks.append(block)
    data = np.vstack(blocks) if blocks else np.empty((0, len(_COLUMNS)))
    return pd.DataFrame(data, columns=_COLUMNS).astype({"frame": int, "track_id": int})

def default_providers() -> list[str]:
    """Pick onnxruntime execution providers, preferring CUDA when available."""
    try:
        import onnxruntime as ort
    except ImportError:
        return ["CPUExecutionProvider"]
    if "CUDAExecutionProvider" in ort.get_available_providers():
        return ["CUDAExecutionProvider", "CPUExecutionProvider"]
    return ["CPUExecutionProvider"]

def detect_faces(
    video_path: str | Path,
    boxes_by_frame: Mapping[int, Sequence[BoundingBox]],
    model_name: str = "buffalo_l",
    det_size: int = 640,
    det_thresh: float = 0.5,
    providers: list[str] | None = None,
) -> Iterator[FaceFrameResult]:
    """Detect one face per tracked person box, yielding a result per frame.

    ``boxes_by_frame`` maps a 0-based frame index to that frame's tracked person
    :class:`BoundingBox`es, as produced by object tracking. For each such frame
    the person crop is run through InsightFace; the highest scoring face above
    ``det_thresh`` is kept, its bounding box and five landmarks offset back into
    full-frame coordinates and tagged with the box's ``track_id``. Every
    requested frame is yielded -- including frames where no face cleared the
    threshold (their ``faces`` is empty) -- so callers can show progress and
    cancel responsively. Stop iterating to cancel early.
    """
    import cv2
    from insightface.app import FaceAnalysis

    if providers is None:
        providers = default_providers()

    root = _insightface_root()
    _ensure_model_available(model_name, root)
    app = FaceAnalysis(name=model_name, providers=providers, root=root)
    ctx_id = 0 if providers[0].startswith("CUDA") else -1
    app.prepare(ctx_id=ctx_id, det_size=(det_size, det_size))

    capture = cv2.VideoCapture(str(video_path))
    if not capture.isOpened():
        raise OSError(f"Could not open video: {video_path}")
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

    try:
        if not boxes_by_frame:
            return
        wanted = set(boxes_by_frame)
        last_frame = max(wanted)
        # Walk frames in order, only decoding (retrieve) the ones we need; this
        # avoids the per-frame seeks that random access would cost.
        frame_idx = -1
        while frame_idx < last_frame:
            if not capture.grab():
                break
            frame_idx += 1
            if frame_idx not in wanted:
                continue
            ok, image = capture.retrieve()
            if not ok:
                break
            faces = _detect_in_boxes(
                app, image, boxes_by_frame[frame_idx], width, height, det_thresh
            )
            yield FaceFrameResult(frame_idx, faces)
    finally:
        capture.release()

def face_box_from_row(row) -> FaceBox:
    """Rebuild a drawable :class:`FaceBox` from a merged DataFrame row."""
    landmarks = [
        (float(getattr(row, f"{name}_x")), float(getattr(row, f"{name}_y")))
        for name in LANDMARK_NAMES
    ]
    box = BoundingBox(
        float(row.face_x1),
        float(row.face_y1),
        float(row.face_x2),
        float(row.face_y2),
        int(row.track_id),
    )
    return FaceBox(box, float(row.face_score), landmarks)

def faces_to_dataframe(frames: Iterable[FaceFrameResult]) -> pd.DataFrame:
    """Stack per-frame face results into a DataFrame keyed on ``(frame, track_id)``.

    Returns a :class:`pandas.DataFrame` with columns ``frame, track_id`` plus
    :data:`FACE_COLUMNS`, ready to left-merge onto the stored tracks. Frames with
    no detected faces contribute no rows.
    """
    rows = []
    for result in frames:
        for face in result.faces:
            row = [
                result.frame_idx,
                face.box.track_id,
                face.score,
                face.box.x1,
                face.box.y1,
                face.box.x2,
                face.box.y2,
            ]
            for px, py in face.landmarks:
                row.extend((px, py))
            rows.append(row)
    data = np.asarray(rows, dtype=float) if rows else np.empty((0, len(_COLUMNS)))
    return pd.DataFrame(data, columns=_COLUMNS).astype({"frame": int, "track_id": int})

def detect_body_poses(
    video_path: str | Path,
    boxes_by_frame: Mapping[int, Sequence[BoundingBox]],
    model_name: str | Path = DEFAULT_MODEL_NAME,
    conf: float = 0.25,
    device: str | None = None,
) -> Iterator[PoseFrameResult]:
    """Detect one body pose per tracked person box, yielding a result per frame.

    ``boxes_by_frame`` maps a 0-based frame index to that frame's tracked person
    :class:`BoundingBox`es. For each requested frame, each person crop is passed
    through an Ultralytics YOLO pose model; the highest scoring pose in the crop
    is kept, its bounding box and COCO keypoints are offset back into full-frame
    coordinates and tagged with the person's ``track_id``. Every requested frame
    is yielded, including frames where no pose cleared ``conf``.
    """
    import cv2
    from ultralytics import YOLO

    if device is None:
        device = default_device()

    model = YOLO(cached_model_path(model_name))

    capture = cv2.VideoCapture(str(video_path))
    if not capture.isOpened():
        raise OSError(f"Could not open video: {video_path}")
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

    try:
        if not boxes_by_frame:
            return
        wanted = set(boxes_by_frame)
        last_frame = max(wanted)
        frame_idx = -1
        while frame_idx < last_frame:
            if not capture.grab():
                break
            frame_idx += 1
            if frame_idx not in wanted:
                continue
            ok, image = capture.retrieve()
            if not ok:
                break
            poses = _detect_in_boxes(
                model, image, boxes_by_frame[frame_idx], width, height, conf, device
            )
            yield PoseFrameResult(frame_idx, poses)
    finally:
        capture.release()

def pose_from_row(row) -> BodyPose:
    """Rebuild a drawable :class:`BodyPose` from a merged DataFrame row."""
    keypoints = [
        (
            float(getattr(row, f"pose_{name}_x")),
            float(getattr(row, f"pose_{name}_y")),
            float(getattr(row, f"pose_{name}_score")),
        )
        for name in KEYPOINT_NAMES
    ]
    box = BoundingBox(
        float(row.pose_x1),
        float(row.pose_y1),
        float(row.pose_x2),
        float(row.pose_y2),
        int(row.track_id),
    )
    return BodyPose(box, float(row.pose_score), keypoints)

def poses_to_dataframe(frames: Iterable[PoseFrameResult]) -> pd.DataFrame:
    """Stack per-frame pose results into a DataFrame keyed on ``(frame, track_id)``.

    Returns a :class:`pandas.DataFrame` with columns ``frame, track_id`` plus
    :data:`POSE_COLUMNS`, ready to left-merge onto the stored tracks. Frames with
    no detected poses contribute no rows.
    """
    rows = []
    for result in frames:
        for pose in result.poses:
            row = [
                result.frame_idx,
                pose.box.track_id,
                pose.score,
                pose.box.x1,
                pose.box.y1,
                pose.box.x2,
                pose.box.y2,
            ]
            for px, py, score in pose.keypoints:
                row.extend((px, py, score))
            rows.append(row)
    data = np.asarray(rows, dtype=float) if rows else np.empty((0, len(_COLUMNS)))
    return pd.DataFrame(data, columns=_COLUMNS).astype({"frame": int, "track_id": int})

def cached_model_path(model_ref: str | Path) -> str:
    """Route a bare weights filename into the shared model cache.

    Given only a bare name (e.g. ``yolo26m.pt``), Ultralytics downloads the
    weights into the current working directory. Rewriting it to an absolute path
    under :func:`model_cache_dir` makes the download land in the shared cache
    instead. An explicit path (one with a directory part) or an already-existing
    file is returned unchanged.
    """
    path = Path(model_ref)
    if path.parent != Path(".") or path.exists():
        return str(model_ref)
    cache_dir = model_cache_dir()
    cache_dir.mkdir(parents=True, exist_ok=True)
    return str(cache_dir / path.name)

def model_cache_dir() -> Path:
    """Directory all downloaded model weights are cached in (cross-platform)."""
    return user_cache_path("body-eye-sync", "SSC") / "models"