From 57446c2f1064a3e8fa140158b8bc253b9d211f5b Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Thu, 10 Sep 2026 16:26:56 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E5=90=88=E5=B9=B6=E4=B8=BA?= =?UTF-8?q?=E5=8D=95=E4=B8=80=E7=9A=84=E6=A1=86=E7=94=9F=E6=88=90=E5=85=A5?= =?UTF-8?q?=E5=8F=A3=EF=BC=8Cgradcam=20/=20shap=20=E5=8F=AF=E9=80=89?= =?UTF-8?q?=EF=BC=8C=E9=BB=98=E8=AE=A4=E5=85=A8=E9=87=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前是两个脚本：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 --- README.md | 68 ++++--- scripts/generate_boxes.py | 368 +++++++++++++++++++++++++++++++++++ scripts/gradcam_boxes.py | 251 ------------------------ scripts/shap_boxes.py | 239 ----------------------- tests/test_generate_boxes.py | 226 +++++++++++++++++++++ 5 files changed, 634 insertions(+), 518 deletions(-) create mode 100644 scripts/generate_boxes.py delete mode 100644 scripts/gradcam_boxes.py delete mode 100644 scripts/shap_boxes.py create mode 100644 tests/test_generate_boxes.py diff --git a/README.md b/README.md index 6fe08e1..24d79e5 100644 --- a/README.md +++ b/README.md @@ -632,52 +632,64 @@ model: 只加载 backbone，分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量 自监督特征本身有多好。 -## 从分类到定位：Grad-CAM 生成框标签 +## 从分类到定位：生成框标签 -分类模型只学过"这是什么缺陷"，没学过"缺陷在哪"。Grad-CAM 把分类分数对卷积特征的梯度 -当作通道重要性，加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 + -取连通域外接矩形，就得到弱监督的框标签： +分类模型只学过"这是什么缺陷"，没学过"缺陷在哪"。归因方法能把分类分数反推回空间位置， +阈值化 + 取连通域外接矩形，就得到弱监督的框标签： ```bash -uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth +# Grad-CAM（默认，全量 11k 张约 2 分钟） +uv run python scripts/generate_boxes.py --out runs/boxes_gradcam --method gradcam + +# SHAP Partition（更精细，全量约 1.7 小时） +uv run python scripts/generate_boxes.py --out runs/boxes_shap --method shap ``` +**默认对全量数据生成**；`--limit` / `--per-class` 只用于快速试参数。 + +### 产出是自成一体的 YOLO 数据集 + ``` -runs/gradcam_boxes/ -├── labels/<类别>/<图名>.txt # YOLO 格式 class_id cx cy w h，可直接喂给检测训练 -├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框，供人工核对 +runs// +├── images/<类别>/<图名>.png # 硬链接到原图，不占额外磁盘 +├── labels/<类别>/<图名>.txt # YOLO：class_id cx cy w h（归一化） +├── heatmaps/<类别>/<图名>.png # 热力图叠加原图 + 框，供人工核对 ├── data.yaml └── summary.json ``` -**框的类别取自图片所在目录**（人工归类结果），不是模型预测——即使模型分错， -要的也是该图真实类别的热力区域。 +`images/` 与 `labels/` 的对应关系符合本仓库的 YOLO 约定，所以 `FilmDefectDataset` +能直接读——这一步把分类的产出接回了检测训练链路。 -### 分辨率是主要限制 +热力图默认针对模型的 **top1 预测类别**；`--target label` 改用图片所在目录的标注类别。 +`summary.json` 里的 `prediction_matches_folder_label` 记录两者的一致率（实测 0.987）。 -resnet18 在 128×128 输入下各层的 CAM 尺寸：layer4 只有 4×4、layer3 是 8×8、layer2 是 16×16。 -层越浅越精细，但类别判别性越弱。`--stage`（4~7 对应 layer1~4）可选，默认 6 即 **layer3**： +### 两种方法的取舍 -| stage | CAM | 实测（每类抽 3 张，共 45 张） | +| | `gradcam` | `shap` | | --- | --- | --- | -| 7 (layer4) | 4×4 | 100% 出 1 个框，但只有 16 个格子，框很粗 | -| **6 (layer3)** | **8×8** | **76% 出 1 框、20% 出 2 框，热力区域与缺陷贴合** | -| 5 (layer2) | 16×16 | 40% 出不了框，响应太散 | +| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） | +| 单张耗时 | ~0.01 秒 | ~0.55 秒 | +| 小目标 | 几像素的黑点在 8×8 上不可分辨 | 能贴合到几十像素的目标 | -### 几个调参点 +实测框面积占比（越小越贴合缺陷本体）：`blackdot` 0.019(shap) vs 0.269(gradcam)、 +`mosquito` 0.030 vs 0.173、`whitedot` 0.076 vs 0.253。小目标类别上 SHAP 明显更准， +代价是慢 50 倍。SHAP 需要 `uv sync --group explain`。 +### 调参 + +- `--max-boxes`（默认 1）：一张裁片通常只对应一处缺陷，多出来的框多半是噪声 - `--threshold`（默认 0.4）：热力图相对最大值的阈值 -- `--max-boxes`（默认 1）：每张图最多保留几个框，按连通域内热力总和排序取前 N 个。 - 一张裁片通常只对应一处缺陷，多出来的框多半是噪声；确实有多缺陷的图再调大 -- `--peak-ratio`（默认 0.75）：次要连通域的峰值要达到全局最大值的这个比例才保留。 - 不加这条会在没有缺陷的角落刷出假框 -- `--padding`（默认 0.1）：框四周外扩。CAM 倾向只覆盖最具判别性的部位（缺陷中心）， - 往往比真实缺陷范围小 -- `--viz-per-class`（默认 20）：每类存多少张可视化；`--viz-all` 全存（很慢很占盘） +- `--min-area-ratio`：**缺省随方法而变**（gradcam 0.02、shap 0.002）。两者归因图分辨率 + 差 16 倍，同一个小黑点在 128×128 上只占百来像素，沿用 0.02（=328 像素）会被整个滤掉 +- `--peak-ratio`（默认 0.75）：次要连通域的峰值需达全局最大值的这个比例，滤掉弱响应假框 +- `--padding`（默认 0.1）：框四周外扩。归因倾向只覆盖最具判别性的部位，往往偏小 +- `--stage`（gradcam，默认 6=layer3）：layer4 只有 4×4 太粗，layer2 响应太散 +- `--no-heatmaps` / `--no-images`：只要标签时可跳过 -> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 这类描述的是整图属性而非 -> 局部缺陷，它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/ -> `grease`/`zigzag` 这些有局部目标的类别。 +> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 描述的是整图属性而非局部缺陷， +> 它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/`grease` 这些 +> 有局部目标的类别。 ## 方法 diff --git a/scripts/generate_boxes.py b/scripts/generate_boxes.py new file mode 100644 index 0000000..dfc593f --- /dev/null +++ b/scripts/generate_boxes.py @@ -0,0 +1,368 @@ +#!/usr/bin/env python +"""用训练好的分类模型为每张图生成热力图与 YOLO 框标签。 + + # Grad-CAM（快，全量 11k 张约 2 分钟） + uv run python scripts/generate_boxes.py --out runs/boxes_gradcam --method gradcam + + # SHAP Partition（慢但更精细，全量约 1.7 小时） + uv run python scripts/generate_boxes.py --out runs/boxes_shap --method shap + +两种归因方法都能定位，取舍不同： + +| | gradcam | shap | +| --- | --- | --- | +| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） | +| 单张耗时 | 一次 backward，约 0.01 秒 | 几百次前向，约 0.55 秒 | +| 小目标 | 几像素的黑点在 8×8 上不可分辨 | 能贴合到几十像素的目标 | + +产出是一个**自成一体的 YOLO 数据集**，可直接喂给 train.py： + +``` +runs// +├── images/<类别>/<图名>.png # 硬链接到原图，不占额外磁盘 +├── labels/<类别>/<图名>.txt # YOLO：class_id cx cy w h（归一化） +├── heatmaps/<类别>/<图名>.png # 热力图叠加原图 + 框，供人工核对 +├── data.yaml +└── summary.json +``` + +images/ 与 labels/ 的对应关系符合本仓库的 YOLO 约定（datasets/yolo.py）， +所以 `FilmDefectDataset` 能直接读。 + +默认对**全量**数据生成；`--limit` / `--per-class` 只用于快速试参数。 +热力图默认针对模型的 **top1 预测类别**；`--target label` 可改用图片所在目录的标注类别。 +""" + +from __future__ import annotations + +import argparse +import errno +import json +import os +import random +import shutil +import sys +import time +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 +from torch.utils.data import DataLoader # noqa: E402 + +from datasets import FolderClassificationDataset, classification_collate # noqa: E402 +from datasets.yolo import write_data_yaml # noqa: E402 +from engine import load_checkpoint # noqa: E402 +from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line # noqa: E402 +from models import build_model # noqa: E402 +from utils import load_config, setup_logging # noqa: E402 +from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 + +logger = None + +METHODS = ("gradcam", "shap") + +#: 两种方法的归因图分辨率差 16 倍，面积阈值不能共用一个默认值： +#: 同一个小黑点在 8×8 上占几格、在 128×128 上只占百来像素， +#: 沿用 0.02（=328 像素）会把它整个滤掉。 +DEFAULT_MIN_AREA_RATIO = {"gradcam": 0.02, "shap": 0.002} + + +class NameCollision(RuntimeError): + """两张不同的源图在产出目录里撞到了同一个文件名。""" + + +def link_or_copy(source: Path, destination: Path) -> str: + """给原图建硬链接；跨文件系统时退回复制。 + + 硬链接不占额外磁盘，却表现得和真实文件一样——产出目录可以整个拷走。 + """ + if destination.exists(): + return "exists" + destination.parent.mkdir(parents=True, exist_ok=True) + try: + os.link(source, destination) + return "hard" + except OSError as exc: + if exc.errno != errno.EXDEV: + raise + shutil.copy2(source, destination) + return "copy" + + +def save_heatmap(path: Path, image_path: Path, heat: np.ndarray, boxes) -> None: + """热力图叠加原图并画框。用 PIL 直接合成而不是 matplotlib： + 全量 11k 张要出图，逐张建 figure 太慢也太占盘。""" + from matplotlib import colormaps + + with Image.open(image_path) as raw: + base = raw.convert("RGB") + width, height = base.size + + resized = np.asarray( + Image.fromarray((heat * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR) + ) / 255.0 + colored = (colormaps["jet"](resized)[..., :3] * 255).astype(np.uint8) + blended = Image.blend(base, Image.fromarray(colored), alpha=0.45) + + from PIL import ImageDraw + + draw = ImageDraw.Draw(blended) + for x1, y1, x2, y2 in boxes: + draw.rectangle( + [x1 * width, y1 * height, x2 * width, y2 * height], outline=(0, 255, 0), width=2 + ) + path.parent.mkdir(parents=True, exist_ok=True) + blended.save(path) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--out", required=True, help="产出目录，建议放在 runs/ 下") + parser.add_argument( + "--method", choices=METHODS, default="gradcam", help="归因方法，默认 gradcam" + ) + parser.add_argument("--checkpoint", default="runs/classify/best.pth", help="分类模型权重") + parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置") + parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") + parser.add_argument("--device", default=None) + parser.add_argument( + "--target", choices=("predicted", "label"), default="predicted", + help="热力图针对哪个类别：predicted=模型 top1 预测（默认），label=图片所在目录", + ) + + parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值（相对最大值）") + parser.add_argument( + "--max-boxes", type=int, default=1, + help="每张图最多保留几个框，按归因总量排序取前 N 个。" + "默认 1：一张裁片通常只对应一处缺陷", + ) + parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例") + parser.add_argument( + "--min-area-ratio", type=float, default=None, + help=f"连通域面积下限（占整图比例）。缺省随方法而定：{DEFAULT_MIN_AREA_RATIO}", + ) + parser.add_argument( + "--peak-ratio", type=float, default=0.75, + help="次要连通域的峰值需达到全局最大值的这个比例才保留，用于滤掉弱响应假框", + ) + + parser.add_argument( + "--stage", type=int, default=6, + help="gradcam 的目标层（4~7 = layer1~4）。默认 6 即 layer3（8×8）；" + "layer4 只有 4×4 太粗，layer2 响应太散", + ) + parser.add_argument("--max-evals", type=int, default=500, help="shap 每张图的前向次数") + parser.add_argument("--smooth-radius", type=int, default=2, help="shap 归因图平滑半径，0=不平滑") + + parser.add_argument("--no-heatmaps", action="store_true", help="不输出热力图，只出标签") + parser.add_argument("--no-images", action="store_true", help="不建原图硬链接") + parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张（试参数用）") + parser.add_argument( + "--per-class", type=int, default=0, help="每类只处理 N 张（试参数用），0=全量" + ) + parser.add_argument("--seed", type=int, default=0, help="--per-class 抽样的随机种子") + return parser.parse_args() + + +def main() -> int: + global logger + args = parse_args() + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + logger = setup_logging(out_dir / "generate_boxes.log") + + cfg = load_config(args.config, args.set) + try: + device = resolve_device(args.device or cfg.get("device", "auto")) + except DeviceNotAvailableError as exc: + logger.error("%s", exc) + return 1 + + data_cfg = cfg["data"] + dataset = FolderClassificationDataset( + root=data_cfg["root"], split="all", classes=data_cfg.get("classes"), + transforms=data_cfg["val_transforms"], exclude=data_cfg.get("exclude") or [], + ) + classes = dataset.classes + + if args.per_class: + grouped: dict[int, list] = {} + for sample in dataset.samples: + grouped.setdefault(sample[1], []).append(sample) + rng = random.Random(args.seed) + dataset.samples = [ + s for group in grouped.values() for s in rng.sample(group, min(args.per_class, len(group))) + ] + if args.limit: + dataset.samples = dataset.samples[: args.limit] + + model_cfg = dict(cfg["model"]) + model_cfg.setdefault("num_classes", len(classes)) + model_cfg["init_from"] = None # 权重来自 checkpoint + model = build_model(model_cfg) + load_checkpoint(args.checkpoint, model, map_location=device) + model.to(device).eval() + for param in model.parameters(): # 归因需要梯度，但不更新权重 + param.requires_grad_(True) + + min_area_ratio = ( + args.min_area_ratio if args.min_area_ratio is not None + else DEFAULT_MIN_AREA_RATIO[args.method] + ) + box_kwargs = { + "threshold": args.threshold, "min_area_ratio": min_area_ratio, + "max_boxes": args.max_boxes, "padding": args.padding, "peak_ratio": args.peak_ratio, + } + + logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint) + logger.info("方法: %s | 目标类别: %s", args.method, args.target) + logger.info("数据集 %s：%d 张，%d 类", data_cfg["root"], len(dataset), len(classes)) + logger.info("框参数: %s", box_kwargs) + if args.method == "shap": + logger.warning( + "shap 单张约 0.55 秒，%d 张预计 %.1f 分钟；只想试参数可用 --per-class / --limit", + len(dataset), len(dataset) * 0.55 / 60, + ) + + started = time.time() + box_counts: Counter[int] = Counter() + link_modes: Counter[str] = Counter() + seen: dict[Path, Path] = {} + agree = 0 + heat_shape = None + + if args.method == "gradcam": + loader = DataLoader( + dataset, + batch_size=int(data_cfg.get("loader", {}).get("batch_size", 64)), + shuffle=False, + num_workers=int(data_cfg.get("loader", {}).get("num_workers", 4)), + collate_fn=classification_collate, + ) + processed = 0 + for step, batch in enumerate(loader): + images = batch["images"].to(device) + labels = batch["labels"].to(device) + with torch.no_grad(): + predicted = model(images).argmax(1) + targets = predicted if args.target == "predicted" else labels + agree += int((predicted == labels).sum()) + heats = grad_cam(model.backbone, model.head, images, targets, stage=args.stage) + heat_shape = tuple(heats.shape[-2:]) + + for heat, target, meta in zip( + heats.cpu().numpy(), targets.cpu().tolist(), batch["meta"], strict=True + ): + try: + _emit(out_dir, Path(meta["path"]), classes, target, heat, box_kwargs, + args, box_counts, link_modes, seen) + except NameCollision as exc: + logger.error("%s", exc) + return 1 + processed += 1 + if step % 20 == 0: + logger.info("处理 [%d/%d]", processed, len(dataset)) + else: + from engine.shap_explain import shap_attribution, smooth + + for index in range(len(dataset)): + sample = dataset[index] + image_chw = sample["image"] + with torch.no_grad(): + predicted = int(model(image_chw[None].to(device)).argmax(1)) + agree += int(predicted == sample["label"]) + target = predicted if args.target == "predicted" else sample["label"] + + heat = shap_attribution( + model, image_chw.permute(1, 2, 0).numpy().astype(np.float32), target, device, + max_evals=args.max_evals, + ) + heat = smooth(heat, args.smooth_radius) + heat_shape = heat.shape + try: + _emit(out_dir, Path(sample["meta"]["path"]), classes, target, heat, box_kwargs, + args, box_counts, link_modes, seen) + except NameCollision as exc: + logger.error("%s", exc) + return 1 + if index % 50 == 0: + elapsed = time.time() - started + logger.info( + "处理 [%d/%d] 已用 %.0fs，平均 %.2fs/张", index, len(dataset), + elapsed, elapsed / max(index, 1), + ) + + total = sum(box_counts.values()) + write_data_yaml(out_dir / "data.yaml", classes, root=out_dir, + train="images", val="images") + summary = { + "method": args.method, + "checkpoint": args.checkpoint, + "dataset": str(data_cfg["root"]), + "target": args.target, + "num_images": total, + "heatmap_shape": list(heat_shape) if heat_shape else None, + "seconds_per_image": (time.time() - started) / max(total, 1), + "box_params": box_kwargs, + "stage": args.stage if args.method == "gradcam" else None, + "max_evals": args.max_evals if args.method == "shap" else None, + "prediction_matches_folder_label": agree / total if total else 0.0, + "boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())}, + "link_modes": dict(link_modes), + } + (out_dir / "summary.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" + ) + + logger.info( + "完成 %d 张，热力图 %s，平均 %.3fs/张", total, heat_shape, summary["seconds_per_image"] + ) + logger.info("模型预测与目录标注一致率: %.4f", summary["prediction_matches_folder_label"]) + for count, number in sorted(box_counts.items()): + logger.info(" 生成 %d 个框: %d 张 (%.1f%%)", count, number, 100 * number / total) + if link_modes: + logger.info("原图链接方式: %s", dict(link_modes)) + logger.info("产出目录: %s", out_dir) + return 0 + + +def _emit(out_dir, source, classes, target, heat, box_kwargs, args, box_counts, link_modes, + seen) -> None: + """把一张图的结果落盘：硬链接 + YOLO 标签 + 热力图。 + + 产出按**预测类别**分目录，而源图来自不同的类别目录，所以两张同名的源图 + （dirA/0.png 与 dirB/0.png）可能落到同一个目标路径上——后写的会覆盖前一个的标签。 + 这里显式挡住，而不是让标签悄悄丢失。 + """ + name = classes[target] + boxes = cam_to_boxes(heat, **box_kwargs) + box_counts[len(boxes)] += 1 + + label_path = out_dir / "labels" / name / f"{source.stem}.txt" + previous = seen.get(label_path) + if previous is not None: + raise NameCollision( + f"{previous} 与 {source} 在产出目录里都要写 {label_path}，后者会覆盖前者。" + "源图文件名需要全局唯一——用 scripts/export_clusters.py --prefix-parent " + "重新导出数据集即可。" + ) + seen[label_path] = source + label_path.parent.mkdir(parents=True, exist_ok=True) + label_path.write_text( + "\n".join(to_yolo_line(box, target) for box in boxes) + ("\n" if boxes else ""), + encoding="utf-8", + ) + if not args.no_images: + link_modes[link_or_copy(source, out_dir / "images" / name / source.name)] += 1 + if not args.no_heatmaps: + save_heatmap(out_dir / "heatmaps" / name / f"{source.stem}.png", source, heat, boxes) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/gradcam_boxes.py b/scripts/gradcam_boxes.py deleted file mode 100644 index 129110b..0000000 --- a/scripts/gradcam_boxes.py +++ /dev/null @@ -1,251 +0,0 @@ -#!/usr/bin/env python -"""用训练好的分类模型 + Grad-CAM 为每张图生成框标签。 - - uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth - -分类模型只学过"这是什么缺陷"。Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性， -加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 + 取连通域外接矩形， -就得到弱监督的框标签。 - -产出： - -``` -runs/gradcam_boxes/ -├── labels/<类别>/<图名>.txt # YOLO 格式：class_id cx cy w h（归一化） -├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框，供人工核对匹配程度 -├── data.yaml # 类别表，可直接喂给检测训练 -└── summary.json -``` - -**框的类别取自图片所在目录**（人工归类的结果），不是模型预测——即使模型分错， -我们要的也是该图真实类别的热力区域。 - -分辨率是主要限制：resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。 -`--stage` 可选层（4~7 对应 layer1~4），层越浅越精细但类别判别性越弱。 -""" - -from __future__ import annotations - -import argparse -import json -import sys -from collections import Counter -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 -from torch.utils.data import DataLoader # noqa: E402 - -from datasets import FolderClassificationDataset, classification_collate # noqa: E402 -from datasets.yolo import write_data_yaml # noqa: E402 -from engine import load_checkpoint # noqa: E402 -from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line # noqa: E402 -from models import build_model # noqa: E402 -from utils import load_config, setup_logging # noqa: E402 -from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 - -logger = None - - -def save_visualization(path: Path, image_path: Path, cam: np.ndarray, boxes, title: str) -> None: - """三联图：原图 | 热力图 | 叠加+框。人工一眼就能看出热力区域与框合不合。""" - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - from matplotlib.patches import Rectangle - - with Image.open(image_path) as raw: - image = np.asarray(raw.convert("RGB")) - height, width = image.shape[:2] - heat = np.asarray( - Image.fromarray((cam * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR) - ) / 255.0 - - fig, axes = plt.subplots(1, 3, figsize=(11, 4)) - axes[0].imshow(image) - axes[0].set_title("original") - axes[1].imshow(heat, cmap="jet", vmin=0, vmax=1) - axes[1].set_title(f"grad-cam {cam.shape[0]}x{cam.shape[1]}") - axes[2].imshow(image) - axes[2].imshow(heat, cmap="jet", alpha=0.45, vmin=0, vmax=1) - for x1, y1, x2, y2 in boxes: - axes[2].add_patch( - Rectangle( - (x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height, - fill=False, edgecolor="lime", linewidth=2, - ) - ) - axes[2].set_title(f"overlay + {len(boxes)} box") - for ax in axes: - ax.axis("off") - fig.suptitle(title, fontsize=10) - fig.tight_layout() - path.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(path, dpi=110) - plt.close(fig) - - -def main() -> int: - global logger - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--checkpoint", default="runs/classify/best.pth", help="分类模型权重") - parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置") - parser.add_argument("--out", default="runs/gradcam_boxes", help="输出目录") - parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") - parser.add_argument("--device", default=None) - parser.add_argument( - "--stage", type=int, default=6, - help="取骨架 features 的第几层做 CAM（4~7 = layer1~4）；默认 6 即 layer3", - ) - parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值（相对最大值）") - parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例") - parser.add_argument( - "--max-boxes", type=int, default=1, - help="每张图最多保留几个框（按连通域内热力总和排序取前 N 个）。" - "默认 1：一张裁片通常只对应一处缺陷，多出来的框多半是噪声", - ) - parser.add_argument( - "--min-area-ratio", type=float, default=0.02, help="连通域面积下限（占整图比例）" - ) - parser.add_argument( - "--peak-ratio", type=float, default=0.75, - help="次要连通域的峰值至少要达到全局最大值的这个比例才保留，用于滤掉弱响应假框", - ) - parser.add_argument( - "--viz-per-class", type=int, default=20, help="每个类别存多少张可视化，0 = 不存" - ) - parser.add_argument("--viz-all", action="store_true", help="每张图都存可视化（很慢很占盘）") - parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张，用于快速试参数") - args = parser.parse_args() - - out_dir = Path(args.out) - out_dir.mkdir(parents=True, exist_ok=True) - logger = setup_logging(out_dir / "gradcam.log") - - cfg = load_config(args.config, args.set) - try: - device = resolve_device(args.device or cfg.get("device", "auto")) - except DeviceNotAvailableError as exc: - logger.error("%s", exc) - return 1 - - data_cfg = cfg["data"] - dataset = FolderClassificationDataset( - root=data_cfg["root"], - split="all", # 为每个样本都生成，不分训练/验证 - classes=data_cfg.get("classes"), - transforms=data_cfg["val_transforms"], - exclude=data_cfg.get("exclude") or [], - ) - if args.limit: - dataset.samples = dataset.samples[: args.limit] - classes = dataset.classes - - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(classes)) - model_cfg["init_from"] = None # 权重从 checkpoint 来，不要再读自监督骨架 - model = build_model(model_cfg) - load_checkpoint(args.checkpoint, model, map_location=device) - model.to(device).eval() - for param in model.parameters(): # Grad-CAM 需要梯度，但不更新权重 - param.requires_grad_(True) - - logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint) - logger.info("数据集 %s：%d 张，%d 类", data_cfg["root"], len(dataset), len(classes)) - logger.info( - "CAM 目标层 features[%d] | 阈值 %.2f | 外扩 %.2f | 每图最多 %d 框", - args.stage, args.threshold, args.padding, args.max_boxes, - ) - - loader = DataLoader( - dataset, - batch_size=int(data_cfg.get("loader", {}).get("batch_size", 64)), - shuffle=False, - num_workers=int(data_cfg.get("loader", {}).get("num_workers", 4)), - collate_fn=classification_collate, - ) - - labels_dir = out_dir / "labels" - viz_dir = out_dir / "viz" - viz_done: Counter[str] = Counter() - box_counts: Counter[int] = Counter() - per_class_boxes: Counter[str] = Counter() - correct = 0 - cam_shape = None - - for step, batch in enumerate(loader): - images = batch["images"].to(device) - targets = batch["labels"].to(device) - cams = grad_cam(model.backbone, model.head, images, targets, stage=args.stage) - with torch.no_grad(): - correct += int((model(images).argmax(1) == targets).sum()) - cam_shape = tuple(cams.shape[-2:]) - - for cam, target, meta in zip(cams.cpu().numpy(), targets.cpu().tolist(), batch["meta"], - strict=True): - name = classes[target] - boxes = cam_to_boxes( - cam, - threshold=args.threshold, - min_area_ratio=args.min_area_ratio, - max_boxes=args.max_boxes, - padding=args.padding, - peak_ratio=args.peak_ratio, - ) - box_counts[len(boxes)] += 1 - per_class_boxes[name] += len(boxes) - - source = Path(meta["path"]) - label_path = labels_dir / name / f"{source.stem}.txt" - label_path.parent.mkdir(parents=True, exist_ok=True) - label_path.write_text( - "\n".join(to_yolo_line(box, target) for box in boxes) + ("\n" if boxes else ""), - encoding="utf-8", - ) - - if args.viz_all or viz_done[name] < args.viz_per_class: - viz_done[name] += 1 - save_visualization( - viz_dir / name / f"{source.stem}.png", source, cam, boxes, - f"{name} | {len(boxes)} box | thr={args.threshold}", - ) - - if step % 20 == 0: - logger.info("处理 [%d/%d]", step, len(loader)) - - total = sum(box_counts.values()) - write_data_yaml(out_dir / "data.yaml", classes, root=out_dir) - summary = { - "checkpoint": args.checkpoint, - "dataset": str(data_cfg["root"]), - "num_images": total, - "cam_shape": list(cam_shape) if cam_shape else None, - "stage": args.stage, - "threshold": args.threshold, - "padding": args.padding, - "max_boxes": args.max_boxes, - "min_area_ratio": args.min_area_ratio, - "peak_ratio": args.peak_ratio, - "classification_accuracy_on_all": correct / total if total else 0.0, - "boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())}, - "boxes_per_class": dict(sorted(per_class_boxes.items())), - } - (out_dir / "summary.json").write_text( - json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" - ) - - logger.info("CAM 分辨率 %s | 全量分类准确率 %.4f", cam_shape, summary["classification_accuracy_on_all"]) - for count, number in sorted(box_counts.items()): - logger.info(" 生成 %d 个框的图片: %d 张 (%.1f%%)", count, number, 100 * number / total) - logger.info("标签写入 %s，可视化 %d 张写入 %s", labels_dir, sum(viz_done.values()), viz_dir) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/shap_boxes.py b/scripts/shap_boxes.py deleted file mode 100644 index 16c01a7..0000000 --- a/scripts/shap_boxes.py +++ /dev/null @@ -1,239 +0,0 @@ -#!/usr/bin/env python -"""用 SHAP Partition explainer 生成框标签，并与 Grad-CAM 并排对比。 - - uv run python scripts/shap_boxes.py --checkpoint runs/classify/best.pth --per-class 10 - -与 Grad-CAM 的取舍： - -| | Grad-CAM | SHAP Partition | -| --- | --- | --- | -| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） | -| 单张耗时 | 一次 backward，全量 11k 张约 2 分钟 | 几百次前向，单张约 1 秒 | - -所以 SHAP 默认**按类别抽样**跑（`--per-class`），用来判断值不值得为它付出时间； -确实要全量请加 `--all`，注意 11k 张需要数小时。 - -产出与 gradcam_boxes.py 同构：labels/ 是 YOLO 标签，viz/ 是对比图 -（原图 | Grad-CAM 热力+框 | SHAP 热力+框），一眼看出谁定得更准。 -""" - -from __future__ import annotations - -import argparse -import json -import random -import sys -import time -from collections import Counter -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -from datasets import FolderClassificationDataset # noqa: E402 -from datasets.yolo import write_data_yaml # noqa: E402 -from engine import load_checkpoint # noqa: E402 -from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line # noqa: E402 -from engine.shap_explain import shap_attribution, smooth # noqa: E402 -from models import build_model # noqa: E402 -from utils import load_config, setup_logging # noqa: E402 -from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 - -logger = None - - -def save_comparison(path: Path, image_path: Path, panels: list[tuple[str, np.ndarray, list]]) -> None: - """一行画出原图与各方法的热力图+框，便于直接比较定位质量。""" - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - from matplotlib.patches import Rectangle - - with Image.open(image_path) as raw: - image = np.asarray(raw.convert("RGB")) - height, width = image.shape[:2] - - fig, axes = plt.subplots(1, len(panels) + 1, figsize=(4 * (len(panels) + 1), 4)) - axes[0].imshow(image) - axes[0].set_title("original") - axes[0].axis("off") - - for ax, (title, heat, boxes) in zip(axes[1:], panels, strict=True): - resized = np.asarray( - Image.fromarray((heat * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR) - ) / 255.0 - ax.imshow(image) - ax.imshow(resized, cmap="jet", alpha=0.45, vmin=0, vmax=1) - for x1, y1, x2, y2 in boxes: - ax.add_patch( - Rectangle((x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height, - fill=False, edgecolor="lime", linewidth=2) - ) - ax.set_title(f"{title} ({heat.shape[0]}x{heat.shape[1]}, {len(boxes)} box)") - ax.axis("off") - - fig.tight_layout() - path.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(path, dpi=110) - plt.close(fig) - - -def main() -> int: - global logger - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--checkpoint", default="runs/classify/best.pth") - parser.add_argument("--config", default="configs/classify/film.yaml") - parser.add_argument("--out", default="runs/shap_boxes") - parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") - parser.add_argument("--device", default=None) - parser.add_argument("--per-class", type=int, default=10, help="每个类别抽多少张") - parser.add_argument("--all", action="store_true", help="跑全量（11k 张需数小时）") - parser.add_argument("--seed", type=int, default=0, help="抽样随机种子") - parser.add_argument("--max-evals", type=int, default=500, help="SHAP 每张图的前向次数") - parser.add_argument("--smooth-radius", type=int, default=2, help="归因图方框平滑半径，0=不平滑") - parser.add_argument("--threshold", type=float, default=0.4) - parser.add_argument("--padding", type=float, default=0.1) - parser.add_argument( - "--max-boxes", type=int, default=1, - help="每张图最多保留几个框（按归因总量排序取前 N 个）。" - "默认 1：一张裁片通常只对应一处缺陷，多出来的框多半是噪声", - ) - parser.add_argument( - "--min-area-ratio", type=float, default=0.002, - help="连通域面积下限（占整图比例）。默认值远小于 gradcam_boxes.py 的 0.02——" - "SHAP 的归因图是 128x128 而不是 8x8，同样一个小黑点在这里只占百来像素，" - "沿用 0.02（=328 像素）会把它整个滤掉", - ) - parser.add_argument("--peak-ratio", type=float, default=0.75) - parser.add_argument("--stage", type=int, default=6, help="对比用的 Grad-CAM 目标层") - parser.add_argument("--no-compare", action="store_true", help="不算 Grad-CAM，只出 SHAP") - args = parser.parse_args() - - out_dir = Path(args.out) - out_dir.mkdir(parents=True, exist_ok=True) - logger = setup_logging(out_dir / "shap.log") - - cfg = load_config(args.config, args.set) - try: - device = resolve_device(args.device or cfg.get("device", "auto")) - except DeviceNotAvailableError as exc: - logger.error("%s", exc) - return 1 - - data_cfg = cfg["data"] - dataset = FolderClassificationDataset( - root=data_cfg["root"], split="all", classes=data_cfg.get("classes"), - transforms=data_cfg["val_transforms"], exclude=data_cfg.get("exclude") or [], - ) - classes = dataset.classes - - if not args.all: - by_class: dict[int, list] = {} - for sample in dataset.samples: - by_class.setdefault(sample[1], []).append(sample) - rng = random.Random(args.seed) - dataset.samples = [ - sample - for samples in by_class.values() - for sample in rng.sample(samples, min(args.per_class, len(samples))) - ] - - model_cfg = dict(cfg["model"]) - model_cfg.setdefault("num_classes", len(classes)) - model_cfg["init_from"] = None - model = build_model(model_cfg) - load_checkpoint(args.checkpoint, model, map_location=device) - model.to(device).eval() - for param in model.parameters(): - param.requires_grad_(True) - - logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint) - logger.info( - "样本 %d 张（%s）| max_evals=%d | 平滑半径 %d", - len(dataset), "全量" if args.all else f"每类 {args.per_class} 张", - args.max_evals, args.smooth_radius, - ) - - box_kwargs = { - "threshold": args.threshold, "min_area_ratio": args.min_area_ratio, - "max_boxes": args.max_boxes, "padding": args.padding, "peak_ratio": args.peak_ratio, - } - labels_dir = out_dir / "labels" - viz_dir = out_dir / "viz" - shap_counts: Counter[int] = Counter() - cam_counts: Counter[int] = Counter() - started = time.time() - - for index in range(len(dataset)): - sample = dataset[index] - target = sample["label"] - name = classes[target] - source = Path(sample["meta"]["path"]) - image_chw = sample["image"] - - attribution = shap_attribution( - model, image_chw.permute(1, 2, 0).numpy().astype(np.float32), target, device, - max_evals=args.max_evals, - ) - attribution = smooth(attribution, args.smooth_radius) - shap_boxes = cam_to_boxes(attribution, **box_kwargs) - shap_counts[len(shap_boxes)] += 1 - - panels = [] - if not args.no_compare: - cam = grad_cam( - model.backbone, model.head, image_chw[None].to(device), - torch.tensor([target], device=device), stage=args.stage, - )[0].cpu().numpy() - cam_boxes = cam_to_boxes(cam, **box_kwargs) - cam_counts[len(cam_boxes)] += 1 - panels.append(("grad-cam", cam, cam_boxes)) - panels.append(("shap", attribution, shap_boxes)) - - label_path = labels_dir / name / f"{source.stem}.txt" - label_path.parent.mkdir(parents=True, exist_ok=True) - label_path.write_text( - "\n".join(to_yolo_line(box, target) for box in shap_boxes) + ("\n" if shap_boxes else ""), - encoding="utf-8", - ) - save_comparison(viz_dir / name / f"{source.stem}.png", source, panels) - - if index % 20 == 0: - elapsed = time.time() - started - logger.info( - "[%d/%d] 已用 %.0fs，平均 %.2fs/张", - index, len(dataset), elapsed, elapsed / max(index, 1), - ) - - total = sum(shap_counts.values()) - write_data_yaml(out_dir / "data.yaml", classes, root=out_dir) - summary = { - "checkpoint": args.checkpoint, - "num_images": total, - "seconds_per_image": (time.time() - started) / max(total, 1), - "max_evals": args.max_evals, - "smooth_radius": args.smooth_radius, - "box_params": box_kwargs, - "shap_boxes_per_image": {str(k): v for k, v in sorted(shap_counts.items())}, - "gradcam_boxes_per_image": {str(k): v for k, v in sorted(cam_counts.items())}, - } - (out_dir / "summary.json").write_text( - json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" - ) - - logger.info("共 %d 张，平均 %.2fs/张", total, summary["seconds_per_image"]) - logger.info("SHAP 出框分布: %s", dict(sorted(shap_counts.items()))) - if cam_counts: - logger.info("Grad-CAM 出框分布: %s", dict(sorted(cam_counts.items()))) - logger.info("对比图写入 %s", viz_dir) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_generate_boxes.py b/tests/test_generate_boxes.py new file mode 100644 index 0000000..cef7329 --- /dev/null +++ b/tests/test_generate_boxes.py @@ -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) # 同名，撞同一个标签路径