Shaun ZhouandClaude Opus 5
86c86b105c
实际工作流是：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>
287 lines
11 KiB
Python
287 lines
11 KiB
Python
"""生成框标签的统一入口（gradcam / shap 两种方法）。"""
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
pytest.importorskip("PIL")
|
|
|
|
import numpy as np # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
|
|
sys.path.insert(0, "scripts")
|
|
import generate_boxes as gb # noqa: E402
|
|
|
|
|
|
@pytest.fixture
|
|
def project(tmp_path):
|
|
"""一个最小的分类数据集 + 训练好的权重 + 配置。"""
|
|
from engine.checkpoint import save_checkpoint
|
|
from models import build_model
|
|
|
|
root = tmp_path / "dataset"
|
|
rng = np.random.default_rng(0)
|
|
for name, count in (("grease", 3), ("normal", 2)):
|
|
directory = root / name
|
|
directory.mkdir(parents=True)
|
|
for i in range(count):
|
|
Image.fromarray(rng.integers(0, 255, (40, 56, 3), dtype=np.uint8)).save(
|
|
directory / f"{name}_{i}.png" # 全局唯一，与 --prefix-parent 的产物一致
|
|
)
|
|
|
|
model = build_model({
|
|
"name": "classifier", "num_classes": 2,
|
|
"backbone": {"name": "torchvision", "arch": "resnet18", "pretrained": False},
|
|
})
|
|
checkpoint = tmp_path / "best.pth"
|
|
save_checkpoint(checkpoint, model)
|
|
|
|
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: torchvision, arch: resnet18, pretrained: false}}
|
|
init_from: null
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
return root, checkpoint, config
|
|
|
|
|
|
def run(argv, monkeypatch):
|
|
monkeypatch.setattr(sys, "argv", ["generate_boxes.py", *argv])
|
|
return gb.main()
|
|
|
|
|
|
def base_argv(project, out, method):
|
|
root, checkpoint, config = project
|
|
return ["--out", str(out), "--method", method,
|
|
"--checkpoint", str(checkpoint), "--config", str(config)]
|
|
|
|
|
|
@pytest.mark.parametrize("method", ["gradcam", "shap"])
|
|
def test_both_methods_produce_a_yolo_dataset(project, tmp_path, monkeypatch, method):
|
|
if method == "shap":
|
|
pytest.importorskip("shap")
|
|
root, _, _ = project
|
|
out = tmp_path / f"out_{method}"
|
|
|
|
extra = ["--max-evals", "60"] if method == "shap" else []
|
|
assert run(base_argv(project, out, method) + extra, monkeypatch) == 0
|
|
|
|
# images/ 与 labels/ 一一对应，符合仓库的 YOLO 约定
|
|
images = sorted(p.relative_to(out / "images") for p in (out / "images").rglob("*.png"))
|
|
labels = sorted(p.relative_to(out / "labels") for p in (out / "labels").rglob("*.txt"))
|
|
assert len(images) == 5
|
|
assert [p.with_suffix("") for p in images] == [p.with_suffix("") for p in labels]
|
|
assert (out / "data.yaml").exists()
|
|
|
|
for label in (out / "labels").rglob("*.txt"):
|
|
for line in label.read_text(encoding="utf-8").strip().splitlines():
|
|
parts = line.split()
|
|
assert len(parts) == 5
|
|
assert all(0.0 <= float(v) <= 1.0 for v in parts[1:])
|
|
|
|
|
|
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)
|
|
|
|
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):
|
|
out = tmp_path / "out"
|
|
run(base_argv(project, out, "gradcam"), monkeypatch)
|
|
|
|
for path in (out / "images").rglob("*.png"):
|
|
assert not path.is_symlink()
|
|
assert path.stat().st_nlink > 1, "应与原图共享 inode，不额外占盘"
|
|
|
|
summary = json.loads((out / "summary.json").read_text("utf-8"))
|
|
assert summary["link_modes"].get("hard") == 5
|
|
|
|
|
|
def test_max_boxes_defaults_to_one(project, tmp_path, monkeypatch):
|
|
out = tmp_path / "out"
|
|
run(base_argv(project, out, "gradcam"), monkeypatch)
|
|
|
|
for label in (out / "labels").rglob("*.txt"):
|
|
assert len(label.read_text(encoding="utf-8").strip().splitlines()) <= 1
|
|
|
|
summary = json.loads((out / "summary.json").read_text("utf-8"))
|
|
assert summary["box_params"]["max_boxes"] == 1
|
|
|
|
|
|
def test_min_area_ratio_default_follows_the_method(project, tmp_path, monkeypatch):
|
|
"""两种方法的归因图分辨率差 16 倍，面积阈值不能共用默认值。"""
|
|
assert gb.DEFAULT_MIN_AREA_RATIO["gradcam"] > gb.DEFAULT_MIN_AREA_RATIO["shap"]
|
|
|
|
out = tmp_path / "out"
|
|
run(base_argv(project, out, "gradcam"), monkeypatch)
|
|
summary = json.loads((out / "summary.json").read_text("utf-8"))
|
|
assert summary["box_params"]["min_area_ratio"] == gb.DEFAULT_MIN_AREA_RATIO["gradcam"]
|
|
|
|
|
|
def test_full_dataset_by_default_and_subsampling_is_opt_in(project, tmp_path, monkeypatch):
|
|
full = tmp_path / "full"
|
|
run(base_argv(project, full, "gradcam"), monkeypatch)
|
|
assert json.loads((full / "summary.json").read_text("utf-8"))["num_images"] == 5
|
|
|
|
limited = tmp_path / "limited"
|
|
run(base_argv(project, limited, "gradcam") + ["--limit", "2"], monkeypatch)
|
|
assert json.loads((limited / "summary.json").read_text("utf-8"))["num_images"] == 2
|
|
|
|
per_class = tmp_path / "per_class"
|
|
run(base_argv(project, per_class, "gradcam") + ["--per-class", "1"], monkeypatch)
|
|
assert json.loads((per_class / "summary.json").read_text("utf-8"))["num_images"] == 2
|
|
|
|
|
|
def test_target_can_be_prediction_or_folder_label(project, tmp_path, monkeypatch):
|
|
for target in ("predicted", "label"):
|
|
out = tmp_path / f"out_{target}"
|
|
assert run(base_argv(project, out, "gradcam") + ["--target", target], monkeypatch) == 0
|
|
summary = json.loads((out / "summary.json").read_text("utf-8"))
|
|
assert summary["target"] == target
|
|
assert 0.0 <= summary["prediction_matches_folder_label"] <= 1.0
|
|
|
|
|
|
def test_optional_outputs_can_be_skipped(project, tmp_path, monkeypatch):
|
|
out = tmp_path / "out"
|
|
run(base_argv(project, out, "gradcam") + ["--no-heatmaps", "--no-images"], monkeypatch)
|
|
|
|
assert not (out / "heatmaps").exists()
|
|
assert not (out / "images").exists()
|
|
assert list((out / "labels").rglob("*.txt")) # 标签仍然产出
|
|
|
|
|
|
def test_heatmaps_are_written_per_image(project, tmp_path, monkeypatch):
|
|
out = tmp_path / "out"
|
|
run(base_argv(project, out, "gradcam"), monkeypatch)
|
|
|
|
heatmaps = list((out / "heatmaps").rglob("*.png"))
|
|
assert len(heatmaps) == 5
|
|
with Image.open(heatmaps[0]) as image:
|
|
assert image.size in {(56, 40)} # 与原图同尺寸
|
|
|
|
|
|
def test_link_or_copy_falls_back_across_filesystems(tmp_path, monkeypatch):
|
|
import os as os_mod
|
|
|
|
source = tmp_path / "a.png"
|
|
source.write_bytes(b"x")
|
|
|
|
def refuse(*_args, **_kwargs):
|
|
raise OSError(errno_module.EXDEV, "Cross-device link")
|
|
|
|
import errno as errno_module
|
|
|
|
monkeypatch.setattr(os_mod, "link", refuse)
|
|
assert gb.link_or_copy(source, tmp_path / "b" / "a.png") == "copy"
|
|
assert (tmp_path / "b" / "a.png").read_bytes() == b"x"
|
|
|
|
|
|
def test_duplicate_source_names_are_rejected(tmp_path):
|
|
"""产出按**预测类别**分目录，而源图来自不同的类别目录：
|
|
dirA/x.png 与 dirB/x.png 若被预测到同一类，就会写同一个标签路径。
|
|
必须报错，而不是让后写的悄悄覆盖前一个。"""
|
|
from argparse import Namespace
|
|
|
|
heat = np.zeros((8, 8), dtype=np.float32)
|
|
heat[3:5, 3:5] = 1.0
|
|
box_kwargs = {"threshold": 0.4, "min_area_ratio": 0.0, "max_boxes": 1,
|
|
"padding": 0.0, "peak_ratio": 0.75}
|
|
args = Namespace(no_images=True, no_heatmaps=True)
|
|
seen: dict = {}
|
|
common = dict(out_dir=tmp_path, classes=["grease"], target=0, heat=heat,
|
|
box_kwargs=box_kwargs, args=args, box_counts={}, link_modes={}, seen=seen)
|
|
|
|
from collections import Counter
|
|
common["box_counts"] = Counter()
|
|
common["link_modes"] = Counter()
|
|
|
|
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
|