"""
YOLOv10 目标检测 Demo（官方 ultralytics >=8.3.0）
首次运行会自动下载 yolov10n.pt 到 ~/.cache/ultralytics
"""
import cv2
from ultralytics import YOLO

def main():
    # 用统一的 YOLO 类加载 YOLOv10 预训练权重
    model = YOLO("yolov10n.pt")  # 自动下载；也可用 yolov10s/m/l/x.pt

    # 1) 对单张图片做推理
    img_path = "img.jpg"
    results = model(img_path, imgsz=640, conf=0.25)
    r = results[0]

    print("检测到", len(r.boxes), "个目标")
    print("类别:", [model.names[int(c)] for c in r.boxes.cls.tolist()])

    # 画框并保存
    annotated = r.plot()
    cv2.imwrite("result.jpg", annotated)
    print("结果已保存为 result.jpg")

    # 2)（可选）训练示例 —— 取消注释即可在 COCO8 小数据集上试跑
    # model.train(data="coco8.yaml", epochs=3, imgsz=640, device="cpu")

if __name__ == "__main__":
    main()