"""前処理(レターボックス)のテスト(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)