ModelTrainEval/detect.py

#!/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())