Shaun ZhouandClaude Opus 5
65bc6b2045
背景差分：大半径方框均值估背景 → 残差（可选极性）→ 平滑 → 有无判定 → 按峰值归一化， 再交给与 gradcam / shap 同一套 cam_to_boxes 出框。只依赖 numpy。 与归因方法的差别：不需要 --checkpoint，分类器还没训练出来就能先跑一版框，也不受分类 精度影响；类别只能来自目录（--target label）；找的是偏离局部背景的区域，而不是模型 据以判断类别的区域，所以对局部目标好、对铺满整图的纹理类给不出有意义的框。 有无缺陷是显式判定的：残差峰值要同时高于绝对对比度下限与噪声水平（MAD 估计）的 15 倍， 否则返回全零图。这条是必需的——实测纯噪声图的"峰值/MAD"稳定在 9 附近（最大 10.9， 几乎与噪声大小无关），有 0.25 对比度黑点时是 25~330；若只按峰值归一化，正常图的噪声 会被拉到 1，满图假框。阈值 15 取在两者之间。 顺带把积分图方框均值从 shap_explain.smooth 提到 engine/classical.box_mean 共用， 避免两处实现各自漂移。 实测：单张 0.011 秒（全量 11k 约 2 分钟）；真实薄膜图上框能套住黑点、虫状物与划痕。 已知局限：背景估计在图像边界处对明暗阶跃有响应，贴边的交界可能出假框。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
379 lines
14 KiB
Python
379 lines
14 KiB
Python
"""生成框标签的统一入口（gradcam / shap / cv 三种方法）。"""
|
|
|
|
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", "cv"])
|
|
def test_every_method_produces_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_defaults_to_folder_label(project, tmp_path, monkeypatch):
|
|
"""框的类别默认用人工标注，不让分类模型的错误预测混进标签。"""
|
|
out = tmp_path / "out_default"
|
|
assert run(base_argv(project, out, "gradcam"), monkeypatch) == 0
|
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
|
assert summary["target"] == "label"
|
|
|
|
|
|
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
|
|
|
|
|
|
# ---------------- 传统 CV 方法 ----------------
|
|
|
|
|
|
def test_cv_runs_without_a_classifier(project, tmp_path, monkeypatch):
|
|
"""cv 不判类别，也就不需要权重——分类器还没训练出来时就能先跑一版框。"""
|
|
root, _, config = project
|
|
out = tmp_path / "out_cv"
|
|
argv = ["--out", str(out), "--method", "cv", "--config", str(config),
|
|
"--checkpoint", str(tmp_path / "does_not_exist.pth")]
|
|
assert run(argv, monkeypatch) == 0
|
|
|
|
assert len(list((out / "labels").rglob("*.txt"))) == 5
|
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
|
assert summary["method"] == "cv"
|
|
# 没加载模型，就没有预测可比
|
|
assert summary["prediction_matches_folder_label"] is None
|
|
assert summary["cv_params"]["polarity"] == "both"
|
|
|
|
|
|
def test_cv_labels_use_the_folder_class(project, tmp_path, monkeypatch):
|
|
"""类别只能来自目录：cv 自己不判类。"""
|
|
root, _, config = project
|
|
out = tmp_path / "out_cv"
|
|
assert run(["--out", str(out), "--method", "cv", "--config", str(config)], monkeypatch) == 0
|
|
|
|
classes = sorted(p.name for p in root.iterdir() if p.is_dir())
|
|
for label in (out / "labels").rglob("*.txt"):
|
|
folder = label.parent.name
|
|
for line in label.read_text(encoding="utf-8").strip().splitlines():
|
|
assert classes[int(line.split()[0])] == folder
|
|
|
|
|
|
def test_cv_finds_a_planted_defect(tmp_path, monkeypatch):
|
|
"""端到端：图上放一个黑点，框要套住它；平坦的图不出框。"""
|
|
root = tmp_path / "dataset"
|
|
(root / "blackdot").mkdir(parents=True)
|
|
(root / "normal").mkdir(parents=True)
|
|
|
|
rng = np.random.default_rng(0)
|
|
flat = np.clip(rng.normal(140, 2, (96, 96)), 0, 255).astype(np.uint8)
|
|
spotted = flat.copy()
|
|
spotted[28:36, 60:68] = 40 # 一个明显的黑点
|
|
Image.fromarray(spotted).save(root / "blackdot" / "spot.png")
|
|
Image.fromarray(flat).save(root / "normal" / "flat.png")
|
|
|
|
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}}}}
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
out = tmp_path / "out"
|
|
assert run(["--out", str(out), "--method", "cv", "--config", str(config)], monkeypatch) == 0
|
|
|
|
spot = (out / "labels" / "blackdot" / "spot.txt").read_text(encoding="utf-8").split()
|
|
assert len(spot) == 5
|
|
_, cx, cy, width, height = (float(v) for v in spot)
|
|
assert abs(cx - 64 / 96) < 0.1 and abs(cy - 32 / 96) < 0.1 # 命中黑点
|
|
assert width * height < 0.15
|
|
|
|
# 平坦的图没有缺陷，应当是空标签
|
|
assert (out / "labels" / "normal" / "flat.txt").read_text(encoding="utf-8") == ""
|
|
|
|
|
|
def test_cv_polarity_is_passed_through(project, tmp_path, monkeypatch):
|
|
root, _, config = project
|
|
out = tmp_path / "out_cv"
|
|
argv = ["--out", str(out), "--method", "cv", "--config", str(config),
|
|
"--cv-polarity", "dark", "--cv-radius", "5", "--cv-min-contrast", "0.02"]
|
|
assert run(argv, monkeypatch) == 0
|
|
params = json.loads((out / "summary.json").read_text(encoding="utf-8"))["cv_params"]
|
|
assert params == {
|
|
"background_radius": 5, "polarity": "dark", "smooth_radius": 2, "min_contrast": 0.02,
|
|
}
|