从0到1，快速训练并使用YOLO模型
简介
YOLO是目前计算机视觉领域最前沿、应用最广泛的目标检测算法框架，他能快速识别区分目标，广泛应用于游戏，无人驾驶，工业等领域。
以识别躲避掉落滑块的游戏的物体图片作为例子。
一，环境配置
pip install ultralytics二，准备数据集
这个格式目录如下：
my_dataset/ ├── data.yaml # 配置文件（定义路径和类别） ├── train/ #训练数据集 │ ├── images/ # 训练图片 │ └── labels/ # 标注文件 (.txt) └── val/ #验证数据集 ├── images/ └── labels/data.yml
path: D:\D_MyProject\Ai\game_ai\my_dataset #数据集路径 train: train/images #训练集图片路径 val: val/images #验证集图片路径 nc: 3 #标记个数 names: #每个标记的名称 0: player 1: enemy 2: game_over下面是用AI生成了数据集的生成脚本
数据量太多了，这里为了演示，或者学习，可以直接使用下面脚本
import pygame import random import sys import os import shutil # =================配置区域================= # 数据集根目录名称 DATASET_ROOT = "my_dataset" # 采集总数量 MAX_IMAGES = 1000 # 训练集占比 (0.8 = 80% 训练, 20% 验证) TRAIN_RATIO = 0.8 # ================= 1. 环境清理与目录创建 ================= print(f"🚀 正在初始化数据集目录: {DATASET_ROOT} ...") # 如果目录已存在，先删除（防止旧数据混入），确保数据纯净 if os.path.exists(DATASET_ROOT): shutil.rmtree(DATASET_ROOT) # 创建 YOLO 标准目录结构 for split in ['train', 'val']: os.makedirs(os.path.join(DATASET_ROOT, split, 'images'), exist_ok=True) os.makedirs(os.path.join(DATASET_ROOT, split, 'labels'), exist_ok=True) # ================= 2. 自动生成 data.yaml ================= yaml_content = f""" path: {os.path.abspath(DATASET_ROOT)} # 使用绝对路径，防止报错 train: train/images val: val/images nc: 3 names: 0: player 1: enemy 2: game_over """ with open(os.path.join(DATASET_ROOT, "data.yaml"), "w", encoding='utf-8') as f: f.write(yaml_content) print("✅ data.yaml 配置文件已生成。") # ================= 3. 游戏与采集逻辑 ================= pygame.init() WIDTH, HEIGHT = 800, 600 # 使用 hidden 模式或正常模式均可，这里用正常模式方便你看到进度 screen = pygame.display.set_mode((WIDTH, HEIGHT)) pygame.display.set_caption("Auto Data Generator") clock = pygame.time.Clock() font_big = pygame.font.SysFont("monospace", 50) font_small = pygame.font.SysFont("monospace", 35) # 游戏状态 player_size, enemy_size = 50, 50 player_x = WIDTH // 2 player_y = HEIGHT - player_size - 10 enemies = [] img_count = 0 GAMEPLAY_LIMIT = int(MAX_IMAGES * 0.9) # 90% 正常游戏，10% Game Over print(f"📸 开始采集 {MAX_IMAGES} 张图片 (自动划分 Train/Val)...") while img_count < MAX_IMAGES: # 处理退出事件，防止卡死 for event in pygame.event.get(): if event.type == pygame.QUIT: pygame.quit() sys.exit() screen.fill((0, 0, 0)) labels = [] # 存储当前帧的标签 dw, dh = 1.0 / WIDTH, 1.0 / HEIGHT # 归一化系数 # --- 逻辑分支：正常游戏 vs Game Over --- if img_count < GAMEPLAY_LIMIT: # A. 正常游戏画面 # 1. 玩家移动 if random.random() < 0.2: # 增加移动频率 player_x += random.choice([-15, 15]) player_x = max(0, min(WIDTH - player_size, player_x)) # 2. 敌人生成与移动 if random.randint(0, 20) == 0: # 增加敌人密度 enemies.append([random.randint(0, WIDTH - enemy_size), 0]) for enemy in enemies: enemy[1] += 15 # 加快下落速度 enemies = [e for e in enemies if e[1] < HEIGHT] # 3. 绘制玩家 (Class 0) pygame.draw.rect(screen, (50, 150, 255), (player_x, player_y, player_size, player_size)) # 计算 YOLO 坐标 (class x_center y_center width height) px, py = (player_x + player_size/2) * dw, (player_y + player_size/2) * dh labels.append(f"0 {px:.6f} {py:.6f} {player_size*dw:.6f} {player_size*dh:.6f}") # 4. 绘制敌人 (Class 1) for e in enemies: pygame.draw.rect(screen, (255, 50, 50), (e[0], e[1], enemy_size, enemy_size)) ex, ey = (e[0] + enemy_size/2) * dw, (e[1] + enemy_size/2) * dh labels.append(f"1 {ex:.6f} {ey:.6f} {enemy_size*dw:.6f} {enemy_size*dh:.6f}") else: # B. Game Over 画面 (Class 2) text_surf = font_big.render("GAME OVER", True, (255, 50, 50)) # 随机抖动位置，防止过拟合 off_x, off_y = random.randint(-50, 50), random.randint(-50, 50) text_x = WIDTH // 2 - text_surf.get_width() // 2 + off_x text_y = HEIGHT // 2 - 100 + off_y screen.blit(text_surf, (text_x, text_y)) # 干扰项 score_surf = font_small.render(f"Score: {random.randint(0,999)}", True, (255,255,255)) screen.blit(score_surf, (WIDTH//2 - score_surf.get_width()//2, HEIGHT//2 + 50)) # 记录标签 tw, th = text_surf.get_width(), text_surf.get_height() tx, ty = (text_x + tw/2) * dw, (text_y + th/2) * dh labels.append(f"2 {tx:.6f} {ty:.6f} {tw*dw:.6f} {th*dh:.6f}") # ================= 4. 保存逻辑 (核心修改) ================= pygame.display.flip() # 采样率：不是每一帧都保存，防止重复度过高 (这里设为 30% 概率保存) if random.random() < 0.3: # A. 决定是去 Train 还是 Val split_folder = "train" if random.random() < TRAIN_RATIO else "val" # B. 生成文件名 filename = f"{img_count:06d}" img_save_path = os.path.join(DATASET_ROOT, split_folder, "images", f"{filename}.jpg") lbl_save_path = os.path.join(DATASET_ROOT, split_folder, "labels", f"{filename}.txt") # C. 保存图片 pygame.image.save(screen, img_save_path) # D. 保存标签 with open(lbl_save_path, "w") as f: f.write("\n".join(labels)) img_count += 1 # 打印进度条 print(f"[{split_folder.upper()}] 进度: {img_count}/{MAX_IMAGES}", end="\r") # 加速模拟，不要垂直同步，越快越好 clock.tick(0) pygame.quit() print(f"\n\n✨ 全部完成！数据集已就绪：{os.path.abspath(DATASET_ROOT)}") print("💡 下一步：直接运行 model.train(data='my_dataset/data.yaml')")三，训练YOLO模型
可以看到，使用ultralytics框架训练YOLO的代码非常简单，只需要几行
注意：这里YOLO会自动下载模型并训练，下载时失败可能需要挂梯子
from ultralytics import YOLO def train_model(): #加载模型 model = YOLO("yolo11n.pt") #开始训练 print("开始训练...") results = model.train( data="my_dataset/data.yaml", #数据集配置文件 epochs=30, #训练轮数 imgsz=640, #图片输入尺寸 batch=16, #显存够大可以改大，比如 32 或 64 device=0, #强制使用第一张显卡 (需要CUDA) workers=0, #Windows下设为0防止多进程报错 project="dodge_project",#保存路径 name="ai_model" #训练运行名称 ) print(f"训练完成！最佳模型保存在: {results.save_dir}/weights/best.pt") if __name__ == "__main__": train_model()四，使用YOLO模型
from ultralytics import YOLO #配置路径 MODEL_PATH = "dodge_project/ai_model/weights/best.pt" PIC_PATH = "my_dataset/val/images/000045.jpg" #加载模型 model = YOLO(MODEL_PATH) print(f"模型加载成功") #识别图片 results = model.predict(PIC_PATH, verbose=False, conf=0.4, imgsz=640) result = results[0] #获取信息并打印 for box in result.boxes: # 获取坐标 (x1, y1, x2, y2)、类别 ID 和置信度 x1, y1, x2, y2 = map(int, box.xyxy[0]) cls_id = int(box.cls[0]) conf = float(box.conf[0]) # 类别名称映射 names = {0: "Player", 1: "Enemy", 2: "Game Over"} label = names.get(cls_id, f"Unknown({cls_id})") print(f"找到物体: [{label:10}] | 置信度: {conf:.2f} | 坐标: ({x1}, {y1}) -> ({x2}, {y2})") #展示图片 result.show()