Shaun ZhouandClaude Opus 5
65bc6b2045
背景差分：大半径方框均值估背景 → 残差（可选极性）→ 平滑 → 有无判定 → 按峰值归一化， 再交给与 gradcam / shap 同一套 cam_to_boxes 出框。只依赖 numpy。 与归因方法的差别：不需要 --checkpoint，分类器还没训练出来就能先跑一版框，也不受分类 精度影响；类别只能来自目录（--target label）；找的是偏离局部背景的区域，而不是模型 据以判断类别的区域，所以对局部目标好、对铺满整图的纹理类给不出有意义的框。 有无缺陷是显式判定的：残差峰值要同时高于绝对对比度下限与噪声水平（MAD 估计）的 15 倍， 否则返回全零图。这条是必需的——实测纯噪声图的"峰值/MAD"稳定在 9 附近（最大 10.9， 几乎与噪声大小无关），有 0.25 对比度黑点时是 25~330；若只按峰值归一化，正常图的噪声 会被拉到 1，满图假框。阈值 15 取在两者之间。 顺带把积分图方框均值从 shap_explain.smooth 提到 engine/classical.box_mean 共用， 避免两处实现各自漂移。 实测：单张 0.011 秒（全量 11k 约 2 分钟）；真实薄膜图上框能套住黑点、虫状物与划痕。 已知局限：背景估计在图像边界处对明暗阶跃有响应，贴边的交界可能出假框。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
460 lines
20 KiB
Python
460 lines
20 KiB
Python
#!/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
|
||
|
||
# 传统 CV：背景差分，不需要分类模型
|
||
uv run python scripts/generate_boxes.py --out runs/boxes_cv --method cv
|
||
|
||
三种方法都能定位，取舍不同：
|
||
|
||
| | gradcam | shap | cv |
|
||
| --- | --- | --- | --- |
|
||
| 依赖分类模型 | 是 | 是 | **否** |
|
||
| 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） | 原图分辨率 |
|
||
| 单张耗时 | 一次 backward，约 0.01 秒 | 几百次前向，约 0.55 秒 | 几次方框滤波，约 0.005 秒 |
|
||
| 小目标 | 几像素的黑点在 8×8 上不可分辨 | 能贴合到几十像素的目标 | 贴合好，但纹理噪声也会被找出来 |
|
||
| 找的是什么 | 模型据以判断类别的区域 | 同左，更精细 | 偏离局部背景的区域 |
|
||
|
||
cv 不判类别，只定位，所以它只能配 `--target label`（默认），且**不需要 --checkpoint**：
|
||
分类器还没训练时就能先跑一版框出来看。它对 blackdot / whitedot / mosquito 这类局部目标
|
||
效果好；ripple / wrinkle 这种铺满整图的纹理类，本来就没有局部目标，它给不出有意义的框。
|
||
|
||
产出是一个**自成一体的 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` 只用于快速试参数。
|
||
热力图与框的类别默认取**图片所在目录的人工标注类别**（`--target label`）：类别的真值是人工归类，
|
||
分类模型只负责回答"为什么是这个类"，也就是框在哪里。`--target predicted` 改用模型的 top1 预测，
|
||
这时模型分错的图，框的类别也会跟着错。
|
||
"""
|
||
|
||
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", "cv")
|
||
|
||
#: 各方法的归因图分辨率差着数量级，面积阈值不能共用一个默认值：
|
||
#: 同一个小黑点在 8×8 上占几格、在 128×128 上只占百来像素，
|
||
#: 沿用 0.02（=328 像素）会把它整个滤掉。cv 在原图分辨率上跑，与 shap 同量级。
|
||
DEFAULT_MIN_AREA_RATIO = {"gradcam": 0.02, "shap": 0.002, "cv": 0.002}
|
||
|
||
#: cv 的分析分辨率上限。框坐标是归一化的，缩小不影响标签，却能挡住
|
||
#: connected_components 那个纯 Python 的 BFS 在大图上变慢（数据里最大有 700×712）。
|
||
CV_MAX_SIDE = 256
|
||
|
||
|
||
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="label",
|
||
help="热力图与框针对哪个类别：label=图片所在目录的人工标注（默认），predicted=模型 top1 预测",
|
||
)
|
||
|
||
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 的目标层（torchvision 骨架：4~7 = layer1~4）。默认 6 即 layer3（8×8）；"
|
||
"layer4 只有 4×4 太粗，layer2 响应太散。-1 表示用骨架最终输出——"
|
||
"tiny_conv 这类不是 nn.Sequential 的骨架只能这样取",
|
||
)
|
||
parser.add_argument("--max-evals", type=int, default=500, help="shap 每张图的前向次数")
|
||
parser.add_argument(
|
||
"--smooth-radius", type=int, default=2, help="shap / cv 归因图平滑半径，0=不平滑",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--cv-radius", type=int, default=0,
|
||
help="cv 的背景估计半径，0=按图取 min(高,宽)//8。太小会把缺陷本身算进背景",
|
||
)
|
||
parser.add_argument(
|
||
"--cv-polarity", choices=("dark", "bright", "both"), default="both",
|
||
help="cv 找哪种缺陷：dark=比背景暗（黑点/蚊虫），bright=更亮（白点/反光），both=两者",
|
||
)
|
||
parser.add_argument(
|
||
"--cv-min-contrast", type=float, default=0.05,
|
||
help="cv 的绝对对比度下限（0~1 灰度，0.05≈13 个灰阶）。"
|
||
"残差峰值低于它就判定整图无缺陷，出 0 个框",
|
||
)
|
||
parser.add_argument(
|
||
"--cv-max-side", type=int, default=CV_MAX_SIDE,
|
||
help="cv 分析前把长边缩到这个上限（框是归一化的，不影响标签），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]
|
||
|
||
# cv 不判类别也不做归因；只有需要模型预测时才加载权重，
|
||
# 这样分类器还没训练出来也能先跑一版 cv 框
|
||
needs_model = args.method != "cv" or args.target == "predicted"
|
||
model = None
|
||
if needs_model:
|
||
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)
|
||
|
||
stage = args.stage if args.stage >= 0 else None
|
||
if args.method == "gradcam" and stage is not None and not isinstance(
|
||
getattr(model.backbone, "features", None), torch.nn.Sequential
|
||
):
|
||
# tiny_conv 这类骨架没有 nn.Sequential 的 features，取不了中间层；
|
||
# 退回最终输出而不是直接崩掉
|
||
logger.warning(
|
||
"骨架 %s 不是 nn.Sequential，无法按 stage 取中间层，改用最终输出"
|
||
"（分辨率更低，如需更精细请换 torchvision 骨架）",
|
||
type(model.backbone).__name__,
|
||
)
|
||
stage = None
|
||
|
||
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,
|
||
}
|
||
cv_kwargs = {
|
||
"background_radius": args.cv_radius, "polarity": args.cv_polarity,
|
||
"smooth_radius": args.smooth_radius, "min_contrast": args.cv_min_contrast,
|
||
}
|
||
|
||
logger.info(
|
||
"设备: %s | 权重: %s", describe_device(device),
|
||
args.checkpoint if needs_model else "不需要（cv + --target label）",
|
||
)
|
||
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 == "cv":
|
||
logger.info("cv 参数: %s", cv_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=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))
|
||
elif args.method == "cv":
|
||
from engine.classical import defect_saliency
|
||
|
||
for index, (path, label) in enumerate(dataset.samples):
|
||
target = label
|
||
if needs_model: # --target predicted：还是得过一遍模型
|
||
with torch.no_grad():
|
||
predicted = int(model(dataset[index]["image"][None].to(device)).argmax(1))
|
||
agree += int(predicted == label)
|
||
target = predicted
|
||
|
||
with Image.open(path) as raw:
|
||
image = raw.convert("L")
|
||
if args.cv_max_side:
|
||
image.thumbnail((args.cv_max_side, args.cv_max_side), Image.BILINEAR)
|
||
heat = defect_saliency(np.asarray(image), **cv_kwargs)
|
||
heat_shape = heat.shape
|
||
try:
|
||
_emit(out_dir, Path(path), classes, target, heat, box_kwargs,
|
||
args, box_counts, link_modes, seen)
|
||
except NameCollision as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
if index % 500 == 0:
|
||
logger.info("处理 [%d/%d]", index, 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": stage if args.method == "gradcam" else None,
|
||
"max_evals": args.max_evals if args.method == "shap" else None,
|
||
"cv_params": cv_kwargs if args.method == "cv" else None,
|
||
# cv + --target label 时根本没加载模型，没有预测可比
|
||
"prediction_matches_folder_label": (agree / total if total else 0.0) if needs_model else None,
|
||
"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"]
|
||
)
|
||
if summary["prediction_matches_folder_label"] is not None:
|
||
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())
|