Newer
Older
TIASJudgeModel / tests / data / test_images.py
"""前処理(レターボックス)のテスト(SPEC_02「前処理」「学習の設定」の train.image_size).

- アスペクト比を保ち,長辺を image_size に合わせる(縮小・拡大とも)
- 余白は黒 (0, 0, 0)
- 読めない画像は止める.例外メッセージに画像のファイル名を出さない
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
from PIL import Image

from tias_judge_model.data.images import letterbox, load_letterboxed_images
from tias_judge_model.data.manifest import ManifestRecord

COLOR = (200, 120, 90)


def _bbox_of_content(arr: np.ndarray) -> tuple[int, int]:
    """非黒画素の (高さ, 幅)."""
    mask = arr.any(axis=2)
    rows = np.where(mask.any(axis=1))[0]
    cols = np.where(mask.any(axis=0))[0]
    return int(rows[-1] - rows[0] + 1), int(cols[-1] - cols[0] + 1)


@pytest.mark.parametrize(
    ("width", "height", "size", "expected_hw"),
    [
        (200, 100, 64, (32, 64)),  # 横長の縮小
        (100, 200, 64, (64, 32)),  # 縦長の縮小
        (20, 10, 64, (32, 64)),  # 横長の拡大
        (90, 90, 32, (32, 32)),  # 正方形
    ],
)
def test_letterbox_keeps_aspect_and_fits_long_side(
    width: int, height: int, size: int, expected_hw: tuple[int, int]
) -> None:
    out = letterbox(Image.new("RGB", (width, height), COLOR), size)
    assert out.shape == (size, size, 3)
    assert out.dtype == np.uint8
    assert _bbox_of_content(out) == expected_hw


def test_letterbox_pads_with_black_and_keeps_color() -> None:
    out = letterbox(Image.new("RGB", (200, 100), COLOR), 64)
    # 上下の余白 16 行ずつが黒
    assert (out[:16] == 0).all()
    assert (out[48:] == 0).all()
    # 内側は元の色のまま
    inner = out[20:44, 4:60].reshape(-1, 3)
    assert (np.abs(inner.astype(int) - np.array(COLOR)) <= 1).all()


def test_letterbox_content_is_centered() -> None:
    out = letterbox(Image.new("RGB", (100, 200), COLOR), 64)
    mask = out.any(axis=2)
    cols = np.where(mask.any(axis=0))[0]
    left, right = cols[0], 64 - 1 - cols[-1]
    assert abs(left - right) <= 1


def test_load_letterboxed_images_order_and_shape(tmp_path: Path) -> None:
    images = tmp_path / "images"
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
    records = []
    for i, c in enumerate(colors):
        d = images / "red_strong"
        d.mkdir(parents=True, exist_ok=True)
        Image.new("RGB", (40 + 10 * i, 30), c).save(d / f"img{i}.png")
        records.append(ManifestRecord(f"red_strong/img{i}.png", f"p{i}", "red", "strong"))
    out = load_letterboxed_images(records, images, 16)
    assert out.shape == (3, 16, 16, 3)
    for i, c in enumerate(colors):
        assert tuple(out[i, 8, 8]) == c


def test_load_letterboxed_images_broken_image_hides_filename(tmp_path: Path) -> None:
    images = tmp_path / "images" / "red_strong"
    images.mkdir(parents=True)
    (images / "SECRETNAME.png").write_bytes(b"not a png")
    records = [ManifestRecord("red_strong/SECRETNAME.png", "SECRETNAME", "red", "strong")]
    with pytest.raises(ValueError) as exc_info:
        load_letterboxed_images(records, tmp_path / "images", 16)
    assert "SECRETNAME" not in str(exc_info.value)


def test_load_letterboxed_images_missing_image_hides_filename(tmp_path: Path) -> None:
    (tmp_path / "images").mkdir()
    records = [ManifestRecord("red_strong/SECRETNAME.png", "SECRETNAME", "red", "strong")]
    with pytest.raises(ValueError) as exc_info:
        load_letterboxed_images(records, tmp_path / "images", 16)
    assert "SECRETNAME" not in str(exc_info.value)