Newer
Older
TIASJudgeModel / tests / segmentation / test_seg_images.py
"""セグメンテーションの評価データの列挙と読み込み(`segmentation/images.py`)のテスト.

(`tests/data/test_images.py` と basename が衝突しないよう `test_seg_images.py` とする)

仕様(実装サマリーを仕様として扱う):

- `<set>/sample/*` と `<set>/label/*` をファイル名の stem で対応づけ,stem の昇順で返す
- 拡張子は jpg・jpeg・png・bmp(大文字小文字を問わない).それ以外は無視する
- 対応の取れない画像・stem の重複・0 枚・ディレクトリが無いと ValueError.
  例外メッセージにファイル名(stem)を出さない(件数のみ)
- 正解マスクはグレースケールで読み,`threshold` 以上を前景とする
- 日本語を含むパスでも読める(`np.fromfile` → `cv2.imdecode`)
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
from PIL import Image

from tias_judge_model.segmentation.images import list_image_pairs, read_bgr, read_label_mask

SECRET = "SECRETPT"


def _save(path: Path, arr: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(arr).save(path)


def _pair(set_dir: Path, stem: str, ext: str = ".png") -> None:
    _save(set_dir / "sample" / f"{stem}{ext}", np.zeros((4, 4, 3), np.uint8))
    _save(set_dir / "label" / f"{stem}.png", np.zeros((4, 4), np.uint8))


def test_pairs_are_matched_by_stem_and_sorted(tmp_path: Path) -> None:
    for stem in ["b", "a", "c"]:
        _pair(tmp_path, stem, ".jpg" if stem == "a" else ".png")
    (tmp_path / "sample" / "notes.txt").write_text("x", encoding="utf-8")
    pairs = list_image_pairs(tmp_path, "sample", "label")
    assert [p.stem for p in pairs] == ["a", "b", "c"]
    assert pairs[0].sample.suffix == ".jpg"
    assert pairs[0].label.name == "a.png"


def test_uppercase_extension_is_accepted(tmp_path: Path) -> None:
    _save(tmp_path / "sample" / "x.JPG", np.zeros((4, 4, 3), np.uint8))
    _save(tmp_path / "label" / "x.BMP", np.zeros((4, 4), np.uint8))
    assert [p.stem for p in list_image_pairs(tmp_path, "sample", "label")] == ["x"]


def test_unmatched_images_raise_without_names(tmp_path: Path) -> None:
    _pair(tmp_path, "ok")
    _save(tmp_path / "sample" / f"{SECRET}1.png", np.zeros((4, 4, 3), np.uint8))
    _save(tmp_path / "label" / f"{SECRET}2.png", np.zeros((4, 4), np.uint8))
    _save(tmp_path / "label" / f"{SECRET}3.png", np.zeros((4, 4), np.uint8))
    with pytest.raises(ValueError) as exc:
        list_image_pairs(tmp_path, "sample", "label")
    msg = str(exc.value)
    assert SECRET not in msg
    assert "1 件" in msg and "2 件" in msg


def test_duplicate_stems_raise_without_names(tmp_path: Path) -> None:
    _pair(tmp_path, SECRET)
    _save(tmp_path / "sample" / f"{SECRET}.jpg", np.zeros((4, 4, 3), np.uint8))
    with pytest.raises(ValueError) as exc:
        list_image_pairs(tmp_path, "sample", "label")
    assert SECRET not in str(exc.value)


def test_empty_set_raises(tmp_path: Path) -> None:
    (tmp_path / "sample").mkdir()
    (tmp_path / "label").mkdir()
    with pytest.raises(ValueError):
        list_image_pairs(tmp_path, "sample", "label")


def test_missing_directory_raises(tmp_path: Path) -> None:
    (tmp_path / "sample").mkdir()
    with pytest.raises(ValueError):
        list_image_pairs(tmp_path, "sample", "label")


def test_read_bgr_returns_bgr_order(tmp_path: Path) -> None:
    rgb = np.zeros((3, 5, 3), np.uint8)
    rgb[..., 0] = 200  # R
    path = tmp_path / "日本語ディレクトリ" / "画像.png"
    _save(path, rgb)
    bgr = read_bgr(path, "評価セット x の 1 件目")
    assert bgr.shape == (3, 5, 3)
    assert np.all(bgr[..., 2] == 200) and np.all(bgr[..., 0] == 0)


def test_read_label_mask_threshold_is_inclusive(tmp_path: Path) -> None:
    mask = np.array([[0, 127, 128, 255]], np.uint8)
    path = tmp_path / "ラベル.png"
    _save(path, mask)
    got = read_label_mask(path, 128, "w")
    assert got.dtype == bool
    assert got.tolist() == [[False, False, True, True]]


def test_read_label_mask_converts_color_to_gray(tmp_path: Path) -> None:
    rgb = np.zeros((2, 2, 3), np.uint8)
    rgb[0, 0] = 255
    path = tmp_path / "m.png"
    _save(path, rgb)
    got = read_label_mask(path, 128, "w")
    assert got.shape == (2, 2)
    assert got.tolist() == [[True, False], [False, False]]


def test_unreadable_image_raises_with_position_only(tmp_path: Path) -> None:
    path = tmp_path / f"{SECRET}.png"
    path.write_bytes(b"not an image")
    with pytest.raises(ValueError) as exc:
        read_bgr(path, "評価セット s の 3 件目")
    assert "3 件目" in str(exc.value)
    assert SECRET not in str(exc.value)
    with pytest.raises(ValueError) as exc:
        read_bgr(tmp_path / f"{SECRET}_missing.png", "評価セット s の 4 件目")
    assert SECRET not in str(exc.value)