Shaun ZhouandClaude Opus 5
d6adea5b05
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>
152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
"""YOLO 风格标注的数据集读取。"""
|
||
|
||
import pytest
|
||
|
||
pytest.importorskip("torch")
|
||
pytest.importorskip("PIL")
|
||
|
||
from pathlib import Path # noqa: E402
|
||
|
||
import numpy as np # noqa: E402
|
||
import torch # noqa: E402
|
||
from PIL import Image # noqa: E402
|
||
|
||
from datasets import FilmDefectDataset, collate_fn # noqa: E402
|
||
|
||
CLASSES = ["scratch", "bubble"]
|
||
TRANSFORMS = [
|
||
{"name": "Resize", "size": [64, 64]},
|
||
{"name": "ToTensor"},
|
||
{"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]},
|
||
]
|
||
|
||
|
||
@pytest.fixture
|
||
def dataset_root(tmp_path):
|
||
"""YOLO 目录：images/ 与同层 labels/，128x128 的图，标注为归一化 cx cy w h。"""
|
||
images = tmp_path / "images"
|
||
labels = tmp_path / "labels"
|
||
images.mkdir()
|
||
labels.mkdir()
|
||
|
||
for name in ("a.png", "b.png"):
|
||
Image.fromarray(np.zeros((128, 128, 3), dtype=np.uint8)).save(images / name)
|
||
# 像素框 (10,10,50,50) -> 中心 (30,30)/128, 宽高 40/128
|
||
(labels / "a.txt").write_text("0 0.234375 0.234375 0.312500 0.312500\n", encoding="utf-8")
|
||
(labels / "b.txt").write_text("", encoding="utf-8") # 空文件 = 无缺陷
|
||
return images, labels
|
||
|
||
|
||
def build(dataset_root, **kwargs):
|
||
images, _ = dataset_root
|
||
return FilmDefectDataset(
|
||
img_dir=images, classes=CLASSES, transforms=TRANSFORMS, stride=4, **kwargs
|
||
)
|
||
|
||
|
||
def test_label_dir_follows_yolo_convention(dataset_root):
|
||
images, labels = dataset_root
|
||
dataset = build(dataset_root)
|
||
assert len(dataset) == 2
|
||
assert dataset.label_file(images / "a.png") == labels / "a.txt"
|
||
|
||
|
||
def test_sample_shapes_and_box_denormalization(dataset_root):
|
||
dataset = build(dataset_root)
|
||
sample = dataset[0]
|
||
|
||
assert sample["image"].shape == (3, 64, 64)
|
||
assert sample["target"]["heatmap"].shape == (2, 16, 16)
|
||
# 归一化框 -> 128px 像素框 (10,10,50,50) -> Resize 到 64 后减半
|
||
assert sample["boxes"][0].tolist() == pytest.approx([5.0, 5.0, 25.0, 25.0], abs=1e-3)
|
||
assert sample["labels"].tolist() == [0]
|
||
assert sample["image_label"].tolist() == [1.0, 0.0]
|
||
|
||
|
||
def test_empty_label_file_is_defect_free(dataset_root):
|
||
sample = build(dataset_root)[1]
|
||
assert sample["boxes"].numel() == 0
|
||
assert sample["target"]["mask"].sum() == 0
|
||
assert sample["image_label"].sum() == 0
|
||
|
||
|
||
def test_missing_label_file_is_defect_free(dataset_root):
|
||
_, labels = dataset_root
|
||
(labels / "b.txt").unlink()
|
||
assert build(dataset_root)[1]["boxes"].numel() == 0
|
||
|
||
|
||
def test_explicit_label_dir(dataset_root, tmp_path):
|
||
images, labels = dataset_root
|
||
other = tmp_path / "yolo_labels"
|
||
other.mkdir()
|
||
(other / "a.txt").write_text("1 0.5 0.5 0.25 0.25\n", encoding="utf-8")
|
||
(other / "b.txt").write_text("", encoding="utf-8")
|
||
|
||
dataset = build(dataset_root, label_dir=other)
|
||
assert dataset[0]["labels"].tolist() == [1]
|
||
|
||
|
||
def test_list_file_controls_order(dataset_root, tmp_path):
|
||
images, _ = dataset_root
|
||
list_file = tmp_path / "val.txt"
|
||
list_file.write_text("images/b.png\nimages/a.png\n", encoding="utf-8")
|
||
|
||
dataset = build(dataset_root, list_file=list_file)
|
||
assert [Path(item).name for item in [dataset[i]["meta"]["path"] for i in range(2)]] == [
|
||
"b.png",
|
||
"a.png",
|
||
]
|
||
|
||
|
||
def test_collate_keeps_boxes_as_list(dataset_root):
|
||
dataset = build(dataset_root)
|
||
batch = collate_fn([dataset[0], dataset[1]])
|
||
|
||
assert batch["images"].shape == (2, 3, 64, 64)
|
||
assert batch["target"]["heatmap"].shape == (2, 2, 16, 16)
|
||
assert isinstance(batch["boxes"], list) and len(batch["boxes"]) == 2
|
||
assert batch["image_labels"].shape == (2, 2)
|
||
|
||
|
||
def test_class_id_out_of_range_raises(dataset_root):
|
||
_, labels = dataset_root
|
||
(labels / "a.txt").write_text("7 0.5 0.5 0.2 0.2\n", encoding="utf-8")
|
||
with pytest.raises(ValueError, match="越界"):
|
||
_ = build(dataset_root)[0]
|
||
|
||
|
||
def test_gray_mode_yields_single_channel(dataset_root):
|
||
images, _ = dataset_root
|
||
transforms = [
|
||
{"name": "Resize", "size": [64, 64]},
|
||
{"name": "ToTensor"},
|
||
{"name": "Normalize", "mean": [0.5], "std": [0.5]},
|
||
]
|
||
dataset = FilmDefectDataset(
|
||
img_dir=images, classes=CLASSES, transforms=transforms, stride=4, gray=True
|
||
)
|
||
assert dataset[0]["image"].shape == (1, 64, 64)
|
||
|
||
|
||
def test_classes_from_data_yaml(dataset_root, tmp_path):
|
||
pytest.importorskip("yaml")
|
||
from datasets import write_data_yaml
|
||
|
||
images, _ = dataset_root
|
||
data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path)
|
||
dataset = FilmDefectDataset(
|
||
img_dir=images, classes=data_yaml, transforms=TRANSFORMS, stride=4
|
||
)
|
||
assert dataset.classes == CLASSES
|
||
|
||
|
||
def test_metrics_on_perfect_prediction():
|
||
from engine.metrics import DetectionMetric
|
||
|
||
metric = DetectionMetric(iou_threshold=0.5)
|
||
boxes = torch.tensor([[0.0, 0.0, 10.0, 10.0]])
|
||
metric.update(boxes, torch.tensor([0]), boxes, torch.tensor([0]))
|
||
result = metric.compute()
|
||
assert result["precision"] == 1.0 and result["recall"] == 1.0
|