yolo训练常用代码

浩峰

已于 2025-07-09 09:25:30 修改

阅读量420

点赞数 4

CC 4.0 BY-SA版权

文章标签： YOLO 深度学习人工智能

于 2025-04-16 23:14:31 首次发布

本文链接：https://blog.youkuaiyun.com/pgp0000/article/details/147288622

1、data.yaml配置文件

data.yaml是YOLO系列目标检测模型中用于配置数据集信息的核心文件基本结构为，train文件夹下的images路径，val文件夹下的images路径，nc类别数量，calsses类别名称。

train: /mnt/workspace/datasets/train/images
val: /mnt/workspace/datasets/val/images
nc: 8
names: ['test00', 'test01', 'test02', 'test03', 
        'test04', 'test05', 'test06', 'test07']

2、train训练文件

train.py文件为模型训练文件，主要逻辑为导入YOLO库，加载预训练权重文件，利用YOLO库的train方法以及data.yaml配置文件，及相关训练参数对我们的数据集进行训练。

from ultralytics import YOLO

if __name__ == "__main__":
    model = YOLO("yolo11n.pt")
    model.train(data="data.yaml", epochs=200, imgsz=640, batch=16, workers=4)

    #  如果batch=-1表示根据设备性能自调整批次大小

3、jieya.py文件

Jieya.py主要作用是将上传到云端的zip文件进行解压，核心为导入zipfile库。

import zipfile

if __name__ == '__main__':
    zip_file_path = 'D:/test.zip'
    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
        zip_ref.extractall('.') # .表示解压到当前目录

!unzip test.zip # 在Notebook环境下可以直接解压

4、export.py文件导出onnx格式文件

此文件主要是将训练好的pt文件导出为想要的格式文件，如onnx文件，openvino文件等。

from ultralytics import YOLO # 加载YOLO模型

# 需要安装onnxslim库，安装方式为pip install onnxslim

if __name__ == '__main__':

    model = YOLO('best.pt') # 加载模型

    model.export(format='onnx') # 导出onnx格式的模型

5、export.py文件导出openvino格式文件

from ultralytics import YOLO

# 导入openvino-dev

if __name__ == '__main__':
    # Load a pretrained model
    model = YOLO('best.pt')

    # Export the model to ONNX format
    model.export(format='openvino')

    # Export the model to TorchScript format

6、利用bt模型对图片进行推理

from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO (r'E:\YOLO\test\best.pt', task='detect')
    results = model.predict(
        source=r'E:\YOLO\test\datasets\train\images',
        save=True, # 检测后的图片会保存在runs文件夹下
        stream=True
    )
for r in results:
    boxes = r.boxes

7、利用openvino模型对视频进行推理

如果需要其他功能，可将此段代码丢给AI，给出想要的功能指令即可

from ultralytics import YOLO

# 在 Ultralytics YOLOv8 中，加载 OpenVINO 模型时应使用 .onnx、.pt 或 .xml 的父目录，而不是直接指定 .xml 文件路径
model = YOLO(r'E:\YOLO\test\datasets\best_openvino_model', task='detect')  # 替换为你的模型路径
results = model.predict(source=r'E:\YOLO\video\test.MP4', show=True, stream=True)

for r in results:
    boxes = r.boxes

8、model.yaml模型标注配置文件

此文件主要用作小模型自动标注大模型时的配置文件

type: yolo11 # 指定模型架构或检测框架类型，名字不能随意起，只能是yolov5/yolov7/yolov8/yolo_nas/yolo11
name: yolo11 # 自定义模型配置名称，随便起名
display_name: yolo2025 # 自定义模型名称，随便起名
model_path: E:\YOLO\sanguo\best.onnx
nms_threshold: 0.45 # 非极大值抑制
confidence_threshold: 0.25 # 置信度
# iou_threshold: 0.5 # 交并比阈值
classes:
  - test00
  - test01
  - test02
  - test03
  - test04
  - test05
  - test06
  - test07

9、获取图片到指定文件夹

功能就是在指定的源图片文件夹路径下，指定图片数量后，会根据总图片数量均匀复制图片到指定的文件夹下,如果目标文件夹下有相同内容，则进行覆盖替换。

import os
import shutil
import logging
from pathlib import Path
from typing import List, Tuple
from PIL import Image

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 配置参数 - 请在此处修改源文件夹和目标文件夹路径
SOURCE_FOLDER = r"E:\YOLO\video\test\test"  # 源图片文件夹路径
TARGET_FOLDER = r"C:\Users\test\Desktop\test_img"  # 目标文件夹路径
IMAGE_COUNT = 10  # 要复制的图片数量
OVERWRITE_EXISTING = True  # 是否覆盖已存在的图片

def is_image_file(file_path: str) -> bool:
    """检查文件是否为图片文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'}
    return Path(file_path).suffix.lower() in image_extensions

def get_image_dimensions(image_path: str) -> Tuple[int, int]:
    """获取图片的尺寸"""
    try:
        with Image.open(image_path) as img:
            return img.size
    except Exception as e:
        logger.error(f"无法获取图片 {image_path} 的尺寸: {e}")
        return (0, 0)

def analyze_images(source_dir: str) -> List[Tuple[str, int, int]]:
    """分析源目录中的所有图片，返回包含路径、宽度和高度的列表"""
    image_list = []
    for root, _, files in os.walk(source_dir):
        for file in files:
            file_path = os.path.join(root, file)
            if is_image_file(file_path):
                width, height = get_image_dimensions(file_path)
                if width > 0 and height > 0:
                    image_list.append((file_path, width, height))
    return image_list

def select_images_uniformly(image_list: List[Tuple[str, int, int]], count: int) -> List[str]:
    """基于尺寸均匀选择图片"""
    if not image_list:
        return []
    
    if len(image_list) <= count:
        logger.warning(f"源文件夹中只有 {len(image_list)} 张图片，将全部复制")
        return [item[0] for item in image_list]
    
    # 按尺寸排序
    sorted_images = sorted(image_list, key=lambda x: (x[1], x[2]))
    
    # 计算采样间隔
    step = len(sorted_images) / count
    
    # 均匀选择图片
    selected_indices = [int(i * step) for i in range(count)]
    return [sorted_images[i][0] for i in selected_indices]

def copy_images(selected_images: List[str], target_dir: str) -> None:
    """复制选中的图片到目标文件夹，根据配置决定是否覆盖已存在的图片"""
    os.makedirs(target_dir, exist_ok=True)
    
    for image_path in selected_images:
        try:
            file_name = os.path.basename(image_path)
            target_path = os.path.join(target_dir, file_name)
            
            # 检查文件是否存在并决定是否覆盖
            if os.path.exists(target_path):
                if OVERWRITE_EXISTING:
                    logger.info(f"目标文件已存在，将覆盖: {target_path}")
                    shutil.copy2(image_path, target_path)
                    logger.info(f"已覆盖复制: {image_path} -> {target_path}")
                else:
                    logger.info(f"目标文件已存在，跳过复制: {target_path}")
            else:
                shutil.copy2(image_path, target_path)
                logger.info(f"已复制: {image_path} -> {target_path}")
                
        except Exception as e:
            logger.error(f"复制图片 {image_path} 失败: {e}")

def main():
    """主函数"""
    # 验证源文件夹是否存在
    if not os.path.exists(SOURCE_FOLDER):
        logger.error(f"源文件夹不存在: {SOURCE_FOLDER}")
        return
    
    # 分析图片
    logger.info(f"正在分析源文件夹 {SOURCE_FOLDER} 中的图片...")
    image_list = analyze_images(SOURCE_FOLDER)
    
    if not image_list:
        logger.error(f"在源文件夹 {SOURCE_FOLDER} 中未找到有效图片")
        return
    
    logger.info(f"找到 {len(image_list)} 张有效图片")
    
    # 均匀选择图片
    logger.info(f"正在均匀选择 {IMAGE_COUNT} 张图片...")
    selected_images = select_images_uniformly(image_list, IMAGE_COUNT)
    
    # 复制图片
    logger.info(f"开始复制图片到 {TARGET_FOLDER}...")
    copy_images(selected_images, TARGET_FOLDER)
    
    logger.info(f"操作完成！已处理 {len(selected_images)} 张图片")

if __name__ == "__main__":
    main()

10、推理视频并逐帧导出推理结果图片

from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(r'E:\YOLO\model_20250708\best.pt', task='detect')
    results = model.predict(
        source=r'E:\YOLO\video\7M2W\VID_20250702_095842.mp4', # 推理图片和视频根据需求启用下方的代码
        #show=True, # 推理结果显示
        save=True,
        stream=True,
        save_frames=True,  # 启用逐帧保存
        vid_stride=200, # 启用每隔5帧保存推理结果
        project=r'C:\Users\HF03\Desktop\test_img',  # 指定保存文件夹路径
        #name='prediction_frames'  # 子文件夹名称
    )
for r in results:
    boxes = r.boxes

11、视频切片

对不同视频格式进行代码分割。

import cv2
import os
from pathlib import Path


def extract_frames(video_path, output_dir, frame_interval=10):
    """从视频中每隔指定帧数提取一帧"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误：无法打开视频文件 {video_path}")
        return

    frame_count = 0
    output_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 每隔frame_interval帧保存一次
        if frame_count % frame_interval == 0:
            output_path = os.path.join(output_dir, f"frame_{output_count:06d}.jpg")
            cv2.imwrite(output_path, frame)
            output_count += 1

        frame_count += 1

    cap.release()
    print(f"处理完成：从{video_path}中提取了{output_count}帧，保存到{output_dir}")


if __name__ == "__main__":
    # 直接指定输入输出路径
    video_path = r"E:\YOLO\video\test\00000.MTS"  # 请替换为实际视频路径
    output_dir = r"E:\YOLO\video\test\00000"  # 请替换为实际输出目录

    # 执行帧提取
    extract_frames(video_path, output_dir)

、参考官方文档

详情请点击Ultralytics YOLO 文档