ModelTrainEval/tests/test_train_smoke.py

"""train.py / eval.py 的端到端冒烟：YOLO 数据目录 -> 训练 -> 评估。

覆盖入口脚本的接线：classes 解析、label_mode 下发、设备解析。固定用 CPU 跑，
免得测试受机器上有没有 GPU 影响。
"""

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

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"

    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)

    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_unavailable_device_fails_loudly(project, monkeypatch, tmp_path):
    """显式请求机器上没有的后端时必须退出，而不是悄悄退回 CPU。"""
    from utils.device import available_backends

    missing = next((b for b in ("cuda", "mps") if b not in available_backends()), None)
    if missing is None:
        pytest.skip("这台机器 cuda 与 mps 都可用")

    assert run(
        train_mod,
        ["train.py", "--config", str(config_path_of(project)), "--work-dir", str(tmp_path / "runs3"),
         "--device", missing],
        monkeypatch,
    ) == 1


def test_device_flag_overrides_config(project, monkeypatch, tmp_path):
    """--device 优先于配置里的 device。"""
    config_path, _ = project
    work_dir = tmp_path / "runs4"
    assert run(
        train_mod,
        ["train.py", "--config", str(config_path), "--work-dir", str(work_dir), "--device", "cpu"],
        monkeypatch,
    ) == 0
    assert (work_dir / "best.pth").exists()


def config_path_of(project):
    return project[0]