feat(detect): 把弱监督框数据集链接到 data/ 并接上 YOLO 训练

第五步产出的框标签按类别分目录,便于人工核对;检测训练要的是 train/val
划分——类别写在标签第一列,不靠目录名。补上两个入口:

- scripts/prepare_yolo_dataset.py:拍平类别目录、按类别分层切分、硬链接到
  data/。分层是必需的,black 只有 53 张,随机切容易让小类在验证集里挂零。
  无框的图默认保留(检测训练里是压低误检的负样本);拍平后的同名冲突直接
  报错而不是静默覆盖。
- detect.py + configs/detect/film.yaml:包一层 ultralytics,沿用本仓库的
  配置继承、设备解析与 SwanLab 可视化。

选 YOLOv8 而不是 v11/v12:CLAUDE.md 限定只用 convolution 模型,YOLO11 骨架
里有 C2PSA、YOLO12 以注意力为主,只有 v8 的 C2f + SPPF 是纯卷积。
mosaic/mixup 关掉——每张图只有一个弱监督框,拼接后小目标被缩得更小,对
blackdot/whitedot 这类只占千分之几面积的类别是净伤害。

ultralytics 的 project 必须传绝对路径,相对路径会被它再拼到全局 settings
的 runs_dir 上,产出会跑到 runs/detect/runs/detect/<name> 去。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-11 12:05:50 +08:00
co-authored by Claude Opus 5
parent 86c86b105c
commit 8dc2fe6a2e
8 changed files with 816 additions and 13 deletions
+150
View File
@@ -0,0 +1,150 @@
#!/usr/bin/env python
"""在弱监督框标签上训练 YOLO 检测模型。
uv run python detect.py --config configs/detect/film.yaml
uv run python detect.py --config configs/detect/film.yaml --eval-only \
--checkpoint runs/detect/film/weights/best.pt
数据是 scripts/prepare_yolo_dataset.py 的产物(train/val 已划分好)。
框来自 Grad-CAM / SHAP 的弱监督标注,不是人工标注——指标该按"与热力区域的一致性"
来读,而不是当成真值上的检测精度。
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import yaml
from utils import load_config, seed_everything, setup_logging
from utils.device import DeviceNotAvailableError, describe_device, resolve_device
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="薄膜缺陷检测(YOLO)训练/评估")
parser.add_argument("--config", required=True, help="yaml 配置路径")
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
parser.add_argument("--work-dir", default=None)
parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu")
parser.add_argument("--resume", action="store_true", help="从 work_dir 里的 last.pt 续训")
parser.add_argument("--eval-only", action="store_true", help="只在验证集上评估")
parser.add_argument("--checkpoint", default=None, help="评估/微调用的权重 .pt")
return parser.parse_args()
def to_ultralytics_device(device) -> str | int:
"""torch.device -> ultralytics 的 device 记法(cuda 用序号,其余用字符串)。"""
if device.type == "cuda":
return device.index or 0
return device.type
def main() -> int:
args = parse_args()
cfg = load_config(args.config, overrides=args.set)
if args.work_dir:
cfg["work_dir"] = args.work_dir
work_dir = Path(cfg.get("work_dir", "runs/detect"))
work_dir.mkdir(parents=True, exist_ok=True)
logger = setup_logging(log_file=work_dir / "detect.log")
seed = int(cfg.get("seed", 42))
seed_everything(seed)
try:
device = resolve_device(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
logger.info("设备: %s", describe_device(device))
data_yaml = Path(cfg["data"]["yaml"])
if not data_yaml.exists():
logger.error(
"找不到 %s;先跑 scripts/prepare_yolo_dataset.py 把 runs/ 下的产出链接到 data/", data_yaml
)
return 1
meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
logger.info("数据集: %s(%d 类)", data_yaml, meta["nc"])
train_cfg = dict(cfg.get("train", {}))
swanlab_cfg = dict(train_cfg.pop("swanlab", {}))
from ultralytics import YOLO
from ultralytics import settings as ul_settings
# ultralytics 默认把数据集根指到 ~/datasets 之类的全局目录;我们的 data.yaml 里写的是
# 绝对路径,但它仍会拿 datasets_dir 去拼相对路径,统一指到项目根省得踩坑。
ul_settings.update({"datasets_dir": str(Path.cwd())})
weights = args.checkpoint or cfg["model"]["weights"]
logger.info("模型: %s", weights)
model = YOLO(weights)
if swanlab_cfg.pop("enabled", False):
try:
from swanlab.integration.ultralytics import add_swanlab_callback
os.environ.setdefault("SWANLAB_PROJ_NAME", str(swanlab_cfg.get("project", "film-detect")))
import swanlab
swanlab.init(
project=swanlab_cfg.get("project", "film-detect"),
name=swanlab_cfg.get("name") or work_dir.name,
mode=swanlab_cfg.get("mode", "local"),
config=cfg,
log_dir=str(work_dir / "swanlog"),
tags=swanlab_cfg.get("tags") or [],
)
add_swanlab_callback(model)
logger.info("SwanLab 已接入 YOLO 回调")
except Exception as exc: # 可视化失败不该挡住训练
logger.warning("SwanLab 接入失败,训练继续: %s", exc)
# ultralytics 自己管 project/name 目录。project 必须给绝对路径:相对路径会被它再拼到
# 全局 settings 的 runs_dir 上,产出就跑到 runs/detect/runs/detect/<name> 去了。
common = {
"data": str(data_yaml.resolve()),
"device": to_ultralytics_device(device),
"project": str(work_dir.resolve().parent),
"name": work_dir.name,
"exist_ok": True,
"seed": seed,
}
if args.eval_only:
metrics = model.val(split="val", plots=True, **common)
else:
model.train(resume=args.resume, **common, **train_cfg)
# train 结束后 ultralytics 已在最优权重上跑过一次 val,这里显式再跑一次拿回指标对象
metrics = model.val(split="val", plots=True, **common)
summary = {
"map50_95": float(metrics.box.map),
"map50": float(metrics.box.map50),
"precision": float(metrics.box.mp),
"recall": float(metrics.box.mr),
"per_class_map50": {
meta["names"][int(c)]: float(v)
for c, v in zip(metrics.box.ap_class_index, metrics.box.ap50, strict=False)
},
}
logger.info(
"mAP50-95 %.4f | mAP50 %.4f | P %.4f | R %.4f",
summary["map50_95"], summary["map50"], summary["precision"], summary["recall"],
)
(work_dir / "metrics.json").write_text(
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
)
logger.info(
"曲线与混淆矩阵见 %s(results.png / *_curve.png / confusion_matrix*.png)",
work_dir.resolve(),
)
return 0
if __name__ == "__main__":
raise SystemExit(main())