Shaun ZhouandClaude Opus 5
9439cbb56a
ultralytics 没有可用的开关：yolo26 是端到端模型，non_max_suppression 对它早退 （utils/nms.py 里 end2end 分支），只按置信度与 max_det 截断，传入的 iou 和 agnostic 都不参与 IoU 抑制；agnostic_nms 只作用于检测头 top-k（每个 anchor 取一个类别）， 压不掉不同 anchor 产生的重叠框。 所以给 DetectionValidator.postprocess 包一层 filter_overlapping_outputs。训练期验证、 早停、best.pt 的选择都随之遵守这条约束，指标与最终推理一致。配置项 train.max_overlap （默认 0，负数关闭），由 detect.py 摘出后自行实现，不传给 ultralytics。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
187 lines
7.6 KiB
Python
187 lines
7.6 KiB
Python
#!/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
|
||
|
||
数据是在 Label Studio 里修正过的标签，导出并划分 train/val 之后放在 data/ 下的数据集。
|
||
框来自 Grad-CAM / SHAP 的弱监督标注再经人工修正。
|
||
|
||
**预测框约束为互不重叠**（`train.max_overlap`，默认 0）。真值里一处缺陷只有一个框、
|
||
零重叠，模型却会输出重叠框，所以验证阶段补一层抑制——早停与 best.pt 的选择都据此进行，
|
||
指标才和最终推理一致。设为负数可关掉。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import yaml
|
||
|
||
from engine.detection import filter_overlapping_outputs
|
||
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 install_overlap_suppression(max_overlap: float) -> None:
|
||
"""给验证器的后处理包一层"框不重叠"约束。
|
||
|
||
只能用打补丁的方式：ultralytics 没有对应的开关——端到端模型（yolo26）的 NMS 早退，
|
||
传入的 iou / agnostic_nms 都不参与 IoU 抑制（见 engine/detection.filter_overlapping_outputs
|
||
的说明）。补在 postprocess 之后，训练期验证、早停、best.pt 的选择就都遵守这条约束。
|
||
"""
|
||
from ultralytics.models.yolo.detect import DetectionValidator
|
||
|
||
original = DetectionValidator.postprocess
|
||
|
||
def postprocess(self, preds):
|
||
return filter_overlapping_outputs(original(self, preds), max_iou=max_overlap)
|
||
|
||
DetectionValidator.postprocess = postprocess
|
||
|
||
|
||
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；先把 Label Studio 里修好的标签导出、划分 train/val 到 data/ 下", data_yaml
|
||
)
|
||
return 1
|
||
meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
|
||
logger.info("数据集: %s（%d 类）", data_yaml, meta["nc"])
|
||
if meta.get("train") == meta.get("val"):
|
||
# generate_boxes.py 的产出就不划分，两者指向同一目录
|
||
logger.warning(
|
||
"data.yaml 的 train 与 val 相同（%s）：验证集就是训练集，mAP 会虚高，"
|
||
"早停与 best.pt 的选择也因此失真。正式训练前请先划分数据集",
|
||
meta.get("train"),
|
||
)
|
||
|
||
train_cfg = dict(cfg.get("train", {}))
|
||
swanlab_cfg = dict(train_cfg.pop("swanlab", {}))
|
||
# 不是 ultralytics 的参数，得先摘出来，否则它会拒绝未知键
|
||
max_overlap = float(train_cfg.pop("max_overlap", 0.0))
|
||
|
||
from ultralytics import YOLO
|
||
from ultralytics import settings as ul_settings
|
||
|
||
if max_overlap >= 0:
|
||
install_overlap_suppression(max_overlap)
|
||
logger.info("验证时约束预测框互不重叠（最大 IoU %.2f）", max_overlap)
|
||
else:
|
||
logger.warning("未约束预测框重叠：指标里会计入真值中不可能出现的重叠框")
|
||
|
||
# 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())
|