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
@@ -632,52 +632,64 @@ model:
|
|||||||
只加载 backbone，分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量
|
只加载 backbone，分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量
|
||||||
自监督特征本身有多好。
|
自监督特征本身有多好。
|
||||||
|
|
||||||
## 从分类到定位：Grad-CAM 生成框标签
|
## 从分类到定位：生成框标签
|
||||||
|
|
||||||
分类模型只学过"这是什么缺陷"，没学过"缺陷在哪"。Grad-CAM 把分类分数对卷积特征的梯度
|
分类模型只学过"这是什么缺陷"，没学过"缺陷在哪"。归因方法能把分类分数反推回空间位置，
|
||||||
当作通道重要性，加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 +
|
阈值化 + 取连通域外接矩形，就得到弱监督的框标签：
|
||||||
取连通域外接矩形，就得到弱监督的框标签：
|
|
||||||
|
|
||||||
```bash
|
```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/
|
runs/<name>/
|
||||||
├── labels/<类别>/<图名>.txt # YOLO 格式 class_id cx cy w h，可直接喂给检测训练
|
├── images/<类别>/<图名>.png # 硬链接到原图，不占额外磁盘
|
||||||
├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框，供人工核对
|
├── labels/<类别>/<图名>.txt # YOLO：class_id cx cy w h（归一化）
|
||||||
|
├── heatmaps/<类别>/<图名>.png # 热力图叠加原图 + 框，供人工核对
|
||||||
├── data.yaml
|
├── data.yaml
|
||||||
└── summary.json
|
└── 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 个格子，框很粗 |
|
| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） |
|
||||||
| **6 (layer3)** | **8×8** | **76% 出 1 框、20% 出 2 框，热力区域与缺陷贴合** |
|
| 单张耗时 | ~0.01 秒 | ~0.55 秒 |
|
||||||
| 5 (layer2) | 16×16 | 40% 出不了框，响应太散 |
|
| 小目标 | 几像素的黑点在 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）：热力图相对最大值的阈值
|
- `--threshold`（默认 0.4）：热力图相对最大值的阈值
|
||||||
- `--max-boxes`（默认 1）：每张图最多保留几个框，按连通域内热力总和排序取前 N 个。
|
- `--min-area-ratio`：**缺省随方法而变**（gradcam 0.02、shap 0.002）。两者归因图分辨率
|
||||||
一张裁片通常只对应一处缺陷，多出来的框多半是噪声；确实有多缺陷的图再调大
|
差 16 倍，同一个小黑点在 128×128 上只占百来像素，沿用 0.02（=328 像素）会被整个滤掉
|
||||||
- `--peak-ratio`（默认 0.75）：次要连通域的峰值要达到全局最大值的这个比例才保留。
|
- `--peak-ratio`（默认 0.75）：次要连通域的峰值需达全局最大值的这个比例，滤掉弱响应假框
|
||||||
不加这条会在没有缺陷的角落刷出假框
|
- `--padding`（默认 0.1）：框四周外扩。归因倾向只覆盖最具判别性的部位，往往偏小
|
||||||
- `--padding`（默认 0.1）：框四周外扩。CAM 倾向只覆盖最具判别性的部位（缺陷中心），
|
- `--stage`（gradcam，默认 6=layer3）：layer4 只有 4×4 太粗，layer2 响应太散
|
||||||
往往比真实缺陷范围小
|
- `--no-heatmaps` / `--no-images`：只要标签时可跳过
|
||||||
- `--viz-per-class`（默认 20）：每类存多少张可视化；`--viz-all` 全存（很慢很占盘）
|
|
||||||
|
|
||||||
> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 这类描述的是整图属性而非
|
> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 描述的是整图属性而非局部缺陷，
|
||||||
> 局部缺陷，它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/
|
> 它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/`grease` 这些
|
||||||
> `grease`/`zigzag` 这些有局部目标的类别。
|
> 有局部目标的类别。
|
||||||
|
|
||||||
## 方法
|
## 方法
|
||||||
|
|
||||||
|
|||||||
@@ -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/<name>/
|
||||||
|
├── 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())
|
||||||
@@ -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())
|
|
||||||
@@ -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())
|
|
||||||
@@ -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