Experiment API¶

class AudioInput(_Input):
    """Audio recorded on its own device, e.g. a directional microphone.

    Embedded audio in video files is handled as part of video playback; this
    input is for separate audio files.
    """

    glasses_video: str | None = Field(
        None,
        description=(
            "Optional id of the glasses video worn by the participant this "
            "recording captures"
        ),
    )

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."""

    version: int = CURRENT_VERSION
    glasses_videos: list[GlassesVideoInput] = Field(default_factory=list)
    fixed_videos: list[FixedVideoInput] = Field(default_factory=list)
    audio: list[AudioInput] = Field(default_factory=list)
    pipeline: Pipeline = Field(default_factory=Pipeline)

    @model_validator(mode="after")
    def _check(self) -> ExperimentConfig:
        ids = (
            [video.id for video in self.glasses_videos]
            + [video.id for video in self.fixed_videos]
            + [audio.id for audio in self.audio]
        )
        duplicates = {i for i in ids if ids.count(i) > 1}
        if duplicates:
            raise ValueError(f"duplicate input ids: {sorted(duplicates)}")

        glasses_ids = {i.id for i in self.glasses_videos}
        unknown = {
            a.glasses_video
            for a in self.audio
            if a.glasses_video is not None and a.glasses_video not in glasses_ids
        }
        if unknown:
            raise ValueError(f"unknown glasses video ids: {sorted(unknown)}")
        return self

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 FixedVideoInput(_Input):
    """Video recorded by a camera at a fixed position in the room."""

class GlassesVideoInput(_Input):
    """Video and gaze data recorded by a participant's glasses-mounted camera."""

    gaze_path: Path = Field(
        description="Gaze samples recorded alongside this video, as a TSV file."
    )

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 Pipeline(_Model):
    """What to run for each type of input."""

    glasses_video: VideoPipeline = Field(default_factory=VideoPipeline)
    fixed_video: VideoPipeline = Field(default_factory=VideoPipeline)
    speech: SpeechPipeline | None = Field(default_factory=SpeechPipeline)
    speech_post_processing: SpeechPostProcessingSettings = Field(
        default_factory=SpeechPostProcessingSettings
    )

class SpeechPipeline(StepPipeline):
    """The stages run over all inputs that contain audio.

    Transcription is the only one: it says what was said, and who said it is
    settled afterwards, by comparing the experiment's recordings with each other.
    """

    transcription: TranscriptionStep = Field(default_factory=TranscriptionStep)

class SpeechPostProcessingSettings(_Model):
    """How transcripts are combined to form experiment-wide speaker turns."""

    split_gap_seconds: float = Field(
        0.75,
        ge=0,
        description=(
            "Split a segment when consecutive words are separated by at "
            "least this many seconds. Sentence-ending punctuation also splits."
        ),
    )
    split_on_sentence_end: bool = Field(
        True,
        description="Split after words ending in sentence punctuation (. ! ? …).",
    )
    split_on_comma: bool = Field(
        True,
        description=(
            "Split at commas when both adjacent clauses meet the minimum word "
            "count and duration below."
        ),
    )
    minimum_clause_words: int = Field(
        4,
        ge=1,
        description=("Minimum number of words required on each side of a comma split."),
    )
    minimum_clause_seconds: float = Field(
        0.5,
        ge=0,
        description=(
            "Minimum duration required on each side of a comma split, in seconds."
        ),
    )
    floor_percentile: float = Field(
        10.0,
        ge=0,
        le=100,
        description=(
            "Use this percentile of each recording's levels as its quiet floor."
        ),
    )
    live_above_floor_db: float = Field(
        10.0,
        ge=0,
        description=(
            "A recording counts as carrying speech when its level is this "
            "many dB above its own quiet floor."
        ),
    )
    ownership_share: float = Field(
        0.5,
        ge=0,
        le=1,
        description=(
            "Keep a piece when this recording is the loudest for more than this "
            "fraction of its active frames, or live for more than this fraction of all frames."
        ),
    )
    fuzzy_agreement: float = Field(
        0.6,
        ge=0,
        le=1,
        description=(
            "Treat overlapping text as the same voice when its "
            "normalized character similarity exceeds this fraction."
        ),
    )

class StepPipeline(_Model):
    """Base class for pipeline configurations edited by the shared GUI."""

class TimelineConfig(_Model):
    """Where one recording sits on the experiment's shared clock."""

    offset: float = Field(
        0.0,
        description=(
            "Seconds to add to this input's own clock to place it on the shared "
            "experiment timeline."
        ),
    )
    rate: float = Field(
        1.0,
        gt=0,
        description=(
            "Experiment seconds per second of this input's own clock. Every "
            "device counts time on its own crystal, and two of them differ by "
            "tens of parts per million, which is tens of milliseconds across a "
            "long recording."
        ),
    )

class TranscriptionStep(_Model):
    """Speech transcription. Fields mirror ``transcribe``."""

    model_name: str = Field(
        "primeline/whisper-large-v3-turbo-german",
        description=(
            "Whisper model. The primeLine models are accuracy-tuned for German; "
            "CrisperWhisper models produce verbatim transcripts; large-v3 is "
            "the strongest general multilingual choice."
        ),
        json_schema_extra={
            "choices": [
                "primeline/whisper-large-v3-turbo-german",
                "primeline/whisper-large-v3-german",
                "nyralabs/CrisperWhisper2.0_large",
                "nyralabs/CrisperWhisper2.0_medium",
                "large-v3",
                "large-v3-turbo",
                "distil-large-v3",
                "tiny",
                "base",
                "small",
                "medium",
            ]
        },
    )
    language: str | None = Field(
        "de",
        description=(
            "ISO 639-1 language code of the recording, e.g. 'de'. German is the "
            "accuracy-first default; leave unset to detect the language from the "
            "first 30 seconds."
        ),
    )
    beam_size: int = Field(5, ge=1, description="Decoding beam width.")
    vad_filter: bool = Field(
        False,
        description=(
            "Skip silent stretches, which speeds up the pass and suppresses text "
            "invented over silence."
        ),
    )

class VideoPipeline(StepPipeline):
    """The stages run over a video input.

    Both video types use this same set of stages, but as independent blocks, so
    e.g. a room camera can be tracked with a different detector than the glasses
    cameras.
    """

    object_tracking: ObjectTrackingStep = Field(default_factory=ObjectTrackingStep)
    face_detection: FaceDetectionStep | None = None
    body_pose: BodyPoseStep | None = None

class Experiment:
    """An experiment: its inputs, their results, and the pipeline to run.

    The inputs own their settings, each in the runtime class for its type;
    :class:`ExperimentConfig` is the on-disk form, converted to and from when
    the experiment is saved and loaded. Input ids are unique across the types,
    and inputs are added, removed and renamed through this class so they stay
    that way.
    """

    def __init__(self, config: ExperimentConfig, folder: str | Path | None = None):
        self.folder = Path(folder) if folder is not None else None
        self.pipeline: Pipeline = config.pipeline
        self.glasses_videos = [
            GlassesVideo(
                id=spec.id,
                path=self._resolve(spec.path),
                gaze_path=self._resolve(spec.gaze_path),
                timeline=Timeline.from_config(spec.timeline),
            )
            for spec in config.glasses_videos
        ]
        self.fixed_videos = [
            FixedVideo(
                id=spec.id,
                path=self._resolve(spec.path),
                timeline=Timeline.from_config(spec.timeline),
            )
            for spec in config.fixed_videos
        ]
        glasses_by_id = {video.id: video for video in self.glasses_videos}
        self.audio = [
            Audio(
                id=spec.id,
                path=self._resolve(spec.path),
                glasses_video=glasses_by_id.get(spec.glasses_video),
                timeline=Timeline.from_config(spec.timeline),
            )
            for spec in config.audio
        ]
        self.speech_turns = SpeechTurns()

    def _load_stored_data(self) -> None:
        """Load any existing outputs owned by this experiment from its folder."""
        for data in self.inputs:
            self._load_results(data)
        self._load_speech_turns()

    @property
    def inputs(self) -> list[Video | Audio]:
        """Every input of every type, for what applies to all of them."""
        return [*self.glasses_videos, *self.fixed_videos, *self.audio]

    def add_glasses_video(self, spec: GlassesVideoInput) -> GlassesVideo:
        """Add a glasses video input, returning its :class:`GlassesVideo`."""
        self._check_id(spec.id)
        video = GlassesVideo(
            id=spec.id,
            path=self._resolve(spec.path),
            gaze_path=self._resolve(spec.gaze_path),
            timeline=Timeline.from_config(spec.timeline),
        )
        self.glasses_videos.append(video)
        return video

    def add_fixed_video(self, spec: FixedVideoInput) -> FixedVideo:
        """Add a fixed video input, returning its :class:`FixedVideo`."""
        self._check_id(spec.id)
        video = FixedVideo(
            id=spec.id,
            path=self._resolve(spec.path),
            timeline=Timeline.from_config(spec.timeline),
        )
        self.fixed_videos.append(video)
        return video

    def add_audio(self, spec: AudioInput) -> Audio:
        """Add an audio input, returning its :class:`Audio`.

        Raises :class:`ValueError` if it names a glasses video that is not in
        this experiment.
        """
        self._check_id(spec.id)
        glasses_video = None
        if spec.glasses_video is not None:
            glasses_video = next(
                (v for v in self.glasses_videos if v.id == spec.glasses_video), None
            )
            if glasses_video is None:
                raise ValueError(f"unknown glasses video id: {spec.glasses_video!r}")
        audio = Audio(
            id=spec.id,
            path=self._resolve(spec.path),
            glasses_video=glasses_video,
            timeline=Timeline.from_config(spec.timeline),
        )
        self.audio.append(audio)
        return audio

    def remove_input(self, data: Video | Audio) -> None:
        """Remove an input from the experiment, leaving its output files alone.

        Raises :class:`ValueError` if it is a glasses video that audio inputs
        still refer to.
        """
        if isinstance(data, GlassesVideo):
            used_by = sorted(a.id for a in self.audio if a.glasses_video is data)
            if used_by:
                raise ValueError(
                    f"glasses video {data.id!r} is still used by audio inputs: {used_by}"
                )
        for inputs in (self.glasses_videos, self.fixed_videos, self.audio):
            if any(existing is data for existing in inputs):
                inputs.remove(data)
                return
        raise ValueError(f"input {data.id!r} is not in this experiment")

    def rename_input(self, data: Video | Audio, new_id: str) -> None:
        """Give an input a new id, moving its output directory if it exists."""
        if new_id == data.id:
            return
        self._check_id(new_id)
        old_id = data.id
        old_output_dir = None
        if self.folder is not None:
            old_output_dir = self.output_dir_for(data)
        data.id = new_id
        if old_output_dir is not None and old_output_dir.exists():
            try:
                old_output_dir.rename(self.output_dir_for(data))
            except OSError:
                data.id = old_id
                raise

    def _check_id(self, input_id: str) -> None:
        """Check an id can name an output directory, and nothing else uses it."""
        validate_input_id(input_id)
        if any(data.id == input_id for data in self.inputs):
            raise ValueError(f"duplicate input id: {input_id!r}")
        if self.folder is not None and (self.output_dir / input_id).exists():
            raise ValueError(f"outputs already exist for input id: {input_id!r}")

    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 _resolve(self, path: Path) -> Path:
        """An input path as used at runtime: absolute where the folder allows."""
        if path.is_absolute() or self.folder is None:
            return path
        return (self.folder / path).resolve()

    def _store(self, path: Path) -> Path:
        """An input path as written to disk: relative where it is under the folder."""
        if self.folder is not None and path.is_relative_to(self.folder.resolve()):
            return path.relative_to(self.folder.resolve())
        return path

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

    def output_dir_for(self, data: Video | Audio) -> Path:
        """The output directory an input owns, inside :attr:`output_dir`."""
        return self.output_dir / data.id

    def _load_results(self, data: Video | Audio) -> None:
        """Fill an input from its stored results, skipping any it cannot read."""
        if self.folder is None:
            return
        directory = self.output_dir_for(data)
        try:
            data.load(directory)
        except (OSError, ValueError) as exc:
            data.clear()
            logger.warning(
                "ignoring unreadable results for input %r in %s: %s",
                data.id,
                directory,
                exc,
            )

    def _load_speech_turns(self) -> None:
        """Fill the experiment's speech turns from its output directory."""
        if self.folder is None:
            return
        try:
            self.speech_turns.load(self.output_dir)
        except (OSError, ValueError) as exc:
            self.speech_turns.clear()
            logger.warning(
                "ignoring unreadable speech turns in %s: %s", self.output_dir, exc
            )

    def config(self) -> ExperimentConfig:
        """The experiment in its on-disk form."""
        return ExperimentConfig(
            glasses_videos=[
                GlassesVideoInput(
                    id=v.id,
                    path=self._store(v.video_path),
                    gaze_path=self._store(v.gaze_path),
                    timeline=v.timeline.to_config(),
                )
                for v in self.glasses_videos
            ],
            fixed_videos=[
                FixedVideoInput(
                    id=v.id,
                    path=self._store(v.video_path),
                    timeline=v.timeline.to_config(),
                )
                for v in self.fixed_videos
            ],
            audio=[
                AudioInput(
                    id=a.id,
                    path=self._store(a.audio_path),
                    timeline=a.timeline.to_config(),
                    glasses_video=(
                        a.glasses_video.id if a.glasses_video is not None else None
                    ),
                )
                for a in self.audio
            ],
            pipeline=self.pipeline,
        )

    @classmethod
    def load(cls, folder: str | Path) -> Experiment:
        """Load the experiment in ``folder``: its inputs and their results."""
        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"
            )
        experiment = cls(ExperimentConfig.model_validate(data), folder)
        experiment._load_stored_data()
        return experiment

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

        ``folder`` defaults to the current :attr:`folder` and becomes it when
        given. The config is always written; an input 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 data in self.inputs:
            if data.has_data():
                data.save(self.output_dir_for(data))
        if self.speech_turns.has_data():
            self.speech_turns.save(self.output_dir)