Shaun ZhouandClaude Opus 5
32e56cb230
框的类别应该来自人工归类（图片所在目录），分类模型只负责定位。原先默认 predicted， 框的类别取模型 top1 预测，模型分错的图会把错误类别写进标签：boxes_shap 那次预测与 目录一致率 0.9915，约 95 张图的类别因此是错的。 两种取值只在预测与目录不一致的图上有区别，其余图解释的是同一个 logit、框完全相同。 --target predicted 仍保留。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
295 lines
11 KiB
Python
295 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_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
|