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:
co-authored by
Claude Opus 5
parent
86c86b105c
commit
8dc2fe6a2e
@@ -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())
|
||||
Reference in New Issue
Block a user