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 文档