From d6adea5b05ebb237bdd1fa8bd14004e3df12bbda Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 2 Sep 2026 17:55:47 +0800 Subject: [PATCH] =?UTF-8?q?test:=20=E5=8D=95=E5=85=83=E6=B5=8B=E8=AF=95?= =?UTF-8?q?=E4=B8=8E=E7=AB=AF=E5=88=B0=E7=AB=AF=E5=86=92=E7=83=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 88 个测试覆盖配置继承、注册表、GPU 守卫、YOLO 格式解析、数据集、模型形状、 标签模式、聚类原语与 SimCLR 损失，以及 train.py / eval.py / 聚类脚本的入口接线。 含两条回归测试： - PCA 白化放大噪声方向导致簇塌陷（默认 whiten=False 的依据） - classes 写成 data.yaml 路径时 num_classes 必须是类别数而非路径字符串长度 Co-Authored-By: Claude Opus 5 --- tests/conftest.py | 4 + tests/test_cluster.py | 126 ++++++++++++++++++++++++++ tests/test_config.py | 48 ++++++++++ tests/test_dataset.py | 151 ++++++++++++++++++++++++++++++++ tests/test_device_guard.py | 26 ++++++ tests/test_label_mode.py | 125 ++++++++++++++++++++++++++ tests/test_models.py | 92 +++++++++++++++++++ tests/test_profile_inference.py | 70 +++++++++++++++ tests/test_registry.py | 29 ++++++ tests/test_simclr.py | 148 +++++++++++++++++++++++++++++++ tests/test_target_encoder.py | 42 +++++++++ tests/test_train_smoke.py | 137 +++++++++++++++++++++++++++++ tests/test_unlabeled_dataset.py | 131 +++++++++++++++++++++++++++ tests/test_visualizer.py | 63 +++++++++++++ tests/test_yolo_format.py | 118 +++++++++++++++++++++++++ 15 files changed, 1310 insertions(+) create mode 100644 tests/conftest.py create mode 100644 tests/test_cluster.py create mode 100644 tests/test_config.py create mode 100644 tests/test_dataset.py create mode 100644 tests/test_device_guard.py create mode 100644 tests/test_label_mode.py create mode 100644 tests/test_models.py create mode 100644 tests/test_profile_inference.py create mode 100644 tests/test_registry.py create mode 100644 tests/test_simclr.py create mode 100644 tests/test_target_encoder.py create mode 100644 tests/test_train_smoke.py create mode 100644 tests/test_unlabeled_dataset.py create mode 100644 tests/test_visualizer.py create mode 100644 tests/test_yolo_format.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..5801bea --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,4 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) diff --git a/tests/test_cluster.py b/tests/test_cluster.py new file mode 100644 index 0000000..4bcc0d5 --- /dev/null +++ b/tests/test_cluster.py @@ -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)))) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..93d5a46 --- /dev/null +++ b/tests/test_config.py @@ -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 diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 0000000..f1d2ac3 --- /dev/null +++ b/tests/test_dataset.py @@ -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 diff --git a/tests/test_device_guard.py b/tests/test_device_guard.py new file mode 100644 index 0000000..3a55994 --- /dev/null +++ b/tests/test_device_guard.py @@ -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" diff --git a/tests/test_label_mode.py b/tests/test_label_mode.py new file mode 100644 index 0000000..b38d9db --- /dev/null +++ b/tests/test_label_mode.py @@ -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) diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..26657cb --- /dev/null +++ b/tests/test_models.py @@ -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 diff --git a/tests/test_profile_inference.py b/tests/test_profile_inference.py new file mode 100644 index 0000000..474656b --- /dev/null +++ b/tests/test_profile_inference.py @@ -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) diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000..d6d71f6 --- /dev/null +++ b/tests/test_registry.py @@ -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}) diff --git a/tests/test_simclr.py b/tests/test_simclr.py new file mode 100644 index 0000000..40f0223 --- /dev/null +++ b/tests/test_simclr.py @@ -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 diff --git a/tests/test_target_encoder.py b/tests/test_target_encoder.py new file mode 100644 index 0000000..9209674 --- /dev/null +++ b/tests/test_target_encoder.py @@ -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 diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py new file mode 100644 index 0000000..e3a7509 --- /dev/null +++ b/tests/test_train_smoke.py @@ -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 diff --git a/tests/test_unlabeled_dataset.py b/tests/test_unlabeled_dataset.py new file mode 100644 index 0000000..d09cb03 --- /dev/null +++ b/tests/test_unlabeled_dataset.py @@ -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) diff --git a/tests/test_visualizer.py b/tests/test_visualizer.py new file mode 100644 index 0000000..f27e80b --- /dev/null +++ b/tests/test_visualizer.py @@ -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 diff --git a/tests/test_yolo_format.py b/tests/test_yolo_format.py new file mode 100644 index 0000000..fa9a797 --- /dev/null +++ b/tests/test_yolo_format.py @@ -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"]