Appearance
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()