Experiment API¶

class BodyPoseStep(_Model):
    """Per-box body-pose detection. Fields mirror ``detect_body_poses``."""

    model_name: str = Field(
        "yolo26m-pose.pt",
        description="Ultralytics YOLO pose weights.",
        json_schema_extra={
            "choices": [
                "yolo26n-pose.pt",
                "yolo26s-pose.pt",
                "yolo26m-pose.pt",
                "yolo26l-pose.pt",
                "yolo26x-pose.pt",
            ]
        },
    )
    conf: float = Field(
        0.25, ge=0.0, le=1.0, description="Minimum pose detection confidence."
    )

class ExperimentConfig(_Model):
    """The serialisable definition of an experiment: its inputs and the pipeline to run."""

    # The pipeline step fields, in run order
    STEP_FIELDS: ClassVar[tuple[str, ...]] = (
        "object_tracking",
        "face_detection",
        "body_pose",
    )

    version: int = CURRENT_VERSION
    name: str
    inputs: list[InputSpec]
    object_tracking: ObjectTrackingStep = Field(default_factory=ObjectTrackingStep)
    face_detection: FaceDetectionStep | None = None
    body_pose: BodyPoseStep | None = None

    @model_validator(mode="after")
    def _check(self) -> ExperimentConfig:
        if not self.inputs:
            raise ValueError("experiment has no inputs")

        ids = [i.id for i in self.inputs]
        duplicates = {i for i in ids if ids.count(i) > 1}
        if duplicates:
            raise ValueError(f"duplicate input ids: {sorted(duplicates)}")
        return self

    @property
    def steps(self) -> list[StepSpec]:
        """The pipeline stages that will run, in order."""
        present = (getattr(self, name) for name in self.STEP_FIELDS)
        return [step for step in present if step is not None]

class FaceDetectionStep(_Model):
    """Per-box face detection. Fields mirror ``detect_faces``."""

    model_name: str = Field(
        "antelopev2",
        description="InsightFace model pack.",
        json_schema_extra={
            "choices": [
                "antelopev2",
                "buffalo_l",
                "buffalo_m",
                "buffalo_s",
                "buffalo_sc",
            ]
        },
    )
    det_size: int = Field(
        640, ge=64, le=2048, description="Detector input size in pixels."
    )
    det_thresh: float = Field(
        0.5, ge=0.0, le=1.0, description="Minimum face detection confidence."
    )
    embeddings_per_track: int = Field(
        32,
        ge=0,
        description=("Number of best face embeddings to keep per tracklet"),
    )

class ObjectTrackingStep(_Model):
    """Object detection + ReID tracking. Fields mirror ``detect_tracklets``.

    ``choices`` in a field's ``json_schema_extra`` are suggested values the GUI
    offers in an editable combobox; a custom value is still allowed.
    """

    detector: str = Field(
        "yolo26m",
        description="Object detector model.",
        json_schema_extra={
            "choices": [
                "yolo26n",
                "yolo26s",
                "yolo26m",
                "yolo26l",
                "yolo26x",
            ]
        },
    )
    reid: str = Field(
        "osnet_x1_0_msmt17",
        description="Re-identification model used to keep track ids stable.",
        json_schema_extra={
            "choices": [
                "osnet_x0_25_msmt17",
                "osnet_x0_5_msmt17",
                "osnet_x0_75_msmt17",
                "osnet_x1_0_msmt17",
                "osnet_ain_x1_0_msmt17",
                "mobilenetv2_x1_0_msmt17",
                "mobilenetv2_x1_4_msmt17",
                "resnet50_msmt17",
                "clip_market1501",
                "clip_duke",
            ]
        },
    )
    tracker: str = Field(
        "botsort",
        description="Multi-object tracking algorithm.",
        json_schema_extra={
            "choices": [
                "botsort",
                "bytetrack",
                "ocsort",
                "deepocsort",
                "hybridsort",
                "strongsort",
                "imprassoc",
                "boosttrack",
            ]
        },
    )
    object_classes: list[int] = Field(
        default=[0],
        description="COCO class ids to detect and track (0 = person).",
    )
    embeddings_per_track: int = Field(
        32,
        ge=0,
        description=(
            "Number of best body-appearance (ReID) embeddings to keep per tracklet"
        ),
    )

class VideoInput(_Model):
    """A single video file to process, referenced by a stable ``id``.

    ``path`` may be relative; it is resolved against the experiment folder by the
    runtime :class:`Experiment`.
    """

    kind: Literal["video"] = "video"
    id: str
    path: Path

class Experiment:
    """An experiment: its :class:`ExperimentConfig` plus the loaded videos."""

    def __init__(self, config: ExperimentConfig, folder: str | Path | None = None):
        self.config = config
        self.folder = Path(folder) if folder is not None else None
        self._videos: dict[str, Video] = {}

    def _require_folder(self) -> Path:
        if self.folder is None:
            raise ValueError("experiment has no folder; load or save it first")
        return self.folder

    def resolved_input_path(self, spec: VideoInput) -> Path:
        """Absolute path for ``spec``; relative paths taken from the folder."""
        if spec.path.is_absolute():
            return spec.path
        return (self._require_folder() / spec.path).resolve()

    @property
    def output_dir(self) -> Path:
        """Where per-input Parquet outputs live, inside the folder."""
        return self._require_folder() / OUTPUTS_DIRNAME

    def output_path(self, spec: VideoInput) -> Path:
        """Parquet output path for ``spec`` under :attr:`output_dir`."""
        return self.output_dir / f"{spec.id}.parquet"

    def video(self, spec: VideoInput) -> Video:
        """The :class:`Video` for ``spec``, created on first access.

        Its cached outputs (data + embeddings) are loaded if they exist.
        """
        video = self._videos.get(spec.id)
        if video is None:
            video = Video()
            video.set_video(self.resolved_input_path(spec))
            if self.folder is not None and self.output_path(spec).exists():
                video.load_parquet(self.output_path(spec))
            self._videos[spec.id] = video
        return video

    @classmethod
    def load(cls, folder: str | Path) -> Experiment:
        """Load the experiment in ``folder`` (its config; videos load on access)."""
        folder = Path(folder)
        with (folder / CONFIG_FILENAME).open("r", encoding="utf-8") as f:
            data = yaml.safe_load(f) or {}
        version = data.get("version", CURRENT_VERSION)
        if version > CURRENT_VERSION:
            raise ValueError(
                f"experiment version {version} is newer than supported "
                f"{CURRENT_VERSION}; please upgrade body-eye-sync"
            )
        return cls(ExperimentConfig.model_validate(data), folder)

    def save(self, folder: str | Path | None = None) -> None:
        """Write the config, and every loaded video's results, into the folder.

        ``folder`` defaults to the current :attr:`folder` and becomes it when
        given. The config is always written; a video is written only once it has
        results.
        """
        if folder is not None:
            self.folder = Path(folder)
        folder = self._require_folder()
        folder.mkdir(parents=True, exist_ok=True)
        with (folder / CONFIG_FILENAME).open("w", encoding="utf-8") as f:
            yaml.safe_dump(self.config.model_dump(mode="json"), f, sort_keys=False)
        for spec in self.config.inputs:
            video = self._videos.get(spec.id)
            if video is not None and video.data is not None:
                self.output_dir.mkdir(parents=True, exist_ok=True)
                video.to_parquet(self.output_path(spec))

class Video:
    """Contains object tracking and vision model outputs for a video.

    Completed results live in a single numeric :attr:`data` DataFrame. While a
    run is in progress, each frame's BoxMOT ``tracks`` array is accumulated and
    collapsed into that DataFrame once :meth:`finish_object_tracking` is called.
    Face detection runs as a later pass over those tracked boxes, accumulating
    per frame and folding its columns onto the matching rows in
    :meth:`finish_face_detection`. Body-pose detection follows the same pattern.
    """

    def __init__(self) -> None:
        # Persistent results.
        self.video_path: Path | None = None
        self._data: pd.DataFrame | None = None
        self._rows_by_frame: dict[int, np.ndarray] = {}
        self._body_embeddings: pd.DataFrame | None = None
        self._face_embeddings: pd.DataFrame | None = None
        # Per-pass scratch: accumulated while a pass runs, then collapsed into the
        # results above and reset. ``_tmp_`` marks them as transient.
        self._tmp_frames: list[tuple[int, np.ndarray]] = []
        self._tmp_face_frames: list[FaceFrameResult] = []
        self._tmp_pose_frames: list[PoseFrameResult] = []
        self._tmp_body_topk = _TrackTopK(0)
        self._tmp_face_topk = _TrackTopK(0)

    def set_video(self, path: str | Path) -> None:
        """Set the current video, invalidating any previous model outputs."""
        self.clear()
        self.video_path = Path(path)

    def begin_object_tracking(self, embeddings_per_track: int = 0) -> None:
        """Drop any previous model outputs.

        ``embeddings_per_track`` keeps that many best body-appearance (ReID)
        embeddings per tracklet, ranked by detection confidence, for later
        identity clustering; ``0`` keeps none.
        """
        self.clear()
        self._tmp_body_topk = _TrackTopK(embeddings_per_track)

    def add_object_tracking_frame(self, frame) -> None:
        """Accumulate a BoxMOT per-frame result, converting to 0-based indices"""
        tracks = np.asarray(frame.tracks)
        self._tmp_frames.append((frame.frame_idx - 1, tracks))
        self._collect_body_embeddings(frame.frame_idx - 1, tracks, frame)

    def _collect_body_embeddings(
        self, frame_index: int, tracks: np.ndarray, frame
    ) -> None:
        """Feed this frame's ReID embeddings (if any) into the per-track top-K."""
        embeddings = getattr(frame, "embeddings", None)
        if embeddings is None:
            return
        embeddings = np.asarray(embeddings)
        for row, vec in zip(tracks, embeddings):
            if not np.any(np.isfinite(vec)):
                continue  # predicted-only track with no detection this frame
            self._tmp_body_topk.add(int(row[4]), frame_index, float(row[5]), vec)

    def finish_object_tracking(self) -> None:
        """Collapse the streamed frames into the stored :attr:`data` DataFrame."""
        self.set_data(tracks_to_dataframe(self._tmp_frames))
        self._body_embeddings = self._tmp_body_topk.to_frame()

    def discard_object_tracking(self) -> None:
        """Drop a cancelled or failed run; its partial output is unusable."""
        self.clear()

    def set_data(self, data: pd.DataFrame) -> None:
        """Replace all results with a complete data DataFrame."""
        self._data = data
        self._rows_by_frame = data.groupby("frame").indices
        self._tmp_frames = []

    def all_boxes_by_frame(self) -> dict[int, list[BoundingBox]]:
        """Tracked person boxes grouped by frame, as later passes consume them."""
        if self._data is None:
            return {}
        return {
            int(frame): self.boxes_for_frame(int(frame))
            for frame in self._rows_by_frame
        }

    def begin_face_detection(self, embeddings_per_track: int = 0) -> None:
        """Drop any previous face columns so a fresh pass starts clean.

        ``embeddings_per_track`` keeps that many best face embeddings per
        tracklet, ranked by face score, for later identity clustering.
        """
        if self._data is not None:
            present = [c for c in FACE_COLUMNS if c in self._data.columns]
            if present:
                self.set_data(self._data.drop(columns=present))
        self._tmp_face_frames = []
        self._tmp_face_topk = _TrackTopK(embeddings_per_track)
        self._face_embeddings = None

    def add_face_detection_frame(self, result: FaceFrameResult) -> None:
        """Accumulate one frame's detected faces for the final merge."""
        self._tmp_face_frames.append(result)
        for face in result.faces:
            self._tmp_face_topk.add(
                face.box.track_id, result.frame_idx, face.score, face.embedding
            )

    def finish_face_detection(self) -> None:
        """Merge the streamed faces onto their ``(frame, track_id)`` rows."""
        if self._data is None:
            return
        faces = faces_to_dataframe(self._tmp_face_frames)
        self.set_data(self._data.merge(faces, on=["frame", "track_id"], how="left"))
        self._tmp_face_frames = []
        self._face_embeddings = self._tmp_face_topk.to_frame()

    def discard_face_detection(self) -> None:
        """Drop a cancelled or failed pass; the tracked boxes are left intact."""
        self._tmp_face_frames = []
        self._tmp_face_topk = _TrackTopK(0)
        self._face_embeddings = None

    def faces_for_frame(self, frame_index: int) -> list[FaceBox]:
        """Detected face boxes for frame ``frame_index`` (0-based)."""
        if self._data is None or "face_score" not in self._data.columns:
            return []
        positions = self._rows_by_frame.get(frame_index)
        if positions is None:
            return []
        rows = self._data.take(positions)
        rows = rows[rows["face_score"].notna()]
        return [face_box_from_row(r) for r in rows.itertuples(index=False)]

    def begin_body_pose_detection(self, embeddings_per_track: int = 0) -> None:
        """Drop any previous pose columns so a fresh pass starts clean.

        ``embeddings_per_track`` is accepted for a uniform ``begin_*`` signature
        across steps but ignored -- pose detection produces no embeddings.
        """
        if self._data is not None:
            present = [c for c in POSE_COLUMNS if c in self._data.columns]
            if present:
                self.set_data(self._data.drop(columns=present))
        self._tmp_pose_frames = []

    def add_body_pose_frame(self, result: PoseFrameResult) -> None:
        """Accumulate one frame's detected body poses for the final merge."""
        self._tmp_pose_frames.append(result)

    def finish_body_pose_detection(self) -> None:
        """Merge the streamed poses onto their ``(frame, track_id)`` rows."""
        if self._data is None:
            return
        poses = poses_to_dataframe(self._tmp_pose_frames)
        self.set_data(self._data.merge(poses, on=["frame", "track_id"], how="left"))
        self._tmp_pose_frames = []

    def discard_body_pose_detection(self) -> None:
        """Drop a cancelled or failed pass; the tracked boxes are left intact."""
        self._tmp_pose_frames = []

    def poses_for_frame(self, frame_index: int) -> list[BodyPose]:
        """Detected body poses for frame ``frame_index`` (0-based)."""
        if self._data is None or "pose_score" not in self._data.columns:
            return []
        positions = self._rows_by_frame.get(frame_index)
        if positions is None:
            return []
        rows = self._data.take(positions)
        rows = rows[rows["pose_score"].notna()]
        return [pose_from_row(r) for r in rows.itertuples(index=False)]

    @property
    def data(self) -> pd.DataFrame | None:
        """All tracked detections as a DataFrame, or ``None`` until complete."""
        return self._data

    @property
    def body_embeddings(self) -> pd.DataFrame | None:
        """Best-K body-appearance embeddings per tracklet, or ``None``."""
        return self._body_embeddings

    @property
    def face_embeddings(self) -> pd.DataFrame | None:
        """Best-K face embeddings per tracklet, or ``None``."""
        return self._face_embeddings

    def boxes_for_frame(self, frame_index: int) -> list[BoundingBox]:
        """Object bounding boxes for frame ``frame_index`` (0-based)."""
        if self._data is None:
            return []
        positions = self._rows_by_frame.get(frame_index)
        if positions is None:
            return []
        rows = self._data.take(positions)
        return [
            BoundingBox(r.x1, r.y1, r.x2, r.y2, int(r.track_id))
            for r in rows.itertuples(index=False)
        ]

    def clear(self) -> None:
        self._data = None
        self._rows_by_frame = {}
        self._tmp_frames = []
        self._tmp_face_frames = []
        self._tmp_pose_frames = []
        self._tmp_body_topk = _TrackTopK(0)
        self._tmp_face_topk = _TrackTopK(0)
        self._body_embeddings = None
        self._face_embeddings = None

    def to_parquet(self, path: str | Path) -> None:
        """Write the completed :attr:`data` to a Parquet file.

        Any collected embeddings are written to companion files beside ``path``
        (``<stem>.body_embeddings.parquet`` / ``<stem>.face_embeddings.parquet``).
        Raises :class:`ValueError` if there is no completed data to write.
        """
        import pyarrow as pa
        import pyarrow.parquet as pq

        if self._data is None:
            raise ValueError("no data to write; run the pipeline first")
        table = pa.Table.from_pandas(self._data, preserve_index=False)
        pq.write_table(table, str(path))
        if self._body_embeddings is not None:
            _write_embeddings(_embeddings_path(path, "body"), self._body_embeddings)
        if self._face_embeddings is not None:
            _write_embeddings(_embeddings_path(path, "face"), self._face_embeddings)

    def load_parquet(self, path: str | Path) -> None:
        """Load results written by :meth:`to_parquet` into this video.

        Replaces any current results with the table at ``path`` and its companion
        embeddings files (``<stem>.body_embeddings.parquet`` /
        ``<stem>.face_embeddings.parquet``), if present.
        """
        self.clear()
        self.set_data(pd.read_parquet(path))
        body_path = _embeddings_path(path, "body")
        if body_path.exists():
            self._body_embeddings = _read_embeddings(body_path)
        face_path = _embeddings_path(path, "face")
        if face_path.exists():
            self._face_embeddings = _read_embeddings(face_path)

    @classmethod
    def from_parquet(cls, path: str | Path) -> "Video":
        """A new :class:`Video` loaded from a Parquet file (see :meth:`load_parquet`)."""
        video = cls()
        video.load_parquet(path)
        return video