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 ONNX Runtime providers, using CUDA only on a usable NVIDIA GPU."""
    try:
        import torch
    except ImportError:
        return ["CPUExecutionProvider"]
    if not torch.cuda.is_available():
        return ["CPUExecutionProvider"]

    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,
        allowed_modules=["detection", "recognition"],
    )
    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 transcribe(
    audio_path: str | Path,
    model_name: str,
    language: str | None = None,
    beam_size: int = 5,
    vad_filter: bool = True,
    device: str = "auto",
    compute_type: str = "default",
) -> Iterator[TranscriptSegment]:
    """Transcribe a whole recording, yielding a result per Whisper segment.

    ``language`` is an ISO 639-1 code such as ``"de"``; ``None`` detects it from
    the first 30 seconds. ``vad_filter`` skips silent stretches, which both
    speeds the pass up and suppresses the text Whisper otherwise invents to fill
    silence. Segments are yielded as they are decoded, so callers can show
    progress; stop iterating to cancel.

    ``device`` is ``"auto"``, ``"cpu"`` or ``"cuda"``; ``"auto"`` uses a GPU when
    the machine has one. CTranslate2 loads its CUDA libraries by name the first
    time it uses a GPU, so a machine whose CUDA does not match the one it was
    built against fails here rather than transcribing more slowly on the CPU.
    """
    if _is_crisper_whisper(model_name):
        yield from _transcribe_crisper(
            audio_path, model_name, language, device, compute_type
        )
        return

    from faster_whisper import WhisperModel

    model = WhisperModel(
        resolve_whisper_model(model_name),
        device=device,
        compute_type=compute_type,
        download_root=str(model_cache_dir()),
    )
    segments, _ = model.transcribe(
        load_audio(audio_path),
        language=language,
        beam_size=beam_size,
        vad_filter=vad_filter,
        word_timestamps=True,
    )
    for segment in segments:
        words = [
            Word(float(w.start), float(w.end), w.word, float(w.probability))
            for w in (segment.words or [])
        ]
        yield TranscriptSegment(
            float(segment.start), float(segment.end), segment.text.strip(), words
        )

def transcript_to_dataframes(
    transcript: Iterable[TranscriptSegment],
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Stack a transcript into a segments table and a per-word table.

    The segments are numbered by ``segment_id`` in the order they were spoken,
    with :data:`SEGMENT_COLUMNS`; each word carries the ``segment_id`` it came
    from and its position within it, with :data:`WORD_COLUMNS`.
    """
    ordered = sorted(transcript, key=lambda s: (s.start, s.end))
    segment_rows = [(index, s.start, s.end, s.text) for index, s in enumerate(ordered)]
    word_rows = [
        (index, position, w.start, w.end, w.word.strip(), w.score)
        for index, segment in enumerate(ordered)
        for position, w in enumerate(segment.words)
    ]

    segments = pd.DataFrame(segment_rows, columns=SEGMENT_COLUMNS)
    words = pd.DataFrame(word_rows, columns=WORD_COLUMNS)
    return (
        segments.astype({"segment_id": int, "start": float, "end": float, "text": str}),
        words.astype(
            {
                "segment_id": int,
                "word_index": int,
                "start": float,
                "end": float,
                "word": str,
                "score": float,
            }
        ),
    )

def measure_loudness(path: str | Path) -> pd.DataFrame:
    """How loud a recording is over time, one row every :data:`HOP_SECONDS`.

    ``level_db`` is the chunk's RMS level in dB and ``time`` is its middle, on
    the recording's own clock. A recording too short to fill a single chunk
    measures as no rows at all.
    """
    samples = load_audio(path, SAMPLE_RATE)
    frame = int(HOP_SECONDS * SAMPLE_RATE)
    if samples.size < frame:
        level_db = np.empty(0)
    else:
        frames = samples[: samples.size // frame * frame].reshape(-1, frame)
        rms = np.sqrt((frames.astype(np.float64) ** 2).mean(axis=1))
        level_db = 20 * np.log10(np.maximum(rms, 1e-8))
    return pd.DataFrame(
        {
            "time": np.arange(level_db.size) * HOP_SECONDS + HOP_SECONDS / 2,
            "level_db": level_db,
        }
    )

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"