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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
f433e5cf59
commit
d6adea5b05
@@ -0,0 +1,4 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
@@ -0,0 +1,126 @@
|
||||
"""无监督聚类的原语。"""
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
from engine.cluster import ( # noqa: E402
|
||||
ClusterResult,
|
||||
cluster_features,
|
||||
cluster_stability,
|
||||
kmeans,
|
||||
normalized_mutual_info,
|
||||
pca_reduce,
|
||||
random_fit_index,
|
||||
silhouette_score,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def blobs():
|
||||
"""4 团分得很开的高斯，外加 200 个纯噪声维度。"""
|
||||
torch.manual_seed(0)
|
||||
centers = torch.tensor([[6.0, 0, 0], [0, 6, 0], [0, 0, 6], [-6, -6, -6]])
|
||||
signal = torch.cat([center + 0.4 * torch.randn(50, 3) for center in centers])
|
||||
noise = 0.01 * torch.randn(200, 200)
|
||||
return signal, torch.cat([signal, noise], dim=1)
|
||||
|
||||
|
||||
def test_kmeans_recovers_well_separated_blobs(blobs):
|
||||
signal, _ = blobs
|
||||
labels, centroids, inertia = kmeans(signal, 4, n_init=5, seed=0)
|
||||
|
||||
assert sorted(torch.bincount(labels).tolist()) == [50, 50, 50, 50]
|
||||
assert centroids.shape == (4, 3)
|
||||
assert inertia > 0
|
||||
assert silhouette_score(signal, labels) > 0.8
|
||||
|
||||
|
||||
def test_kmeans_rejects_bad_arguments(blobs):
|
||||
signal, _ = blobs
|
||||
with pytest.raises(ValueError, match="至少为 2"):
|
||||
kmeans(signal, 1)
|
||||
with pytest.raises(ValueError, match="少于簇数"):
|
||||
kmeans(signal[:3], 4)
|
||||
|
||||
|
||||
def test_nmi_is_invariant_to_label_renaming():
|
||||
labels = [0, 0, 1, 1, 2, 2]
|
||||
assert normalized_mutual_info(labels, labels) == pytest.approx(1.0)
|
||||
assert normalized_mutual_info(labels, [2, 2, 0, 0, 1, 1]) == pytest.approx(1.0)
|
||||
|
||||
torch.manual_seed(0)
|
||||
unrelated = torch.randint(0, 2, (200,)).tolist()
|
||||
assert normalized_mutual_info([0] * 100 + [1] * 100, unrelated) < 0.1
|
||||
|
||||
|
||||
def test_whitening_amplifies_noise_directions(blobs):
|
||||
"""回归测试：pca 白化 + 较大的 dim 会把噪声方向抬到与信号同权，簇随即塌掉。
|
||||
|
||||
这正是默认 `whiten=False` 的原因——早期版本默认开白化，真实数据上
|
||||
最大的簇吃掉了 96% 的样本。
|
||||
"""
|
||||
_, noisy = blobs
|
||||
|
||||
clean = cluster_features(noisy, 4, pca_dim=32, whiten=False)
|
||||
whitened = cluster_features(noisy, 4, pca_dim=32, whiten=True)
|
||||
|
||||
assert clean.silhouette > 0.8
|
||||
assert sorted(clean.sizes) == [50, 50, 50, 50]
|
||||
assert whitened.silhouette < clean.silhouette
|
||||
assert max(whitened.sizes) > max(clean.sizes)
|
||||
|
||||
|
||||
def test_cluster_features_is_insensitive_to_pca_dim_without_whitening(blobs):
|
||||
_, noisy = blobs
|
||||
scores = [cluster_features(noisy, 4, pca_dim=dim).silhouette for dim in (8, 32, 64)]
|
||||
assert all(score > 0.8 for score in scores)
|
||||
|
||||
|
||||
def test_pca_reduce_shapes_and_whitening_effect(blobs):
|
||||
_, noisy = blobs
|
||||
projected = pca_reduce(noisy, 16)
|
||||
assert projected.shape == (200, 16)
|
||||
# 不白化时各主成分方差递减；白化后大致拉平
|
||||
variances = projected.var(dim=0)
|
||||
assert variances[0] > variances[-1]
|
||||
whitened = pca_reduce(noisy, 16, whiten=True).var(dim=0)
|
||||
assert float(whitened.max() / whitened.min()) < float(variances.max() / variances.min())
|
||||
|
||||
|
||||
def test_stability_is_high_on_real_structure_and_low_on_noise(blobs):
|
||||
signal, _ = blobs
|
||||
torch.manual_seed(1)
|
||||
assert cluster_stability(signal, 4, seeds=(0, 7), pca_dim=0) > 0.9
|
||||
assert cluster_stability(torch.randn(300, 16), 4, seeds=(0, 7), pca_dim=0) < 0.5
|
||||
|
||||
|
||||
def test_cluster_features_returns_full_result(blobs):
|
||||
signal, _ = blobs
|
||||
result = cluster_features(signal, 4, pca_dim=0)
|
||||
|
||||
assert isinstance(result, ClusterResult)
|
||||
assert result.labels.shape == (200,)
|
||||
assert result.centroids.shape[0] == 4
|
||||
assert result.distances.shape == (200,)
|
||||
assert sum(result.sizes) == 200
|
||||
|
||||
|
||||
def test_fit_index_only_fits_a_subset_but_assigns_everything(blobs):
|
||||
signal, _ = blobs
|
||||
index = random_fit_index(200, 60, seed=0)
|
||||
|
||||
assert index is not None and len(index) == 60
|
||||
assert random_fit_index(200, 500) is None # 样本本来就少于上限
|
||||
assert random_fit_index(200, 0) is None # 0 表示不采样
|
||||
|
||||
result = cluster_features(signal, 4, pca_dim=0, fit_index=index)
|
||||
assert result.labels.shape == (200,) # 拟合用 60 个，分配仍是全部
|
||||
assert sorted(result.sizes) == [50, 50, 50, 50]
|
||||
|
||||
|
||||
def test_silhouette_subsamples_large_inputs(blobs):
|
||||
signal, _ = blobs
|
||||
labels = kmeans(signal, 4, n_init=3, seed=0)[0]
|
||||
assert silhouette_score(signal, labels, max_samples=50) > 0.5
|
||||
assert torch.isnan(torch.tensor(silhouette_score(signal, torch.zeros_like(labels))))
|
||||
@@ -0,0 +1,48 @@
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("yaml")
|
||||
|
||||
from utils.config import Config, apply_overrides, deep_merge, load_config, parse_override # noqa: E402
|
||||
|
||||
|
||||
def test_deep_merge_is_recursive():
|
||||
merged = deep_merge({"a": {"b": 1, "c": 2}}, {"a": {"c": 3}, "d": 4})
|
||||
assert merged == {"a": {"b": 1, "c": 3}, "d": 4}
|
||||
|
||||
|
||||
def test_parse_override_literals():
|
||||
assert parse_override("train.lr=1e-3") == ("train.lr", 1e-3)
|
||||
assert parse_override("train.amp=False") == ("train.amp", False)
|
||||
assert parse_override("model.backbone.name=tiny_conv") == ("model.backbone.name", "tiny_conv")
|
||||
|
||||
|
||||
def test_apply_overrides_creates_missing_nodes():
|
||||
cfg = apply_overrides({}, ["a.b.c=1"])
|
||||
assert cfg == {"a": {"b": {"c": 1}}}
|
||||
|
||||
|
||||
def test_load_config_with_base(tmp_path):
|
||||
(tmp_path / "base.yaml").write_text("train:\n lr: 0.1\n epochs: 10\n", encoding="utf-8")
|
||||
(tmp_path / "child.yaml").write_text(
|
||||
"_base_: ./base.yaml\ntrain:\n lr: 0.01\nwork_dir: runs/x\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
cfg = load_config(tmp_path / "child.yaml", ["train.epochs=2"])
|
||||
assert isinstance(cfg, Config)
|
||||
assert cfg.train.lr == 0.01
|
||||
assert cfg.train.epochs == 2
|
||||
assert cfg.get_path("work_dir") == "runs/x"
|
||||
|
||||
|
||||
def test_load_config_detects_cycle(tmp_path):
|
||||
(tmp_path / "a.yaml").write_text("_base_: ./b.yaml\n", encoding="utf-8")
|
||||
(tmp_path / "b.yaml").write_text("_base_: ./a.yaml\n", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match="循环引用"):
|
||||
load_config(tmp_path / "a.yaml")
|
||||
|
||||
|
||||
def test_project_configs_are_loadable():
|
||||
cfg = load_config("configs/train/baseline.yaml")
|
||||
assert cfg["model"]["name"] == "defect_net"
|
||||
assert cfg["data"]["train"]["dataset"]["name"] == "film_defect"
|
||||
assert len(cfg["classes"]) > 0
|
||||
@@ -0,0 +1,151 @@
|
||||
"""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
|
||||
@@ -0,0 +1,26 @@
|
||||
"""GPU 守卫：项目禁止在没有 Nvidia GPU 的设备上训练。"""
|
||||
|
||||
import pytest
|
||||
|
||||
from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda
|
||||
|
||||
|
||||
def test_require_cuda_matches_hardware():
|
||||
devices = cuda_devices()
|
||||
if devices:
|
||||
assert require_cuda() == devices
|
||||
else:
|
||||
with pytest.raises(NoNvidiaGpuError):
|
||||
require_cuda()
|
||||
|
||||
|
||||
def test_resolve_device_rejects_cuda_without_gpu():
|
||||
torch = pytest.importorskip("torch")
|
||||
from utils.device import resolve_device
|
||||
|
||||
if torch.cuda.is_available():
|
||||
assert resolve_device("cuda").type == "cuda"
|
||||
else:
|
||||
with pytest.raises(NoNvidiaGpuError):
|
||||
resolve_device("cuda")
|
||||
assert resolve_device("cpu").type == "cpu"
|
||||
@@ -0,0 +1,125 @@
|
||||
"""多类(multilabel) / 单类(multiclass) 由 yaml 的 label_mode 决定，两条路径都要能跑通。"""
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
pytest.importorskip("PIL")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
from datasets import FilmDefectDataset, collate_fn # noqa: E402
|
||||
from engine.losses import DetectionLoss # noqa: E402
|
||||
from engine.metrics import SingleLabelMetric # noqa: E402
|
||||
from models import build_model # noqa: E402
|
||||
from models.defect_net import classification_outputs # noqa: E402
|
||||
|
||||
CLASSES = ["scratch", "bubble"]
|
||||
TRANSFORMS = [{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}]
|
||||
|
||||
|
||||
def make_model(label_mode: str, num_classes: int = 2):
|
||||
return build_model(
|
||||
{
|
||||
"name": "defect_net",
|
||||
"num_classes": num_classes,
|
||||
"label_mode": label_mode,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
||||
"neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
|
||||
"head": {"hidden_channels": 8},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dataset_factory(tmp_path):
|
||||
"""YOLO 目录：a.png 上 scratch 1 个框、bubble 2 个框；b.png 无缺陷。"""
|
||||
images = tmp_path / "images"
|
||||
labels = tmp_path / "labels"
|
||||
images.mkdir()
|
||||
labels.mkdir()
|
||||
for name in ("a.png", "b.png"):
|
||||
Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / name)
|
||||
|
||||
(labels / "a.txt").write_text(
|
||||
"0 0.1875 0.1875 0.25 0.25\n" # scratch
|
||||
"1 0.5000 0.1875 0.25 0.25\n" # bubble
|
||||
"1 0.2500 0.5000 0.25 0.25\n", # bubble
|
||||
encoding="utf-8",
|
||||
)
|
||||
(labels / "b.txt").write_text("", encoding="utf-8")
|
||||
|
||||
def build(label_mode: str):
|
||||
return FilmDefectDataset(
|
||||
img_dir=images,
|
||||
classes=CLASSES,
|
||||
transforms=TRANSFORMS,
|
||||
stride=2,
|
||||
label_mode=label_mode,
|
||||
)
|
||||
|
||||
return build
|
||||
|
||||
|
||||
def test_multilabel_dataset_yields_multi_hot(dataset_factory):
|
||||
dataset = dataset_factory("multilabel")
|
||||
assert dataset[0]["image_label"].tolist() == [1.0, 1.0]
|
||||
assert dataset[1]["image_label"].tolist() == [0.0, 0.0]
|
||||
|
||||
|
||||
def test_multiclass_dataset_yields_exclusive_index(dataset_factory):
|
||||
dataset = dataset_factory("multiclass")
|
||||
# bubble 出现 2 次、scratch 1 次 -> 取 bubble(下标 1)
|
||||
assert dataset[0]["image_label"].item() == 1
|
||||
assert dataset[0]["image_label"].dtype == torch.long
|
||||
# 无缺陷 -> 背景下标 = len(classes)
|
||||
assert dataset[1]["image_label"].item() == len(CLASSES)
|
||||
|
||||
|
||||
def test_collate_shapes_differ_by_mode(dataset_factory):
|
||||
multilabel = collate_fn([dataset_factory("multilabel")[i] for i in range(2)])
|
||||
multiclass = collate_fn([dataset_factory("multiclass")[i] for i in range(2)])
|
||||
assert multilabel["image_labels"].shape == (2, len(CLASSES))
|
||||
assert multiclass["image_labels"].shape == (2,)
|
||||
# 定位监督与 label_mode 无关
|
||||
assert torch.equal(multilabel["target"]["heatmap"], multiclass["target"]["heatmap"])
|
||||
|
||||
|
||||
def test_classification_head_width():
|
||||
assert classification_outputs(5, "multilabel") == 5
|
||||
assert classification_outputs(5, "multiclass") == 6
|
||||
assert make_model("multilabel").cls_head.num_outputs == 2
|
||||
assert make_model("multiclass").cls_head.num_outputs == 3
|
||||
|
||||
|
||||
def test_invalid_label_mode_raises():
|
||||
with pytest.raises(ValueError, match="label_mode"):
|
||||
make_model("single")
|
||||
with pytest.raises(ValueError, match="label_mode"):
|
||||
DetectionLoss(label_mode="single")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("label_mode", ["multilabel", "multiclass"])
|
||||
def test_end_to_end_loss_backward(dataset_factory, label_mode):
|
||||
model = make_model(label_mode)
|
||||
criterion = DetectionLoss(label_mode=label_mode)
|
||||
dataset = dataset_factory(label_mode)
|
||||
batch = collate_fn([dataset[0], dataset[1]])
|
||||
|
||||
losses = criterion(model(batch["images"]), batch)
|
||||
losses["total"].backward()
|
||||
|
||||
assert set(losses) == {"heatmap", "wh", "offset", "cls", "total"}
|
||||
assert torch.isfinite(losses["total"])
|
||||
assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.cls_head.parameters())
|
||||
|
||||
|
||||
def test_single_label_metric_counts_only_present_classes():
|
||||
metric = SingleLabelMetric(num_outputs=3)
|
||||
logits = torch.tensor([[9.0, 0.0, 0.0], [0.0, 9.0, 0.0], [9.0, 0.0, 0.0]])
|
||||
metric.update(logits, torch.tensor([0, 1, 1]))
|
||||
result = metric.compute()
|
||||
|
||||
assert result["cls_accuracy"] == pytest.approx(2 / 3)
|
||||
# 类别 2 没出现，不参与 macro 平均
|
||||
assert result["cls_f1"] == pytest.approx((2 / 3 + 2 / 3) / 2)
|
||||
@@ -0,0 +1,92 @@
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
from models import build_model # noqa: E402
|
||||
from models.postprocess import decode_heatmap # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_cfg():
|
||||
return {
|
||||
"name": "defect_net",
|
||||
"num_classes": 3,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]},
|
||||
"neck": {"name": "upsample", "channels": 16, "num_upsamples": 2},
|
||||
"head": {"hidden_channels": 16},
|
||||
}
|
||||
|
||||
|
||||
def test_forward_shapes_and_stride(model_cfg):
|
||||
model = build_model(model_cfg).eval()
|
||||
# stem(2) * 2 个下采样 stage = 8，再上采样 4 倍 => stride 2
|
||||
assert model.out_stride == 2
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(2, 3, 64, 64))
|
||||
|
||||
size = 64 // model.out_stride
|
||||
assert out["heatmap"].shape == (2, 3, size, size)
|
||||
assert out["wh"].shape == (2, 2, size, size)
|
||||
assert out["offset"].shape == (2, 2, size, size)
|
||||
assert out["logits"].shape == (2, 3)
|
||||
|
||||
|
||||
def test_backbone_is_convolution_only(model_cfg):
|
||||
model = build_model(model_cfg)
|
||||
forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM, torch.nn.GRU)
|
||||
assert not any(isinstance(m, forbidden) for m in model.modules())
|
||||
|
||||
|
||||
def test_without_classification_head(model_cfg):
|
||||
model = build_model({**model_cfg, "with_classification": False}).eval()
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(1, 3, 64, 64))
|
||||
assert "logits" not in out
|
||||
|
||||
|
||||
def test_decode_heatmap_recovers_planted_box():
|
||||
heatmap = torch.full((1, 2, 32, 32), -10.0)
|
||||
heatmap[0, 1, 8, 10] = 10.0
|
||||
wh = torch.zeros(1, 2, 32, 32)
|
||||
wh[0, :, 8, 10] = torch.tensor([4.0, 6.0])
|
||||
offset = torch.zeros(1, 2, 32, 32)
|
||||
|
||||
(result,) = decode_heatmap(heatmap, wh, offset, stride=4, topk=5, score_threshold=0.5)
|
||||
|
||||
assert result["labels"].tolist() == [1]
|
||||
assert result["boxes"][0].tolist() == pytest.approx([32.0, 20.0, 48.0, 44.0])
|
||||
|
||||
|
||||
def test_losses_decrease_on_a_single_batch(model_cfg):
|
||||
from engine.losses import DetectionLoss
|
||||
|
||||
torch.manual_seed(0)
|
||||
model = build_model(model_cfg)
|
||||
criterion = DetectionLoss()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
|
||||
|
||||
images = torch.randn(2, 3, 64, 64)
|
||||
size = 64 // model.out_stride
|
||||
target = {
|
||||
"heatmap": torch.zeros(2, 3, size, size),
|
||||
"wh": torch.zeros(2, 2, size, size),
|
||||
"offset": torch.zeros(2, 2, size, size),
|
||||
"mask": torch.zeros(2, size, size),
|
||||
}
|
||||
target["heatmap"][0, 1, 5, 5] = 1.0
|
||||
target["mask"][0, 5, 5] = 1.0
|
||||
target["wh"][0, :, 5, 5] = 3.0
|
||||
batch = {"target": target, "image_labels": torch.zeros(2, 3)}
|
||||
|
||||
first = last = None
|
||||
for step in range(5):
|
||||
losses = criterion(model(images), batch)
|
||||
optimizer.zero_grad()
|
||||
losses["total"].backward()
|
||||
optimizer.step()
|
||||
total = float(losses["total"].detach())
|
||||
first = total if step == 0 else first
|
||||
last = total
|
||||
|
||||
assert last < first
|
||||
@@ -0,0 +1,70 @@
|
||||
"""推理 profiler 脚本：CPU 上也要能跑出耗时结果（有 GPU 时额外给出显存峰值）。"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
from models import build_model # noqa: E402
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "profile_inference.py"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def script():
|
||||
spec = importlib.util.spec_from_file_location("profile_inference", SCRIPT)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model():
|
||||
return build_model(
|
||||
{
|
||||
"name": "defect_net",
|
||||
"num_classes": 2,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
||||
"neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
|
||||
"head": {"hidden_channels": 8},
|
||||
}
|
||||
).eval()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def args():
|
||||
return argparse.Namespace(iters=2, warmup=1, amp=False, postprocess=True, rows=3, trace=None)
|
||||
|
||||
|
||||
def test_measure_reports_latency_and_memory(script, model, args):
|
||||
images = torch.randn(2, 3, 32, 32)
|
||||
result = script.measure(model, images, args)
|
||||
|
||||
assert result["batch_size"] == 2
|
||||
assert result["latency_ms"] > 0
|
||||
assert result["throughput_fps"] > 0
|
||||
if torch.cuda.is_available():
|
||||
assert result["peak_allocated_mb"] > 0
|
||||
assert result["peak_reserved_mb"] >= result["peak_allocated_mb"]
|
||||
else:
|
||||
assert "peak_allocated_mb" not in result
|
||||
|
||||
|
||||
def test_profile_exports_trace(script, model, args, tmp_path, capsys):
|
||||
args.trace = str(tmp_path / "trace.json")
|
||||
script.profile_once(model, torch.randn(1, 3, 32, 32), args)
|
||||
|
||||
assert Path(args.trace).exists()
|
||||
out = capsys.readouterr().out
|
||||
assert "显存占用 Top" in out and "耗时 Top" in out
|
||||
|
||||
|
||||
def test_postprocess_toggle_runs_both_paths(script, model, args):
|
||||
images = torch.randn(1, 3, 32, 32)
|
||||
args.postprocess = False
|
||||
script.run_once(model, images, args)
|
||||
args.postprocess = True
|
||||
script.run_once(model, images, args)
|
||||
@@ -0,0 +1,29 @@
|
||||
import pytest
|
||||
|
||||
from utils.registry import Registry
|
||||
|
||||
|
||||
def test_register_and_build():
|
||||
registry = Registry("demo")
|
||||
|
||||
@registry.register("adder")
|
||||
class Adder:
|
||||
def __init__(self, base: int = 0):
|
||||
self.base = base
|
||||
|
||||
assert "adder" in registry
|
||||
assert len(registry) == 1
|
||||
assert registry.build({"name": "adder", "base": 3}).base == 3
|
||||
assert registry.build({"name": "adder"}, base=5).base == 5
|
||||
|
||||
|
||||
def test_duplicate_and_missing():
|
||||
registry = Registry("demo")
|
||||
registry.register("x")(lambda: None)
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
registry.register("x")(lambda: None)
|
||||
with pytest.raises(KeyError, match="找不到"):
|
||||
registry.get("y")
|
||||
with pytest.raises(KeyError):
|
||||
registry.build({"base": 1})
|
||||
@@ -0,0 +1,148 @@
|
||||
"""SimCLR 自监督：损失的数学性质 + 聚类脚本的 simclr 分支端到端。"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
pytest.importorskip("PIL")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
from engine.simclr import SimCLRTrainer, nt_xent_loss # noqa: E402
|
||||
from models.heads.projection import ProjectionHead # noqa: E402
|
||||
from models.registry import BACKBONES # noqa: E402
|
||||
from utils.device import NoNvidiaGpuError # noqa: E402
|
||||
|
||||
|
||||
def test_nt_xent_is_minimal_when_views_agree():
|
||||
"""两个视图完全一致且各样本互相正交时，损失应接近理论下界 log(2B−1) 的反面——
|
||||
即正样本对独占概率质量，损失趋近 0。"""
|
||||
z = torch.eye(4) # 4 个互相正交的样本
|
||||
aligned = nt_xent_loss(z, z.clone(), temperature=0.05)
|
||||
misaligned = nt_xent_loss(z, z.roll(1, dims=0), temperature=0.05)
|
||||
|
||||
assert aligned < 0.1
|
||||
assert misaligned > aligned
|
||||
|
||||
|
||||
def test_nt_xent_upper_bound_at_chance_level():
|
||||
"""所有表征相同时无法区分正负样本，损失 = log(2B−1)。"""
|
||||
z = torch.ones(8, 16)
|
||||
loss = float(nt_xent_loss(z, z.clone(), temperature=0.2))
|
||||
assert loss == pytest.approx(float(torch.log(torch.tensor(2 * 8 - 1))), abs=1e-4)
|
||||
|
||||
|
||||
def test_nt_xent_is_symmetric_and_temperature_sensitive():
|
||||
torch.manual_seed(0)
|
||||
z1, z2 = torch.randn(6, 8), torch.randn(6, 8)
|
||||
assert float(nt_xent_loss(z1, z2, 0.2)) == pytest.approx(float(nt_xent_loss(z2, z1, 0.2)), abs=1e-5)
|
||||
# 温度越低，随机表征的损失越大（softmax 更尖，惩罚更重）
|
||||
assert nt_xent_loss(z1, z2, 0.05) > nt_xent_loss(z1, z2, 0.5)
|
||||
|
||||
|
||||
def test_nt_xent_gradient_pulls_positive_pair_together():
|
||||
torch.manual_seed(0)
|
||||
z1 = torch.randn(4, 8, requires_grad=True)
|
||||
z2 = torch.randn(4, 8, requires_grad=True)
|
||||
before = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
|
||||
|
||||
optimizer = torch.optim.SGD([z1, z2], lr=1.0)
|
||||
for _ in range(20):
|
||||
optimizer.zero_grad()
|
||||
nt_xent_loss(z1, z2, 0.2).backward()
|
||||
optimizer.step()
|
||||
|
||||
after = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
|
||||
assert after > before
|
||||
|
||||
|
||||
def test_projection_head_accepts_feature_map_and_vector():
|
||||
head = ProjectionHead(16, hidden_channels=32, out_channels=8)
|
||||
assert head(torch.randn(4, 16, 5, 5)).shape == (4, 8)
|
||||
assert head(torch.randn(4, 16)).shape == (4, 8)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tiny_image_dir(tmp_path):
|
||||
root = tmp_path / "imgs"
|
||||
root.mkdir()
|
||||
rng = np.random.default_rng(0)
|
||||
for i in range(8):
|
||||
Image.fromarray(rng.integers(0, 255, (40, 60, 3), dtype=np.uint8)).save(root / f"{i}.png")
|
||||
return root
|
||||
|
||||
|
||||
def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
|
||||
|
||||
dataset = UnlabeledImageDataset(
|
||||
tiny_image_dir,
|
||||
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
|
||||
views=2,
|
||||
view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.5, 1.0]}],
|
||||
)
|
||||
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
|
||||
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
|
||||
|
||||
trainer = SimCLRTrainer(
|
||||
backbone,
|
||||
loader,
|
||||
{"epochs": 2, "amp": False, "log_interval": 100, "projection": {"hidden_channels": 16,
|
||||
"out_channels": 8}},
|
||||
torch.device("cpu"),
|
||||
tmp_path / "run",
|
||||
)
|
||||
path = trainer.fit()
|
||||
|
||||
assert path.exists()
|
||||
state = torch.load(path, map_location="cpu", weights_only=False)
|
||||
# 只存骨架，不含投影头
|
||||
assert all("projector" not in key and "mlp" not in key for key in state["model"])
|
||||
backbone.load_state_dict(state["model"]) # 能被干净地加载回去
|
||||
|
||||
|
||||
def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
|
||||
"""method=simclr 全流程：训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
|
||||
sys.path.insert(0, "scripts")
|
||||
import cluster_unlabeled as cu
|
||||
|
||||
monkeypatch.setattr(cu, "require_cuda", lambda: ["FAKE-GPU"])
|
||||
work_dir = tmp_path / "run"
|
||||
monkeypatch.setattr(
|
||||
sys, "argv",
|
||||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||||
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
|
||||
"device=cpu", "method=simclr", "loader.num_workers=0", "loader.batch_size=4",
|
||||
"simclr.epochs=1", "simclr.batch_size=4", "simclr.amp=False",
|
||||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
|
||||
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
|
||||
"num_clusters=2", "n_init=2"],
|
||||
)
|
||||
|
||||
assert cu.main() == 0
|
||||
assert (work_dir / "simclr_backbone.pth").exists()
|
||||
assert (work_dir / "assignments.csv").exists()
|
||||
assert (work_dir / "summary.json").exists()
|
||||
assert (work_dir / "clusters" / "cluster_00").is_dir()
|
||||
|
||||
|
||||
def test_simclr_branch_is_blocked_without_gpu(tiny_image_dir, tmp_path, monkeypatch):
|
||||
"""无 Nvidia GPU 时 simclr 分支必须被守卫拦下（pretrained 分支不受影响）。"""
|
||||
sys.path.insert(0, "scripts")
|
||||
import cluster_unlabeled as cu
|
||||
|
||||
def refuse():
|
||||
raise NoNvidiaGpuError("没有 Nvidia GPU")
|
||||
|
||||
monkeypatch.setattr(cu, "require_cuda", refuse)
|
||||
monkeypatch.setattr(
|
||||
sys, "argv",
|
||||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||||
"--work-dir", str(tmp_path / "run"), "--root", str(tiny_image_dir),
|
||||
"--set", "device=cpu", "method=simclr", "loader.num_workers=0"],
|
||||
)
|
||||
assert cu.main() == 1
|
||||
@@ -0,0 +1,42 @@
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("numpy")
|
||||
pytest.importorskip("torch")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
|
||||
from datasets.target_encoder import CenterTargetEncoder, gaussian_radius # noqa: E402
|
||||
|
||||
|
||||
def test_gaussian_radius_is_positive_and_scales():
|
||||
assert gaussian_radius(40, 40) > gaussian_radius(10, 10) > 0
|
||||
|
||||
|
||||
def test_encoder_marks_single_center():
|
||||
encoder = CenterTargetEncoder(num_classes=3, stride=4)
|
||||
boxes = np.array([[40.0, 40.0, 80.0, 80.0]])
|
||||
labels = np.array([1])
|
||||
|
||||
target = encoder(boxes, labels, (256, 256))
|
||||
|
||||
assert target["heatmap"].shape == (3, 64, 64)
|
||||
assert target["mask"].sum() == 1
|
||||
# 中心 (60, 60) / stride 4 = (15, 15)
|
||||
assert target["mask"][15, 15] == 1
|
||||
assert target["heatmap"][1].max() == pytest.approx(1.0)
|
||||
assert target["heatmap"][0].max() == 0
|
||||
assert target["wh"][:, 15, 15].tolist() == pytest.approx([10.0, 10.0])
|
||||
|
||||
|
||||
def test_encoder_skips_invalid_and_out_of_range_boxes():
|
||||
encoder = CenterTargetEncoder(num_classes=2, stride=4)
|
||||
boxes = np.array([[10.0, 10.0, 5.0, 5.0], [900.0, 900.0, 950.0, 950.0]])
|
||||
target = encoder(boxes, np.array([0, 1]), (256, 256))
|
||||
assert target["mask"].sum() == 0
|
||||
|
||||
|
||||
def test_encoder_respects_max_objects():
|
||||
encoder = CenterTargetEncoder(num_classes=1, stride=4, max_objects=2)
|
||||
boxes = np.array([[i * 20.0, 0.0, i * 20.0 + 16, 16.0] for i in range(5)])
|
||||
target = encoder(boxes, np.zeros(5, dtype=np.int64), (256, 256))
|
||||
assert target["mask"].sum() == 2
|
||||
@@ -0,0 +1,137 @@
|
||||
"""train.py / eval.py 的端到端冒烟：YOLO 数据目录 -> 训练 -> 评估。
|
||||
|
||||
覆盖入口脚本的接线（classes 解析、label_mode 下发、GPU 守卫），
|
||||
所以在 CPU 上跑，`require_cuda` 用 monkeypatch 放行。
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch")
|
||||
pytest.importorskip("yaml")
|
||||
pytest.importorskip("PIL")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
import yaml # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
import eval as eval_mod # noqa: E402
|
||||
import train as train_mod # noqa: E402
|
||||
from datasets import write_data_yaml # noqa: E402
|
||||
from utils.device import NoNvidiaGpuError # noqa: E402
|
||||
|
||||
CLASSES = ["scratch", "bubble"]
|
||||
TRANSFORMS = [
|
||||
{"name": "Resize", "size": [32, 32]},
|
||||
{"name": "ToTensor"},
|
||||
{"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project(tmp_path):
|
||||
"""造一个最小的 YOLO 数据集 + 配置，classes 用 data.yaml 的路径给。"""
|
||||
for split in ("train", "val"):
|
||||
images = tmp_path / "images" / split
|
||||
labels = tmp_path / "labels" / split
|
||||
images.mkdir(parents=True)
|
||||
labels.mkdir(parents=True)
|
||||
for i in range(2):
|
||||
Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / f"{i}.png")
|
||||
content = "0 0.5 0.5 0.25 0.25\n" if i else ""
|
||||
(labels / f"{i}.txt").write_text(content, encoding="utf-8")
|
||||
|
||||
data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path)
|
||||
|
||||
def dataset(split):
|
||||
return {
|
||||
"dataset": {"name": "film_defect", "img_dir": str(tmp_path / "images" / split),
|
||||
"transforms": TRANSFORMS},
|
||||
"loader": {"batch_size": 2, "num_workers": 0, "shuffle": split == "train"},
|
||||
}
|
||||
|
||||
config = {
|
||||
"classes": str(data_yaml),
|
||||
"label_mode": "multilabel",
|
||||
"image_size": [32, 32],
|
||||
"device": "cpu",
|
||||
"seed": 0,
|
||||
"model": {
|
||||
"name": "defect_net",
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
||||
"neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
|
||||
"head": {"hidden_channels": 8},
|
||||
},
|
||||
"data": {"train": dataset("train"), "val": dataset("val")},
|
||||
"train": {
|
||||
"epochs": 1,
|
||||
"amp": False,
|
||||
"log_interval": 10,
|
||||
"optimizer": {"name": "adamw", "lr": 0.001},
|
||||
"swanlab": {"enabled": False},
|
||||
},
|
||||
}
|
||||
config_path = tmp_path / "config.yaml"
|
||||
config_path.write_text(yaml.safe_dump(config, allow_unicode=True), encoding="utf-8")
|
||||
return config_path, tmp_path
|
||||
|
||||
|
||||
def run(module, argv, monkeypatch):
|
||||
monkeypatch.setattr(sys, "argv", argv)
|
||||
return module.main()
|
||||
|
||||
|
||||
def test_train_then_eval(project, monkeypatch, tmp_path):
|
||||
config_path, root = project
|
||||
work_dir = tmp_path / "runs"
|
||||
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
|
||||
|
||||
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
||||
monkeypatch) == 0
|
||||
assert (work_dir / "best.pth").exists()
|
||||
assert (work_dir / "config.yaml").exists()
|
||||
|
||||
metrics_path = tmp_path / "metrics.json"
|
||||
assert run(
|
||||
eval_mod,
|
||||
["eval.py", "--config", str(config_path), "--checkpoint", str(work_dir / "best.pth"),
|
||||
"--set", "device=cpu", "--out", str(metrics_path)],
|
||||
monkeypatch,
|
||||
) == 0
|
||||
assert metrics_path.exists()
|
||||
|
||||
|
||||
def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
|
||||
"""classes 写成 data.yaml 路径时，num_classes 必须是 2，而不是路径字符串的长度。"""
|
||||
config_path, _ = project
|
||||
work_dir = tmp_path / "runs2"
|
||||
captured = {}
|
||||
|
||||
real_build = train_mod.build_model
|
||||
|
||||
def spy(cfg):
|
||||
captured["cfg"] = dict(cfg)
|
||||
return real_build(cfg)
|
||||
|
||||
monkeypatch.setattr(train_mod, "build_model", spy)
|
||||
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
|
||||
|
||||
run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], monkeypatch)
|
||||
|
||||
assert captured["cfg"]["num_classes"] == len(CLASSES)
|
||||
assert captured["cfg"]["label_mode"] == "multilabel"
|
||||
|
||||
|
||||
def test_training_is_blocked_without_nvidia_gpu(project, monkeypatch, tmp_path):
|
||||
config_path, _ = project
|
||||
|
||||
def refuse():
|
||||
raise NoNvidiaGpuError("没有 Nvidia GPU")
|
||||
|
||||
monkeypatch.setattr(train_mod, "require_cuda", refuse)
|
||||
assert run(
|
||||
train_mod,
|
||||
["train.py", "--config", str(config_path), "--work-dir", str(tmp_path / "runs3")],
|
||||
monkeypatch,
|
||||
) == 1
|
||||
@@ -0,0 +1,131 @@
|
||||
"""无标注数据集与聚类相关的预处理。"""
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,63 @@
|
||||
"""SwanLab 可视化封装：可用时记录，不可用时必须静默降级、不影响训练。"""
|
||||
|
||||
import builtins
|
||||
|
||||
import pytest
|
||||
|
||||
from engine.visualizer import SwanLabLogger
|
||||
|
||||
|
||||
def test_disabled_logger_is_noop(tmp_path):
|
||||
logger = SwanLabLogger(enabled=False, work_dir=tmp_path)
|
||||
assert logger.enabled is False
|
||||
logger.log({"loss": 1.0}, step=1)
|
||||
logger.finish()
|
||||
|
||||
|
||||
def test_mode_disabled_is_treated_as_off(tmp_path):
|
||||
assert SwanLabLogger(work_dir=tmp_path, mode="disabled").enabled is False
|
||||
|
||||
|
||||
def test_invalid_mode_raises(tmp_path):
|
||||
with pytest.raises(ValueError, match="swanlab mode"):
|
||||
SwanLabLogger(work_dir=tmp_path, mode="cloud")
|
||||
|
||||
|
||||
def test_missing_swanlab_degrades_gracefully(tmp_path, monkeypatch):
|
||||
real_import = builtins.__import__
|
||||
|
||||
def fake_import(name, *args, **kwargs):
|
||||
if name == "swanlab":
|
||||
raise ImportError("no swanlab")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", fake_import)
|
||||
logger = SwanLabLogger(work_dir=tmp_path)
|
||||
assert logger.enabled is False
|
||||
logger.log({"loss": 1.0}) # 不应抛异常
|
||||
|
||||
|
||||
def test_offline_run_records_metrics(tmp_path):
|
||||
pytest.importorskip("swanlab")
|
||||
logger = SwanLabLogger(
|
||||
work_dir=tmp_path, project="mte-tests", mode="offline", config={"lr": 0.1}
|
||||
)
|
||||
if not logger.enabled: # 环境不允许起 run 时跳过，不让测试变脆
|
||||
pytest.skip("swanlab 无法在当前环境初始化")
|
||||
logger.log({"total": 1.0, "lr": 0.1}, step=1, prefix="train/")
|
||||
logger.finish()
|
||||
assert (tmp_path / "swanlog").exists()
|
||||
|
||||
|
||||
def test_log_failure_disables_instead_of_raising(tmp_path):
|
||||
logger = SwanLabLogger(enabled=False, work_dir=tmp_path)
|
||||
logger.enabled = True # 伪造已启用但底层不可用的状态
|
||||
|
||||
class Broken:
|
||||
@staticmethod
|
||||
def log(*args, **kwargs):
|
||||
raise RuntimeError("network down")
|
||||
|
||||
logger._swanlab = Broken()
|
||||
logger.log({"loss": 1.0}, step=1)
|
||||
assert logger.enabled is False
|
||||
@@ -0,0 +1,118 @@
|
||||
"""YOLO 标注格式的解析与写出。"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("numpy")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
|
||||
from datasets.yolo import ( # noqa: E402
|
||||
denormalize,
|
||||
format_label_lines,
|
||||
label_path_for,
|
||||
list_images,
|
||||
load_classes,
|
||||
parse_label_file,
|
||||
write_data_yaml,
|
||||
)
|
||||
|
||||
|
||||
def test_label_path_swaps_images_for_labels():
|
||||
img_dir = Path("/data/film/images/train")
|
||||
assert label_path_for(img_dir / "0001.png", img_dir) == Path("/data/film/labels/train/0001.txt")
|
||||
|
||||
|
||||
def test_label_path_keeps_subdirs_and_honours_explicit_dir():
|
||||
img_dir = Path("/data/film/images/train")
|
||||
assert label_path_for(img_dir / "cam1/0001.png", img_dir) == Path(
|
||||
"/data/film/labels/train/cam1/0001.txt"
|
||||
)
|
||||
assert label_path_for(img_dir / "0001.png", img_dir, Path("/tmp/lab")) == Path("/tmp/lab/0001.txt")
|
||||
|
||||
|
||||
def test_parse_center_format(tmp_path):
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("0 0.5 0.5 0.4 0.2\n# 注释\n\n1 0.25 0.25 0.1 0.1\n", encoding="utf-8")
|
||||
|
||||
boxes, labels = parse_label_file(path)
|
||||
|
||||
assert labels.tolist() == [0, 1]
|
||||
assert boxes[0].tolist() == pytest.approx([0.3, 0.4, 0.7, 0.6])
|
||||
assert denormalize(boxes, 100, 200)[0].tolist() == pytest.approx([30.0, 80.0, 70.0, 120.0])
|
||||
|
||||
|
||||
def test_missing_file_means_no_defect(tmp_path):
|
||||
boxes, labels = parse_label_file(tmp_path / "nope.txt")
|
||||
assert boxes.shape == (0, 4) and labels.shape == (0,)
|
||||
|
||||
|
||||
def test_polygon_rows_fall_back_to_bounding_box(tmp_path):
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("2 0.1 0.1 0.4 0.2 0.3 0.5\n", encoding="utf-8")
|
||||
|
||||
boxes, labels = parse_label_file(path)
|
||||
|
||||
assert labels.tolist() == [2]
|
||||
assert boxes[0].tolist() == pytest.approx([0.1, 0.1, 0.4, 0.5])
|
||||
|
||||
|
||||
def test_coordinates_are_clamped_and_degenerate_rows_dropped(tmp_path):
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("0 0.5 0.5 2.0 2.0\n0 0.5 0.5 0.0 0.0\n", encoding="utf-8")
|
||||
|
||||
boxes, labels = parse_label_file(path)
|
||||
|
||||
assert len(labels) == 1
|
||||
assert boxes[0].tolist() == pytest.approx([0.0, 0.0, 1.0, 1.0])
|
||||
|
||||
|
||||
def test_malformed_rows_report_file_and_line(tmp_path):
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("0 0.5 0.5 0.4\n", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match=r"a\.txt:1"):
|
||||
parse_label_file(path)
|
||||
|
||||
path.write_text("scratch 0.5 0.5 0.4 0.2\n", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match="不是合法的 YOLO 标注行"):
|
||||
parse_label_file(path)
|
||||
|
||||
|
||||
def test_class_id_range_is_checked(tmp_path):
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("3 0.5 0.5 0.4 0.2\n", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match="越界"):
|
||||
parse_label_file(path, num_classes=2)
|
||||
|
||||
|
||||
def test_format_lines_roundtrips_pixel_boxes(tmp_path):
|
||||
boxes = np.array([[10.0, 20.0, 50.0, 60.0]])
|
||||
lines = format_label_lines(boxes, [1], width=100, height=200)
|
||||
assert lines == ["1 0.300000 0.200000 0.400000 0.200000"]
|
||||
|
||||
path = tmp_path / "a.txt"
|
||||
path.write_text("\n".join(lines), encoding="utf-8")
|
||||
parsed, labels = parse_label_file(path)
|
||||
assert labels.tolist() == [1]
|
||||
assert denormalize(parsed, 100, 200)[0].tolist() == pytest.approx([10.0, 20.0, 50.0, 60.0])
|
||||
|
||||
|
||||
def test_data_yaml_roundtrip(tmp_path):
|
||||
pytest.importorskip("yaml")
|
||||
path = write_data_yaml(tmp_path / "data.yaml", ["scratch", "bubble"], root=tmp_path)
|
||||
assert load_classes(path) == ["scratch", "bubble"]
|
||||
assert load_classes(["a", "b"]) == ["a", "b"]
|
||||
|
||||
|
||||
def test_list_images_supports_dir_and_list_file(tmp_path):
|
||||
images = tmp_path / "images"
|
||||
images.mkdir()
|
||||
for name in ("b.png", "a.jpg", "notes.txt"):
|
||||
(images / name).write_bytes(b"")
|
||||
|
||||
assert [p.name for p in list_images(images)] == ["a.jpg", "b.png"]
|
||||
|
||||
list_file = tmp_path / "train.txt"
|
||||
list_file.write_text("images/b.png\nimages/a.jpg\n", encoding="utf-8")
|
||||
assert [p.name for p in list_images(images, list_file)] == ["b.png", "a.jpg"]
|
||||
Reference in New Issue
Block a user