clip-annotator/src/river_annotation_tool/config.py

from dataclasses import dataclass, field
from pathlib import Path

import yaml


@dataclass
class FilenameConfig:
    video_in_zip: str = "left.mp4"
    video_tmp_suffix: str = ".mp4"
    zip_extension: str = ".zip"
    mask: str = "mask.png"
    metadata: str = "metadata.json"
    frame: str = "frame.png"
    overlay: str = "overlay.png"
    mask_vis: str = "mask_vis.png"
    gif_original_hires: str = "video_original_hires.gif"
    gif_original_lowres: str = "video_original_lowres.gif"
    gif_overlay_hires: str = "video_overlay_hires.gif"
    gif_overlay_lowres: str = "video_overlay_lowres.gif"


@dataclass
class AppConfig:
    storage: str  # required: 'local' or 's3'
    display_max: int = 480
    fps_fallback: int = 25
    max_frames: int = 100
    data_dir: str = "data/clips"
    out_dir: str = "data/annotation_results"
    clips_file: str = "config/clips.txt"
    optical_flow_config_file: str = ""
    questions: list = field(default_factory=list)
    filenames: FilenameConfig = field(default_factory=FilenameConfig)

    def get_questions(self):
        return [
            (
                s["section"],
                [
                    (
                        item["key"],
                        item["label"],
                        [str(o) for o in item["options"]],
                        str(item["default"])
                        if item.get("default") is not None
                        else None,
                    )
                    for item in s["items"]
                ],
            )
            for s in self.questions
        ]


@dataclass
class OpticalFlowConfig:
    enabled: bool = False
    norm_squared_threshold: float = 0.3
    gaussian_kernel: tuple[int, int] = (5, 5)
    brightness_range: tuple[int, int] = (20, 235)


def load_optical_flow_config(path: Path) -> OpticalFlowConfig:
    with open(path) as f:
        data = yaml.safe_load(f)
    data["gaussian_kernel"] = tuple(data["gaussian_kernel"])
    data["brightness_range"] = tuple(data["brightness_range"])
    return OpticalFlowConfig(**data)


def load_config(path: Path) -> AppConfig:
    with open(path) as f:
        data = yaml.safe_load(f)
    if "storage" not in data:
        raise ValueError(
            f"{path}: missing required field 'storage'. Set it to 'local' or 's3'."
        )
    fn_data = data.pop("filenames", {})
    cfg = AppConfig(**data)
    cfg.filenames = FilenameConfig(**fn_data)
    return cfg