"""データ拡張と Dataset のテスト（SPEC_02「データ拡張」「学習の設定」）.
- **色を変える拡張は行わない**: 一様な色の画像に拡張をかけても，舌領域の画素は元の色のまま
（はみ出しの黒との境界の補間を除き，黒と元の色の混合＝元の色の定数倍にしかならない）
- 拡張は幾何変換（左右反転等）であり，実際に画像が変わりうる
- 評価時（augment=False）は拡張しない
- 正規化は ImageNet の mean / std
"""
from __future__ import annotations
import numpy as np
import pytest
import torch
from torchvision.transforms import v2
from tias_judge_model.data.dataset import (
TongueImageDataset,
make_train_augmentation,
normalize_batch,
)
COLOR = np.array([190, 110, 95], dtype=np.uint8)
N_TRIALS = 40
def _uniform_tongue(size: int = 64) -> np.ndarray:
"""黒背景の中央に一様な色の領域がある画像 `(H, W, 3)`."""
img = np.zeros((size, size, 3), dtype=np.uint8)
img[8 : size - 8, 12 : size - 12] = COLOR
return img
def _flatten_transforms(t: object) -> list[object]:
inner = getattr(t, "transforms", None)
if inner is None:
return [t]
out: list[object] = []
for x in inner:
out.extend(_flatten_transforms(x))
return out
def test_augmentation_contains_no_color_transforms() -> None:
forbidden = (
v2.ColorJitter,
v2.Grayscale,
v2.RandomGrayscale,
v2.RandomAutocontrast,
v2.RandomEqualize,
v2.RandomPosterize,
v2.RandomSolarize,
v2.RandomInvert,
v2.RandomAdjustSharpness,
v2.RandomPhotometricDistort,
v2.RandomChannelPermutation,
v2.AutoAugment,
v2.RandAugment,
v2.TrivialAugmentWide,
v2.AugMix,
)
for t in _flatten_transforms(make_train_augmentation()):
assert not isinstance(t, forbidden), f"色を変える変換が含まれる: {type(t).__name__}"
def test_augmentation_does_not_change_pixel_colors() -> None:
torch.manual_seed(0)
ds = TongueImageDataset(_uniform_tongue()[None], np.array([2]), augment=True)
color = COLOR.astype(np.float64)
for _ in range(N_TRIALS):
image, label = ds[0]
assert label == 2
arr = image.permute(1, 2, 0).numpy().astype(np.float64).reshape(-1, 3)
nonblack = arr[arr.any(axis=1)]
assert len(nonblack) > 0
# 画素は黒と元の色の混合（= 元の色の t 倍）でしかない（色相・彩度が変わらない）
t = nonblack.sum(axis=1, keepdims=True) / color.sum()
assert np.abs(nonblack - t * color).max() <= 2.0
assert t.max() <= 1.0 + 1e-2 # 明るくならない
# 舌領域の大半は元の色そのまま
exact = np.all(np.abs(nonblack - color) <= 1, axis=1)
assert exact.mean() > 0.8
def test_augmentation_keeps_shape_and_dtype() -> None:
torch.manual_seed(1)
ds = TongueImageDataset(_uniform_tongue(48)[None], np.array([0]), augment=True)
image, _ = ds[0]
assert image.shape == (3, 48, 48)
assert image.dtype == torch.uint8
def test_augmentation_actually_transforms_geometry() -> None:
# 左右非対称な画像で，拡張をかけると少なくとも 1 回は元と異なる
img = np.zeros((32, 32, 3), dtype=np.uint8)
img[4:28, 2:12] = COLOR
ds = TongueImageDataset(img[None], np.array([0]), augment=True)
original = torch.from_numpy(img).permute(2, 0, 1)
torch.manual_seed(2)
assert any(not torch.equal(ds[0][0], original) for _ in range(N_TRIALS))
def test_eval_dataset_is_not_augmented() -> None:
img = _uniform_tongue(32)
img[0, 0] = (1, 2, 3) # 非対称にする
ds = TongueImageDataset(img[None], np.array([4]), augment=False)
original = torch.from_numpy(img).permute(2, 0, 1)
for _ in range(10):
image, label = ds[0]
assert torch.equal(image, original)
assert label == 4
def test_dataset_length_mismatch_raises() -> None:
with pytest.raises(Exception): # noqa: B017 件数不一致で止まれば型は問わない
TongueImageDataset(np.zeros((2, 8, 8, 3), dtype=np.uint8), np.array([0]), augment=False)
def test_normalize_batch_uses_imagenet_stats() -> None:
mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
batch = torch.tensor(np.stack([_uniform_tongue(8)] * 2)).permute(0, 3, 1, 2)
out = normalize_batch(batch)
assert out.dtype == torch.float32
expected = (batch.float() / 255.0 - mean) / std
assert torch.allclose(out, expected, atol=1e-5)
# 黒画素は -mean/std
assert torch.allclose(out[0, :, 0, 0], (-mean / std).view(3), atol=1e-5)
