YOLO的部署和微调

YOLO（You Only Look Once）是一种流行的物体检测和图像分割模型，目前已经更新到YOLOV11。但是用法基本类似，所以下面以YOLOV8为例来介绍其使用方法。

YOLO简介

Github: https://github.com/ultralytics/ultralytics （目前已经更新到YOLOV11）

数据集结构

参考：《数据集制作-目标检测》 和 《数据集制作-LabelImg的使用》

dataset_root/                  # 数据集根目录
├── images/                    # 图像文件夹
│   ├── train/                 # 训练集图像
│   │   ├── img1.jpg
│   │   └── ...
│   ├── val/                   # 验证集图像
│   └── test/                  # 测试集图像（可选）
│
├── labels/                    # 标注文件夹
│   ├── train/                 # 训练集标注文件
│   │   ├── img1.txt
│   │   └── ...
│   ├── val/                   # 验证集标注文件
│   └── test/                  # 测试集标注文件（可选）
│
└── data.yaml                  # 数据集配置文件

训练（微调）

from ultralytics import YOLO

# model = YOLO("yolov8n.yaml").load("runs/detect/train36/weights/last.pt")  # 加载自己训练好的模型
model = YOLO("yolov8s.yaml").load("yolov8s.pt")  # build from YAML and transfer weights

if __name__ == "__main__":
    # Train the model
    results = model.train(data="F:\\datasets\\mydataset3\\custom.yaml", epochs=200, imgsz=640)

推理

from ultralytics import YOLO
import os

# 加载训练好的YOLO模型
model = YOLO("runs/detect/train44/weights/best.pt")

# 设置测试图像目录路径
image_dir = "dataset4/test/class1"
# 获取目录下所有.jpg图像文件路径
image_files = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith('.jpg')]

# 批量推理（前300张图像）
results = model(image_files[0:300])

# 处理推理结果
for i, result in enumerate(results, start=1):
    # 获取模型输出信息（可根据需要取消注释使用）
    # boxes = result.boxes    # 边界框信息
    # masks = result.masks    # 分割掩码信息
    # keypoints = result.keypoints  # 关键点信息
    # probs = result.probs    # 分类概率
    # obb = result.obb        # 旋转边界框
    
    # 保存结果图像到results目录
    result.save(filename=f"results/result-detect-{i}.jpg")

使用streamlit创建UI

import base64
import streamlit as st
import numpy as np
from PIL import Image #Image 是PIL 库中代表一个图像的类（对象）
import cv2
from ultralytics import YOLO
import os
 
st.markdown('<h1 style="color:black;">YoloV8 Image classification model</h1>', unsafe_allow_html=True)
 
 
# 背景图片background image to streamlit
 
@st.cache_data()
# 以base64的方式传输文件
def get_base64_of_bin_file(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()
#设置背景图片,颜色等
def set_png_as_page_bg(png_file):
    bin_str = get_base64_of_bin_file(png_file) 
    page_bg_img = '''
    <style>
    .stApp {
    background-image: url("data:image/png;base64,%s");
    background-size: cover;
    background-repeat: no-repeat;
    background-attachment: scroll; # doesn't work
    }
    </style>
    ''' % bin_str
    
    st.markdown(page_bg_img, unsafe_allow_html=True)
    return
 
set_png_as_page_bg('./dog.png')
 
# Load a model
model = YOLO("runs/detect/train44/weights/best.pt")
#预测
image_dir = "dataset4/test/class2"
# 上传png/jpg的照片
upload= st.file_uploader('Insert image for classification', type=['png','jpg'])
c1, c2= st.columns(2)
if upload is not None:
    results = model(os.path.join(image_dir, upload.name))
    results[0].save(filename="results/result-dectet-0.jpg")
    c1.image(os.path.join(image_dir, upload.name))
    c2.image('results/result-dectet-0.jpg')

附录-开发环境

# 虚拟环境
python -m venv venv
.\venv\Scripts\activate
# 清华源： -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install ultralytics
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121