Newer
Older
TIASJudgeModel / src / tias_judge_model / data / manifest.py
"""manifest(学習が読む唯一の入口)の生成・書き出し・読み込み・検証.

manifest の列は SPEC_01「データ契約」に従う.
例外メッセージ・ログには患者 ID・画像ファイル名を含めず,行番号・件数・クラス名で示す
(CLAUDE.md 必須ルール).
"""

from __future__ import annotations

import csv
import logging
from collections import Counter
from collections.abc import Iterable, Sequence
from dataclasses import asdict, dataclass
from pathlib import Path, PurePosixPath

from tias_judge_model.data.classes import (
    CLASS_SLUGS,
    SEVERITIES,
    is_valid_combination,
    label_to_id,
    parse_dir_name,
)

logger = logging.getLogger(__name__)

MANIFEST_COLUMNS: tuple[str, ...] = ("image_path", "patient_id", "label", "severity", "note")
REQUIRED_COLUMNS: tuple[str, ...] = ("image_path", "patient_id", "label", "severity")
IMAGE_SUFFIX = ".png"

# CSV ではヘッダが 1 行目なので,レコードの添字 0 は 2 行目になる.
_FIRST_DATA_LINE = 2


@dataclass(frozen=True)
class ManifestRecord:
    """manifest の 1 行.

    Attributes:
        image_path: `images/` からの相対パス(区切りは `/`).
        patient_id: 患者 ID.分割の単位.
        label: クラス slug.
        severity: 程度(`strong` / `mild` / `normal`).
        note: 任意の備考(除外理由等).
    """

    image_path: str
    patient_id: str
    label: str
    severity: str
    note: str = ""

    @property
    def label_id(self) -> int:
        """クラス ID(0〜4)."""
        return label_to_id(self.label)


def build_manifest_records(images_dir: Path) -> list[ManifestRecord]:
    """`images/<label>_<severity>/*.png` を走査して manifest のレコードを作る.

    `patient_id` はファイル名の stem とする(1 患者 1 枚が確定しているため.SPEC_01).
    この前提が崩れる兆候(同じ stem が複数ある)があれば自動では続行しない.

    Args:
        images_dir: 画像ディレクトリ(`<data_root>/images`).

    Returns:
        検証済みのレコード(`image_path` 昇順).

    Raises:
        ValueError: ディレクトリが無い・想定外のディレクトリ名やファイルがある・
            画像が 0 枚のクラスディレクトリがある・画像が 1 枚も無い・stem が重複する・
            レコードの検証に失敗した場合.
    """
    if not images_dir.is_dir():
        raise ValueError(f"画像ディレクトリが存在しない: {images_dir}")

    stray_files = [p for p in images_dir.iterdir() if not p.is_dir()]
    if stray_files:
        raise ValueError(
            f"画像ディレクトリ直下にディレクトリ以外が {len(stray_files)} 件ある"
            "(画像は <label>_<severity>/ の下に置く)"
        )

    records: list[ManifestRecord] = []
    for class_dir in sorted(images_dir.iterdir(), key=lambda p: p.name):
        label, severity = parse_dir_name(class_dir.name)
        entries = list(class_dir.iterdir())
        subdirs = [p for p in entries if p.is_dir()]
        if subdirs:
            raise ValueError(
                f"{class_dir.name}/ にサブディレクトリが {len(subdirs)} 件ある(想定外の構成)"
            )
        non_png = [p for p in entries if p.suffix.lower() != IMAGE_SUFFIX]
        if non_png:
            suffixes = sorted({p.suffix or "(拡張子なし)" for p in non_png})
            raise ValueError(
                f"{class_dir.name}/ に PNG 以外のファイルが {len(non_png)} 件ある"
                f"(拡張子: {', '.join(suffixes)})"
            )
        pngs = sorted(entries, key=lambda p: p.name)
        if not pngs:
            # 空のまま続行すると,そのクラスが欠けた状態で学習・評価が進んでしまう.
            raise ValueError(f"{class_dir.name}/ に画像が 1 枚も無い(空のクラスディレクトリ)")
        for png in pngs:
            records.append(
                ManifestRecord(
                    image_path=f"{class_dir.name}/{png.name}",
                    patient_id=png.stem,
                    label=label,
                    severity=severity,
                )
            )
        logger.info("%s/: %d 件", class_dir.name, len(pngs))

    if not records:
        raise ValueError("画像が 1 枚も見つからない")

    stem_counts = Counter(r.patient_id for r in records)
    duplicated = sum(1 for n in stem_counts.values() if n > 1)
    if duplicated:
        # stem を患者 ID と見なす前提(1 患者 1 枚)が崩れている可能性があるため止める.
        raise ValueError(
            f"ファイル名(stem)が重複する画像が {duplicated} 種類ある."
            "stem を patient_id と見なせないため,実際の患者 ID を確認すること"
        )

    records.sort(key=lambda r: r.image_path)
    validate_records(records)
    logger.info("manifest のレコードを %d 件生成した", len(records))
    return records


def write_manifest(records: Iterable[ManifestRecord], path: Path) -> None:
    """manifest を UTF-8 の CSV として書き出す.

    Args:
        records: 書き出すレコード.
        path: 出力先の CSV パス.親ディレクトリは作成される.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=MANIFEST_COLUMNS)
        writer.writeheader()
        for r in records:
            writer.writerow(asdict(r))
            count += 1
    logger.info("manifest を書き出した(%d 件)", count)


def read_manifest(path: Path) -> list[ManifestRecord]:
    """manifest CSV を読み込み,検証済みのレコードを返す.

    `note` 列は省略可(無ければ空文字とする).それ以外の列は必須.

    Args:
        path: manifest CSV のパス.

    Returns:
        検証済みのレコード(ファイル内の順序のまま).

    Raises:
        ValueError: ファイルが無い・必須列の欠落・列数の不一致・レコードの検証に失敗した場合.
    """
    if not path.is_file():
        raise ValueError(f"manifest が存在しない: {path}")

    # BOM 付きで保存された場合(Excel 等)もヘッダを正しく読むため utf-8-sig を使う.
    with path.open(encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        header = reader.fieldnames or []
        missing = [c for c in REQUIRED_COLUMNS if c not in header]
        if missing:
            raise ValueError(f"manifest に必須列が無い: {', '.join(missing)}")

        records: list[ManifestRecord] = []
        for line_number, row in enumerate(reader, start=_FIRST_DATA_LINE):
            if None in row:
                raise ValueError(f"manifest の {line_number} 行目: 列数がヘッダより多い")
            if any(row.get(c) is None for c in REQUIRED_COLUMNS):
                raise ValueError(f"manifest の {line_number} 行目: 列数がヘッダより少ない")
            records.append(
                ManifestRecord(
                    image_path=row["image_path"],
                    patient_id=row["patient_id"],
                    label=row["label"],
                    severity=row["severity"],
                    note=row.get("note") or "",
                )
            )

    if not records:
        raise ValueError("manifest にデータ行が 1 行も無い")
    validate_records(records)
    logger.info(
        "manifest を読み込んだ(%d 件・患者 %d 人)", len(records), _count_patients(records)
    )
    return records


def validate_records(records: Sequence[ManifestRecord]) -> None:
    """manifest レコードの整合性を検証する.

    エラー位置は CSV の行番号(ヘッダを 1 行目とし,レコードの添字 0 を 2 行目とする)で示す.

    Args:
        records: 検証するレコード.

    Raises:
        ValueError: `image_path` が空・不正な形式・重複,`patient_id` が空,
            未知の label / severity,無効な組み合わせ,同一患者に異なる label がある場合.
    """
    seen_paths: dict[str, int] = {}
    patient_labels: dict[str, tuple[str, int]] = {}
    for index, r in enumerate(records):
        line = index + _FIRST_DATA_LINE
        where = f"manifest の {line} 行目"

        if not r.image_path.strip():
            raise ValueError(f"{where}: image_path が空である")
        _check_image_path_format(r.image_path, where)
        if r.image_path in seen_paths:
            raise ValueError(
                f"{where}: image_path が {seen_paths[r.image_path]} 行目と重複している"
            )
        seen_paths[r.image_path] = line

        if not r.patient_id.strip():
            raise ValueError(f"{where}: patient_id が空である")
        if r.label not in CLASS_SLUGS:
            raise ValueError(f"{where}: 未知の label である: {r.label!r}")
        if r.severity not in SEVERITIES:
            raise ValueError(f"{where}: 未知の severity である: {r.severity!r}")
        if not is_valid_combination(r.label, r.severity):
            raise ValueError(
                f"{where}: label と severity の組み合わせが無効である: {r.label}, {r.severity}"
            )

        previous = patient_labels.get(r.patient_id)
        if previous is None:
            patient_labels[r.patient_id] = (r.label, line)
        elif previous[0] != r.label:
            # 同一患者に異なる label があると,患者単位の層化分割が成り立たない.
            raise ValueError(
                f"{where}: 同一患者に異なる label が付いている"
                f"({previous[1]} 行目は {previous[0]},この行は {r.label})"
            )


def resolve_image_path(images_dir: Path, record: ManifestRecord) -> Path:
    """レコードの画像の実パスを返す.

    Args:
        images_dir: 画像ディレクトリ(`<data_root>/images`).
        record: manifest レコード.

    Returns:
        画像ファイルのパス.
    """
    return images_dir.joinpath(*PurePosixPath(record.image_path).parts)


def _check_image_path_format(image_path: str, where: str) -> None:
    # images/ の外を指すパスや OS 依存の区切りは,サーバと手元で解釈が変わるため拒否する.
    if "\\" in image_path:
        raise ValueError(f"{where}: image_path の区切りは '/' でなければならない")
    pure = PurePosixPath(image_path)
    if pure.is_absolute() or ".." in pure.parts or ":" in image_path:
        raise ValueError(f"{where}: image_path は images/ からの相対パスでなければならない")


def _count_patients(records: Sequence[ManifestRecord]) -> int:
    return len({r.patient_id for r in records})