Skip to content

YOLO26 + COCO128 目标检测 跑通环境

目标

跑通环境

yolo预设权重,已经使用coco数据集训练过。 coco128数据集训练,本身没有意义。

在线演示

在线演示

项目目录

- coco128.py
- test_image.jpg
- train/

完整代码

  • coco128.py
python
"""
YOLO26 + COCO128 目标检测完整流程
===================================
模型:    YOLO26 (Ultralytics 最新版, 端到端无NMS, MuSGD优化器)
数据集:  coco128 (128张图, COCO 80类, 自动下载)
参考:   https://docs.ultralytics.com/zh/models/yolo26/

使用方法:
    1. 安装: pip3 install ultralytics
    2. 运行: python3 coco128.py

设备说明:
    - Mac (M系列): device="mps"  (GPU加速)
    - NVIDIA GPU:  device="0"
    - CPU:         device="cpu"
"""

from ultralytics import YOLO
import os
import sys


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 配置区 —— 改动参数在这里改
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
CONFIG = {
    # YOLO26 模型尺寸: n(nano) / s(small) / m(medium) / l(large) / x(超大型)
    # 精度参考: n=40.9mAP / s=48.6 / m=53.1 / l=55.0 / x=57.5
    "model_name": "yolo26n.pt",

    # 数据集路径: coco128(快速验证128张) / coco.yaml(完整118k张)
    "data_yaml": "/Users/ly/Downloads/dataset/coco128/data.yaml",

    # 训练参数
    "epochs": 5,                # 环境验证 5 轮足够; 正式训练 50-300
    "imgsz": 640,              # 输入尺寸: 320/416/640/1280
    "batch": 8,                # 批次大小, 显存不够调小
    "device": "mps",           # 运行设备

    # 其他
    "project": "train",    # 输出目录
    "name": "yolo26_coco128",   # 实验名称
    "test_img": "./test_image.jpg",  # 测试图片
}



# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 1: 加载模型
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step1_load_model():
    print("\n" + "─" * 50)
    print("Step 1: 加载 YOLO26 预训练模型")
    print("─" * 50)
    model = YOLO(CONFIG["model_name"])
    print(f"  模型:        {CONFIG['model_name']}")
    print(f"  任务类型:    {model.task}")
    print(f"  类别数:      {len(model.names)}")
    print(f"  模型说明:    端到端无NMS, MuSGD优化器")
    print(f"  类别列表:   {list(model.names.values())}")
    return model


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 2: 推理测试(验证环境正常)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step2_predict(model):
    print("\n" + "─" * 50)
    print("Step 2: 推理测试 (预训练权重)")
    print("─" * 50)
    results = model.predict(
        source=CONFIG["test_img"],
        conf=0.25,
        save=True,
        verbose=False,
    )
    count = 0
    for r in results:
        if r.boxes is None or len(r.boxes) == 0:
            print("  未检测到目标(可能网络图片加载失败)")
        else:
            for box in r.boxes:
                label = model.names[int(box.cls)]
                conf = float(box.conf)
                xyxy = [round(x, 1) for x in box.xyxy[0].tolist()]
                print(f"  ✅ {label:<15s} conf={conf:.2f}  box={xyxy}")
                count += 1
    print(f"  共检测到 {count} 个目标")
    print("  结果图已保存到 runs/detect/predict/")
    return results


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 3: 训练
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step3_train(model):
    print("\n" + "─" * 50)
    print("Step 3: 开始训练 (COCO128)")
    print("─" * 50)
    print(f"  模型:        {CONFIG['model_name']} (YOLO26)")
    print(f"  数据集:      {CONFIG['data_yaml']} (自动下载)")
    print(f"  训练轮数:    {CONFIG['epochs']}")
    print(f"  图片尺寸:    {CONFIG['imgsz']}")
    print(f"  批次大小:    {CONFIG['batch']}")
    print(f"  运行设备:    {CONFIG['device']}")
    print("  训练中,请稍候...\n")

    results = model.train(
        data=CONFIG["data_yaml"],
        epochs=CONFIG["epochs"],
        imgsz=CONFIG["imgsz"],
        batch=CONFIG["batch"],
        device=CONFIG["device"],
        project=CONFIG["project"],
        name=CONFIG["name"],
        workers=4,
        pretrained=True,
        verbose=True,
    )

    save_dir = results.save_dir
    best_path = os.path.join(save_dir, "weights", "best.pt")
    last_path = os.path.join(save_dir, "weights", "last.pt")

    print(f"\n  ✅ 训练完成!")
    print(f"  最佳权重: {best_path}")
    print(f"  最终权重: {last_path}")
    return results, save_dir


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 4: 验证评估
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step4_validate(best_path):
    print("\n" + "─" * 50)
    print("Step 4: 验证评估")
    print("─" * 50)
    best_model = YOLO(best_path)
    metrics = best_model.val(
        data=CONFIG["data_yaml"],
        imgsz=CONFIG["imgsz"],
        batch=CONFIG["batch"],
        device=CONFIG["device"],
    )
    print(f"  mAP50:      {metrics.box.map50:.4f}")
    print(f"  mAP50-95:   {metrics.box.map:.4f}")
    print(f"  Precision:  {metrics.box.mp:.4f}")
    print(f"  Recall:     {metrics.box.mr:.4f}")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 5: 微调模型推理
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step5_test_finetuned(best_path):
    print("\n" + "─" * 50)
    print("Step 5: 微调模型推理测试")
    print("─" * 50)


    model = YOLO(best_path)

    results = model(CONFIG["test_img"])

    count = 0
    for result in results:
        boxes = result.boxes
        for box in boxes:
            cls_id = int(box.cls[0])             # 类别 ID
            cls_name = result.names[cls_id]       # 类别名称
            confidence = float(box.conf[0])       # 置信度
            x1, y1, x2, y2 = box.xyxy[0].tolist()  # 边界框坐标
            print(f"检测到: {cls_name}, 置信度: {confidence:.2f}, 坐标: ({x1:.0f},{y1:.0f})-({x2:.0f},{y2:.0f})")
            count += 1

        result.show()   # 显示结果图片
        result.save()   # 保存结果到 runs/detect/predict/

    print(f"  共检测到 {count} 个目标")

    # results = model.predict(
    #     source=CONFIG["test_img"],
    #     conf=0.25,
    #     save=True,
    #     verbose=False,
    # )
    # count = 0
    # for r in results:
    #     if r.boxes is None or len(r.boxes) == 0:
    #         print("  未检测到目标")
    #     else:
    #         for box in r.boxes:
    #             label = model.names[int(box.cls)]
    #             conf = float(box.conf)
    #             xyxy = [round(x, 1) for x in box.xyxy[0].tolist()]
    #             print(f"  ✅ {label:<15s} conf={conf:.2f}  box={xyxy}")
    #             count += 1
    # print(f"  共检测到 {count} 个目标")
    # print("  结果图已保存到 runs/detect/predict/")
    return results


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Step 6: 导出部署(可选)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def step6_export(best_path):
    print("\n" + "─" * 50)
    print("Step 6: 导出模型")
    print("─" * 50)
    export_model = YOLO(best_path)
    formats = ["onnx", "torchscript", "coreml"]
    for fmt in formats:
        try:
            path = export_model.export(format=fmt)
            print(f"  ✅ {fmt:12s}{path}")
        except Exception as e:
            print(f"  ⚠️  {fmt:12s} 导出失败: {e}")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 主函数
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
def main():
    print("=" * 50)
    print("  YOLO26 + COCO128 目标检测完整流程")
    print("=" * 50)

    try:
        # Step 1: 加载模型
        model = step1_load_model()

        # Step 2: 推理测试
        step2_predict(model)

        # Step 3: 训练
        train_results, save_dir = step3_train(model)

        # Step 4: 验证
        best_path = os.path.join(save_dir, "weights", "best.pt")
        step4_validate(best_path)

        # Step 5: 微调后推理
        step5_test_finetuned(best_path)

        # Step 6: 导出(可选,取消注释启用)
        step6_export(best_path)

        # 完成
        print("\n" + "=" * 50)
        print("  🎉 全部完成!")
        print("=" * 50)
        print(f"""
  模型:        {CONFIG['model_name']} (YOLO26)
  数据集:      {CONFIG['data_yaml']}
  训练轮数:    {CONFIG['epochs']}
  输出目录:    {save_dir}
  最佳权重:    {best_path}


  下一步建议:
  1. 完整训练: epochs=100, data="coco.yaml"
  2. 换大模型: "yolo26s.pt" / "yolo26m.pt"
  3. 自定义数据: 准备 YOLO 格式数据集 + 自定义 yaml
  4. 导出部署: 取消 Step 6 注释,导出 onnx/coreml/torchscript
        """)

    except FileNotFoundError as e:
        print(f"\n❌ 找不到文件: {e}")
        print("   确认 ultralytics 已安装: pip3 install ultralytics")
        sys.exit(1)
    except RuntimeError as e:
        err = str(e).lower()
        if "mps" in err or "m1" in err or "gpu" in err:
            print(f"\n⚠️  GPU/MPS 不可用,修改配置: CONFIG['device'] = 'cpu'")
        raise
    except KeyboardInterrupt:
        print("\n\n⛔ 用户主动中断")


if __name__ == "__main__":
    main()

京ICP备2024093538号-1