Newer
Older
TIASJudgeModel / src / tias_judge_model / data / classes.py
"""舌色クラスと程度(severity)の固定定義.

整数 ID は混同行列の軸をぶらさないために固定しており,順序を変更してはならない
(SPEC_01「クラス定義」).
"""

from __future__ import annotations

# ID 順に並べる.この並びがそのままクラス ID になる.
CLASS_SLUGS: tuple[str, ...] = ("pale", "light_red", "red", "dark_red", "purple")

CLASS_DISPLAY_NAMES: dict[str, str] = {
    "pale": "淡白",
    "light_red": "淡紅",
    "red": "紅",
    "dark_red": "暗赤紅",
    "purple": "紫",
}

NUM_CLASSES: int = len(CLASS_SLUGS)

SEVERITIES: tuple[str, ...] = ("strong", "mild", "normal")

# 淡紅(正常)のみ程度の区分が無い(TEST_01「フォルダ構成とラベル対応」).
VALID_SEVERITIES: dict[str, tuple[str, ...]] = {
    "pale": ("strong", "mild"),
    "light_red": ("normal",),
    "red": ("strong", "mild"),
    "dark_red": ("strong", "mild"),
    "purple": ("strong", "mild"),
}

_SLUG_TO_ID: dict[str, int] = {slug: i for i, slug in enumerate(CLASS_SLUGS)}


def label_to_id(label: str) -> int:
    """クラス slug を整数 ID に変換する.

    Args:
        label: クラス slug(例: `"dark_red"`).

    Returns:
        クラス ID(0〜4).

    Raises:
        ValueError: 未知の slug の場合.
    """
    if label not in _SLUG_TO_ID:
        raise ValueError(f"未知のクラス slug である: {label!r}(有効値: {', '.join(CLASS_SLUGS)})")
    return _SLUG_TO_ID[label]


def id_to_label(class_id: int) -> str:
    """整数 ID をクラス slug に変換する.

    Args:
        class_id: クラス ID(0〜4).

    Returns:
        クラス slug.

    Raises:
        ValueError: 範囲外の ID の場合.
    """
    if not 0 <= class_id < NUM_CLASSES:
        raise ValueError(f"クラス ID が範囲外である: {class_id}(0〜{NUM_CLASSES - 1})")
    return CLASS_SLUGS[class_id]


def is_valid_combination(label: str, severity: str) -> bool:
    """クラスと程度の組み合わせが有効かを返す.

    Args:
        label: クラス slug.
        severity: 程度(`strong` / `mild` / `normal`).

    Returns:
        有効な組み合わせなら True.未知の値を含む場合も False.
    """
    return severity in VALID_SEVERITIES.get(label, ())


def make_dir_name(label: str, severity: str) -> str:
    """クラスと程度から画像ディレクトリ名 `<label>_<severity>` を作る.

    Args:
        label: クラス slug.
        severity: 程度.

    Returns:
        ディレクトリ名(例: `"light_red_normal"`).

    Raises:
        ValueError: 無効な組み合わせの場合.
    """
    if not is_valid_combination(label, severity):
        raise ValueError(f"無効なクラスと程度の組み合わせである: {label!r}, {severity!r}")
    return f"{label}_{severity}"


def parse_dir_name(name: str) -> tuple[str, str]:
    """画像ディレクトリ名 `<label>_<severity>` を (label, severity) に解釈する.

    `light_red_normal` のように label 自体が `_` を含むため,末尾の `_` で区切る.

    Args:
        name: ディレクトリ名.

    Returns:
        (label, severity) のタプル.

    Raises:
        ValueError: 書式が不正,または無効な組み合わせの場合.
    """
    label, sep, severity = name.rpartition("_")
    if not sep or not is_valid_combination(label, severity):
        valid = ", ".join(ALL_DIR_NAMES)
        raise ValueError(f"想定外の画像ディレクトリ名である: {name!r}(有効値: {valid})")
    return label, severity


# 有効なディレクトリ名の一覧(クラス ID 順,程度は VALID_SEVERITIES の順).
ALL_DIR_NAMES: tuple[str, ...] = tuple(
    make_dir_name(label, severity) for label in CLASS_SLUGS for severity in VALID_SEVERITIES[label]
)