refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
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
57446c2f10
commit
86c86b105c
+33
-37
@@ -66,10 +66,13 @@ def test_load_config_detects_cycle(tmp_path):
|
||||
|
||||
|
||||
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
|
||||
cluster = load_config("configs/cluster/film.yaml")
|
||||
assert cluster["dataset"]["name"] == "unlabeled"
|
||||
assert cluster["num_clusters"] > 0
|
||||
|
||||
classify = load_config("configs/classify/film.yaml")
|
||||
assert classify["model"]["name"] == "classifier"
|
||||
assert classify["train"]["monitor"] == "macro_f1"
|
||||
|
||||
|
||||
def test_replace_key_swaps_a_whole_subdict(tmp_path):
|
||||
@@ -116,46 +119,39 @@ def test_registry_reports_incompatible_fields_actionably():
|
||||
def test_every_shipped_config_builds_its_components(config_path):
|
||||
"""仓库里每个配置都必须能真的构建出组件。
|
||||
|
||||
回归测试：configs/train/resnet18.yaml 曾经无法使用——它同时继承 tiny_conv 与
|
||||
resnet18 的模型配置，合并后 backbone 混了两边的字段，构建时直接报错。
|
||||
回归测试：派生配置曾经无法使用——它同时继承两份骨架配置，
|
||||
合并后 backbone 混了两边的字段（tiny_conv 的 widths 混进 torchvision），构建时直接报错。
|
||||
"""
|
||||
torch = pytest.importorskip("torch")
|
||||
from datasets.yolo import load_classes
|
||||
from models import build_model
|
||||
pytest.importorskip("torch")
|
||||
from models.registry import BACKBONES
|
||||
|
||||
cfg = load_config(config_path)
|
||||
|
||||
if "model" in cfg:
|
||||
if "classes" not in cfg:
|
||||
# configs/model/*.yaml 是片段，只能被 _base_ 组合，单独用没有 classes
|
||||
pytest.skip(f"{config_path} 是配置片段，不含 classes")
|
||||
model_cfg = dict(cfg["model"])
|
||||
model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
|
||||
# 单测里不下载 ImageNet 权重
|
||||
# 聚类配置：pretrained 用顶层 backbone，simclr 用 simclr.backbone
|
||||
for backbone_cfg in (cfg.get("backbone"), cfg.get_path("simclr.backbone")):
|
||||
if not backbone_cfg:
|
||||
continue
|
||||
backbone_cfg = dict(backbone_cfg)
|
||||
if backbone_cfg.get("name") == "torchvision":
|
||||
backbone_cfg["pretrained"] = False # 单测里不下载 ImageNet 权重
|
||||
backbone = BACKBONES.build(backbone_cfg)
|
||||
assert backbone.out_channels > 0
|
||||
|
||||
# 分类配置
|
||||
model_cfg = cfg.get("model")
|
||||
if model_cfg:
|
||||
import torch
|
||||
|
||||
from models import build_model
|
||||
|
||||
model_cfg = dict(model_cfg)
|
||||
model_cfg.setdefault("num_classes", 3)
|
||||
model_cfg["init_from"] = None
|
||||
if model_cfg.get("backbone", {}).get("name") == "torchvision":
|
||||
model_cfg["backbone"] = {**model_cfg["backbone"], "pretrained": False}
|
||||
model = build_model(model_cfg)
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(1, 3, 64, 64))
|
||||
assert "heatmap" in out
|
||||
|
||||
simclr_backbone = cfg.get_path("simclr.backbone")
|
||||
if simclr_backbone:
|
||||
backbone_cfg = dict(simclr_backbone)
|
||||
if backbone_cfg.get("name") == "torchvision":
|
||||
backbone_cfg["pretrained"] = False
|
||||
backbone = BACKBONES.build(backbone_cfg)
|
||||
assert backbone.out_channels > 0
|
||||
|
||||
|
||||
def test_eval_config_matches_its_training_backbone():
|
||||
"""评估配置构建出的骨架必须与训练配置一致，否则权重根本加载不上。"""
|
||||
for train_cfg, eval_cfg in [
|
||||
("configs/train/baseline.yaml", "configs/eval/baseline.yaml"),
|
||||
("configs/train/resnet18.yaml", "configs/eval/resnet18.yaml"),
|
||||
]:
|
||||
assert load_config(train_cfg)["model"] == load_config(eval_cfg)["model"], eval_cfg
|
||||
assert model(torch.randn(1, 3, 64, 64)).shape[1] == model_cfg["num_classes"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -169,13 +165,13 @@ def test_eval_config_matches_its_training_backbone():
|
||||
def test_set_rejects_typos_with_suggestion(override, expected_hint):
|
||||
"""--set 打错字母会凭空造出没人读的键，训练照跑但改动无效——必须拦下并给出建议。"""
|
||||
with pytest.raises(KeyError) as excinfo:
|
||||
load_config("configs/cluster/film_simclr_resnet50.yaml", [override])
|
||||
load_config("configs/cluster/film_simclr_resnet18.yaml", [override])
|
||||
assert expected_hint in str(excinfo.value)
|
||||
|
||||
|
||||
def test_set_accepts_existing_keys():
|
||||
cfg = load_config(
|
||||
"configs/cluster/film_simclr_resnet50.yaml",
|
||||
"configs/cluster/film_simclr_resnet18.yaml",
|
||||
["simclr.optimizer.lr=0.001", "simclr.optimizer.weight_decay=0.0", "num_clusters=7"],
|
||||
)
|
||||
assert cfg["simclr"]["optimizer"]["lr"] == 0.001
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
"""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
|
||||
@@ -94,21 +94,25 @@ def test_both_methods_produce_a_yolo_dataset(project, tmp_path, monkeypatch, met
|
||||
assert all(0.0 <= float(v) <= 1.0 for v in parts[1:])
|
||||
|
||||
|
||||
def test_output_is_readable_by_the_detection_dataset(project, tmp_path, monkeypatch):
|
||||
"""产出要能直接喂给检测训练——这是这一步存在的意义。"""
|
||||
from datasets import FilmDefectDataset
|
||||
def test_output_is_a_wellformed_yolo_dataset(project, tmp_path, monkeypatch):
|
||||
"""产出是给外部检测训练用的交换格式，必须自洽：
|
||||
data.yaml 的类别与 labels/ 对得上，每行标签可解析且坐标在 [0,1] 内。"""
|
||||
import yaml
|
||||
|
||||
out = tmp_path / "out"
|
||||
run(base_argv(project, out, "gradcam"), monkeypatch)
|
||||
|
||||
dataset = FilmDefectDataset(
|
||||
img_dir=out / "images", classes=str(out / "data.yaml"),
|
||||
transforms=[{"name": "Resize", "size": [64, 64]}, {"name": "ToTensor"}], stride=4,
|
||||
)
|
||||
assert len(dataset) == 5
|
||||
sample = dataset[0]
|
||||
assert sample["image"].shape == (3, 64, 64)
|
||||
assert sample["boxes"].shape[1] == 4
|
||||
meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
|
||||
names = [meta["names"][i] for i in sorted(meta["names"])]
|
||||
assert meta["nc"] == len(names) == 2
|
||||
|
||||
for label in (out / "labels").rglob("*.txt"):
|
||||
assert (out / "images" / label.parent.name / f"{label.stem}.png").is_file()
|
||||
for line in label.read_text(encoding="utf-8").strip().splitlines():
|
||||
class_id, cx, cy, w, h = line.split()
|
||||
assert 0 <= int(class_id) < meta["nc"]
|
||||
assert all(0.0 <= float(v) <= 1.0 for v in (cx, cy, w, h))
|
||||
assert float(w) > 0 and float(h) > 0
|
||||
|
||||
|
||||
def test_images_are_hard_links_not_copies(project, tmp_path, monkeypatch):
|
||||
@@ -224,3 +228,59 @@ def test_duplicate_source_names_are_rejected(tmp_path):
|
||||
gb._emit(source=Path("/data/dirA/x.png"), **common)
|
||||
with pytest.raises(gb.NameCollision, match="prefix-parent"):
|
||||
gb._emit(source=Path("/data/dirB/x.png"), **common) # 同名，撞同一个标签路径
|
||||
|
||||
|
||||
def test_gradcam_falls_back_when_backbone_has_no_sequential_features(tmp_path, monkeypatch):
|
||||
"""tiny_conv 不是 nn.Sequential，取不了中间层——应当退回最终输出并警告，而不是崩掉。"""
|
||||
from engine.checkpoint import save_checkpoint
|
||||
from models import build_model
|
||||
|
||||
root = tmp_path / "dataset"
|
||||
(root / "a").mkdir(parents=True)
|
||||
rng = np.random.default_rng(0)
|
||||
for i in range(2):
|
||||
Image.fromarray(rng.integers(0, 255, (40, 56, 3), dtype=np.uint8)).save(
|
||||
root / "a" / f"a_{i}.png"
|
||||
)
|
||||
|
||||
backbone = {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}
|
||||
checkpoint = tmp_path / "best.pth"
|
||||
save_checkpoint(checkpoint, build_model(
|
||||
{"name": "classifier", "num_classes": 1, "backbone": dict(backbone)}
|
||||
))
|
||||
|
||||
config = tmp_path / "cfg.yaml"
|
||||
config.write_text(
|
||||
f"""
|
||||
device: cpu
|
||||
data:
|
||||
root: {root}
|
||||
classes: null
|
||||
exclude: []
|
||||
val_transforms: [{{name: Resize, size: [64, 64]}}, {{name: ToTensor}}]
|
||||
loader: {{batch_size: 2, num_workers: 0}}
|
||||
model:
|
||||
name: classifier
|
||||
backbone: {{name: tiny_conv, widths: [8, 16], depths: [1, 1]}}
|
||||
init_from: null
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
out = tmp_path / "out"
|
||||
monkeypatch.setattr(
|
||||
sys, "argv",
|
||||
["generate_boxes.py", "--out", str(out), "--method", "gradcam",
|
||||
"--checkpoint", str(checkpoint), "--config", str(config), "--no-heatmaps"],
|
||||
)
|
||||
assert gb.main() == 0 # 不崩
|
||||
|
||||
summary = json.loads((out / "summary.json").read_text("utf-8"))
|
||||
assert summary["stage"] is None # 已退回最终输出
|
||||
assert len(list((out / "labels").rglob("*.txt"))) == 2
|
||||
|
||||
|
||||
def test_stage_minus_one_selects_the_final_output(project, tmp_path, monkeypatch):
|
||||
out = tmp_path / "out"
|
||||
assert run(base_argv(project, out, "gradcam") + ["--stage", "-1"], monkeypatch) == 0
|
||||
assert json.loads((out / "summary.json").read_text("utf-8"))["stage"] is None
|
||||
|
||||
@@ -1,125 +0,0 @@
|
||||
"""多类(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)
|
||||
+26
-102
@@ -1,96 +1,12 @@
|
||||
"""卷积骨架与分类模型。"""
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
from models import build_model # noqa: E402
|
||||
from models.backbones.resnet import _SUPPORTED as _SUPPORTED_ARCHS # 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
|
||||
from models.heads.projection import ProjectionHead # noqa: E402
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arch", sorted(_SUPPORTED_ARCHS))
|
||||
@@ -118,10 +34,23 @@ def test_unsupported_arch_is_rejected():
|
||||
BACKBONES.build({"name": "torchvision", "arch": "vit_b_16"})
|
||||
|
||||
|
||||
def test_tiny_conv_shapes():
|
||||
from models.registry import BACKBONES
|
||||
|
||||
backbone = BACKBONES.build(
|
||||
{"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]}
|
||||
).eval()
|
||||
with torch.no_grad():
|
||||
feat = backbone(torch.randn(2, 3, 64, 64))
|
||||
assert feat.shape[:2] == (2, 32)
|
||||
assert backbone.out_channels == 32
|
||||
# stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍
|
||||
assert backbone.out_stride == 8
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"])
|
||||
def test_alternative_backbones_drive_simclr(arch):
|
||||
"""换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。"""
|
||||
from models.heads.projection import ProjectionHead
|
||||
from models.registry import BACKBONES
|
||||
|
||||
backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval()
|
||||
@@ -130,18 +59,13 @@ def test_alternative_backbones_drive_simclr(arch):
|
||||
assert head(backbone(torch.randn(2, 3, 64, 64))).shape == (2, 16)
|
||||
|
||||
|
||||
def test_alternative_backbone_drives_defect_net():
|
||||
"""检测模型换骨架同样要能跑通（neck 的 in_channels 由 backbone 决定）。"""
|
||||
model = build_model(
|
||||
{
|
||||
"name": "defect_net",
|
||||
"num_classes": 3,
|
||||
"backbone": {"name": "torchvision", "arch": "shufflenet_v2_x1_0", "pretrained": False},
|
||||
"neck": {"name": "upsample", "channels": 16, "num_upsamples": 3},
|
||||
"head": {"hidden_channels": 16},
|
||||
}
|
||||
).eval()
|
||||
assert model.out_stride == 4
|
||||
def test_classifier_shapes_and_convolution_only():
|
||||
model = build_model({
|
||||
"name": "classifier", "num_classes": 4,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
||||
"classification": {"dropout": 0.1},
|
||||
}).eval()
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(1, 3, 64, 64))
|
||||
assert out["heatmap"].shape == (1, 3, 16, 16)
|
||||
assert model(torch.randn(3, 3, 64, 64)).shape == (3, 4)
|
||||
forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM)
|
||||
assert not any(isinstance(m, forbidden) for m in model.modules())
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
"""推理 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)
|
||||
@@ -1,42 +0,0 @@
|
||||
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
|
||||
@@ -1,217 +0,0 @@
|
||||
"""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 torch # 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]
|
||||
|
||||
|
||||
def test_resume_takes_hyperparams_from_config(project, monkeypatch, tmp_path):
|
||||
"""检测训练续训：lr / weight_decay / 总轮数以 yaml 为准，检查点只提供进度。
|
||||
|
||||
回归测试：load_checkpoint 直接 optimizer.load_state_dict，把超参连同动量一起还原，
|
||||
改了配置再续训会被静默覆盖；而且此前根本没保存调度器状态，学习率调度会从头开始。
|
||||
"""
|
||||
import math
|
||||
|
||||
import yaml as yaml_mod
|
||||
|
||||
config_path, root = project
|
||||
cfg = yaml_mod.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
cfg["train"].update({"epochs": 4, "optimizer": {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
|
||||
"scheduler": {"name": "cosine"}})
|
||||
config_path.write_text(yaml_mod.safe_dump(cfg, allow_unicode=True), encoding="utf-8")
|
||||
|
||||
work_dir = tmp_path / "runs"
|
||||
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
||||
monkeypatch) == 0
|
||||
|
||||
captured = {}
|
||||
real_fit = train_mod.Trainer.fit
|
||||
|
||||
def probe(self):
|
||||
captured["groups"] = [dict(g) for g in self.optimizer.param_groups]
|
||||
captured["scheduler"] = self.scheduler
|
||||
captured["start_epoch"] = self.start_epoch
|
||||
return 0.0
|
||||
|
||||
monkeypatch.setattr(train_mod.Trainer, "fit", probe)
|
||||
assert run(
|
||||
train_mod,
|
||||
["train.py", "--config", str(config_path), "--work-dir", str(work_dir),
|
||||
"--resume", str(work_dir / "last.pth"), "--set",
|
||||
"train.epochs=10", "train.optimizer.lr=5e-4", "train.optimizer.weight_decay=0.05"],
|
||||
monkeypatch,
|
||||
) == 0
|
||||
monkeypatch.setattr(train_mod.Trainer, "fit", real_fit)
|
||||
|
||||
groups = captured["groups"]
|
||||
assert groups[0]["initial_lr"] == pytest.approx(5e-4) # 不是检查点里的 1e-3
|
||||
assert groups[0]["weight_decay"] == pytest.approx(0.05) # 不是检查点里的 1e-4
|
||||
assert groups[1]["weight_decay"] == 0.0 # BN/bias 分组结构没被冲掉
|
||||
|
||||
scheduler = captured["scheduler"]
|
||||
assert scheduler.T_max == 10
|
||||
assert scheduler.last_epoch == captured["start_epoch"] == 4 # 调度进度已快进，不是从 0 开始
|
||||
assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 10)) / 2)
|
||||
|
||||
|
||||
def test_checkpoint_carries_scheduler_and_progress(project, monkeypatch, tmp_path):
|
||||
"""检查点要带上调度器状态与 global_step，否则续训无从恢复调度进度。"""
|
||||
config_path, _ = project
|
||||
work_dir = tmp_path / "runs"
|
||||
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
||||
monkeypatch) == 0
|
||||
|
||||
state = torch.load(work_dir / "last.pth", map_location="cpu", weights_only=False)
|
||||
expected = {"model", "optimizer", "scheduler", "epoch", "global_step", "best_metric", "train_cfg"}
|
||||
assert expected <= set(state)
|
||||
assert state["global_step"] > 0
|
||||
assert state["train_cfg"]["optimizer"]["lr"] is not None
|
||||
@@ -1,118 +0,0 @@
|
||||
"""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