Files
ModelTrainEval/tests/test_unlabeled_dataset.py
T
Shaun ZhouandClaude Opus 5 d6adea5b05 test: 单元测试与端到端冒烟
88 个测试覆盖配置继承、注册表、GPU 守卫、YOLO 格式解析、数据集、模型形状、
标签模式、聚类原语与 SimCLR 损失,以及 train.py / eval.py / 聚类脚本的入口接线。

含两条回归测试:
- PCA 白化放大噪声方向导致簇塌陷(默认 whiten=False 的依据)
- classes 写成 data.yaml 路径时 num_classes 必须是类别数而非路径字符串长度

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 17:55:47 +08:00

132 lines
4.8 KiB
Python

"""无标注数据集与聚类相关的预处理。"""
import pytest
pytest.importorskip("torch")
pytest.importorskip("PIL")
import numpy as np # noqa: E402
import torch # noqa: E402
from PIL import Image # noqa: E402
from datasets.build import build_dataset # noqa: E402
from datasets.transforms import build_transforms # noqa: E402
from datasets.unlabeled import UnlabeledImageDataset, shape_features, unlabeled_collate # noqa: E402
TRANSFORMS = [
{"name": "Resize", "size": [32, 32]},
{"name": "ToTensor"},
{"name": "PerImageStandardize"},
]
@pytest.fixture
def image_root(tmp_path):
"""尺寸各异、含 RGBA 与灰度、还有一个要跳过的子目录。"""
(tmp_path / "a").mkdir()
(tmp_path / "b" / "deep").mkdir(parents=True)
(tmp_path / "images").mkdir()
Image.fromarray(np.zeros((113, 152, 3), dtype=np.uint8)).save(tmp_path / "a" / "1.png")
Image.fromarray(np.zeros((60, 900, 4), dtype=np.uint8), mode="RGBA").save(tmp_path / "a" / "2.png")
Image.fromarray(np.zeros((80, 40), dtype=np.uint8)).save(tmp_path / "b" / "deep" / "3.png")
Image.fromarray(np.zeros((10, 10, 3), dtype=np.uint8)).save(tmp_path / "images" / "skip.png")
(tmp_path / "a" / "notes.txt").write_text("忽略", encoding="utf-8")
return tmp_path
def test_recursive_scan_skips_excluded_dirs_and_non_images(image_root):
dataset = UnlabeledImageDataset(image_root, transforms=TRANSFORMS, exclude=["images"])
names = [path.name for path in dataset.images]
assert names == ["1.png", "2.png", "3.png"]
def test_mixed_sizes_and_modes_become_one_tensor_shape(image_root):
dataset = UnlabeledImageDataset(image_root, transforms=TRANSFORMS, exclude=["images"])
for index in range(len(dataset)):
assert dataset[index]["views"][0].shape == (3, 32, 32) # RGBA / L 都统一到 3 通道
def test_gray_mode_and_metadata(image_root):
dataset = UnlabeledImageDataset(
image_root,
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
gray=True,
exclude=["images"],
)
sample = dataset[0]
assert sample["views"][0].shape == (1, 32, 32)
assert sample["meta"]["size"] == (152, 113)
assert sample["meta"]["source_dir"] == "a"
def test_two_views_are_augmented_independently(image_root):
dataset = UnlabeledImageDataset(
image_root,
transforms=TRANSFORMS,
views=2,
view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.3, 0.6]}],
exclude=["images"],
)
sample = dataset[0]
assert len(sample["views"]) == 2
batch = unlabeled_collate([dataset[0], dataset[1]])
assert len(batch["views"]) == 2
assert batch["views"][0].shape == (2, 3, 32, 32)
assert batch["shape"].shape == (2, 4)
assert len(batch["meta"]) == 2
def test_registry_build(image_root):
dataset = build_dataset(
{"name": "unlabeled", "root": str(image_root), "transforms": TRANSFORMS, "exclude": ["images"]}
)
assert len(dataset) == 3
def test_empty_dir_raises(tmp_path):
(tmp_path / "empty").mkdir()
with pytest.raises(FileNotFoundError, match="没有找到图片"):
UnlabeledImageDataset(tmp_path / "empty")
with pytest.raises(FileNotFoundError, match="目录不存在"):
UnlabeledImageDataset(tmp_path / "nope")
def test_shape_features_are_log_scaled():
features = shape_features(200, 100)
assert features.shape == (4,)
assert features[2] == pytest.approx(np.log(2.0)) # 宽高比
assert features[3] == pytest.approx(np.log(20000))
def test_per_image_standardize_removes_brightness_and_contrast():
"""两张只差整体亮度/增益的图,标准化后应当几乎一致。"""
transform = build_transforms([{"name": "ToTensor"}, {"name": "PerImageStandardize"}])
base = np.random.default_rng(0).integers(40, 80, (32, 32, 3), dtype=np.uint8)
dark = transform({"image": Image.fromarray(base)})["image"]
bright = transform({"image": Image.fromarray(np.clip(base * 1.5 + 40, 0, 255).astype(np.uint8))})[
"image"
]
assert torch.allclose(dark, bright, atol=0.05)
assert float(dark.mean()) == pytest.approx(0.0, abs=1e-4)
assert float(dark.std()) == pytest.approx(1.0, abs=1e-2)
def test_letterbox_and_centercrop_handle_extreme_aspect_ratio():
"""40:1 的长条,两种预处理都得吐出固定尺寸。"""
strip = Image.fromarray(np.zeros((113, 4464, 3), dtype=np.uint8))
letterboxed = build_transforms([{"name": "Letterbox", "size": [128, 128]}, {"name": "ToTensor"}])(
{"image": strip}
)
assert letterboxed["image"].shape == (3, 128, 128)
assert letterboxed["letterbox"]["pad"][1] > 0 # 上下补边
cropped = build_transforms([{"name": "CenterCrop", "size": [128, 128]}, {"name": "ToTensor"}])(
{"image": strip}
)
assert cropped["image"].shape == (3, 128, 128)