refactor: 合并为单一的框生成入口，gradcam / shap 可选，默认全量
此前是两个脚本：gradcam_boxes.py 全量跑但只有一种方法，shap_boxes.py 默认抽样 且强制输出对比图。两者合并为 scripts/generate_boxes.py： - --method gradcam|shap 二选一，默认 gradcam - 默认对全量数据生成；--limit / --per-class 降级为试参数用的可选项 - --min-area-ratio 缺省随方法而变（gradcam 0.02、shap 0.002）：两者归因图分辨率差 16 倍，同一个小黑点在 128x128 上只占百来像素，共用阈值会被整个滤掉 产出改为符合 CLAUDE.md 工作流第 8 步的规格，是一个自成一体的 YOLO 数据集： images/ 是原图硬链接、labels/ 是 YOLO 标签、heatmaps/ 是热力图叠加框。 images/ 与 labels/ 的对应关系符合仓库的 YOLO 约定，FilmDefectDataset 可直接读取， 把分类产出接回了检测训练链路（已加测试验证这条闭环）。 热力图目标类别改为模型 top1 预测（原为目录标注），--target label 可切回； summary.json 记录两者一致率（实测 0.9873）。 修掉一个会静默丢标签的 bug：产出按预测类别分目录，不同源目录下的同名图 （dirA/x.png 与 dirB/x.png）若被预测到同一类会写同一个标签路径，后写的覆盖前者。 现在显式检测并中止，提示用 export_clusters.py --prefix-parent 重新导出。 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
9aa87d31a6
commit
57446c2f10
@@ -0,0 +1,226 @@
|
||||
"""生成框标签的统一入口（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_readable_by_the_detection_dataset(project, tmp_path, monkeypatch):
|
||||
"""产出要能直接喂给检测训练——这是这一步存在的意义。"""
|
||||
from datasets import FilmDefectDataset
|
||||
|
||||
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
|
||||
|
||||
|
||||
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) # 同名，撞同一个标签路径
|
||||
Reference in New Issue
Block a user