YOLOv8猫狗检测:从SwanLab可视化训练到Gradio Demo网站

news2024/10/6 20:35:48

基于YOLO模型在自定义数据上做训练，实现对特定目标的识别和检测，是CV领域非常经典的任务，也是AI项目落地最热门的方向之一。

这篇文章我将带大家使用Ultralytics、SwanLab、Gradio这两个开源工具，完成从数据集准备、代码编写、可视化训练到推理Demo的全过程。

观察了一下，中文互联网上似乎很少有针对自定义数据的，能直接跑起来的YOLO训练代码和教程，所以也希望这篇文章可以帮到在科研一线的大家～

在这里插入图片描述

代码：完整代码直接看本文第3、4节
模型与数据集：百度云，提取码: f238
实验过程：猫狗检测-SwanLab
SwanLab：https://swanlab.cn
Ultralytics：https://github.com/huggingface/transformers
Gradio：https://github.com/gradio-app/gradio

1.环境安装

我们需要安装以下这4个Python库：

ultralytics
swanlab>=0.3.6
gradio

一键安装命令：

pip install ultralytics swanlab gradio

他们的作用分别是：

Ultralytics：YOLO官方团队推出的CV训练与推理框架，不仅支持目标检测任务，还支持分割、姿态识别、分类等更多任务。本项目用Ultralytics作为训练框架。

在这里插入图片描述

swanlab: 一个深度学习实验管理与训练可视化工具，由西安电子科技大学团队打造，官网, 融合了Weights & Biases与Tensorboard的特点，可以记录整个实验的超参数、指标、训练环境、Python版本等，并可视化图表，帮助你分析训练的表现。本项目用swanlab主要用于记录指标和可视化。

在这里插入图片描述

gradio: HuggingFace出品的推理Demo构建工具，是深度学习最流行的Demo框架之一，可以用python代码轻松搭建网页。本项目用Gradio作为推理Demo框架。

在这里插入图片描述

整个项目的目录结构如下：

在这里插入图片描述

2.准备数据集

数据集这里我使用的是Kaggle上的Dog and Cat Detection数据集，包含3686张带标注的猫狗图像。

在这里插入图片描述

这里除了下载数据集以外，还需要对格式做处理，所以我把做好的数据集放到百度云（提取码: f238）里了，推荐大家直接下载。

下面重点介绍一下怎么让你的自定义数据集适配Ultralytics，掌握之后，几乎所有自定义数据集的处理就都学会了。

首先，Ultralytics推荐的数据集结构是这样的：

datasets
├── images
│   ├── train
│      ├── 00001.jpg
│      ├── ...
│   ├── val
│   ├── test
├── labels
│   ├── train
│      ├── 00001.txt
│      ├── ...
│   ├── val
│   ├── test
├── data.yaml

这里面是一个数据集文件夹，包含images、labels两个文件夹和一个data.yaml配置文件：

images文件夹放图像，labels文件夹放标注文件，图像和标注文件的名称要一一对应
images和labels文件夹下分别放train、val、test三个子文件夹，作为训练集、验证集和测试集
data.yaml的格式如下：

path:  path/to/datasets # 这里填写你数据集所在的绝对路径
train: images/train
val: images/train
test: images/test

# 标签和对应的类别
names:
  0: cat
  1: dog

标注文件的格式如下：

0 0.618 0.127 0.299 0.226
1 0.491 0.333 0.506 0.545

每一行的第一个数字代表标签，后续的四个数字是标注框的<x1> <y1> <x2> <y2>相对于图像shape的归一化值（或者说比例）。
标签文件中出现多行则代表着图像中有多个检测到的目标。

由于Kaggle的数据集并不是按YOLO的标注格式来的，所以需要写脚本进行处理，推荐直接从百度云（提取码: f238）下载。

这里我们只做train和val。最后处理好的数据集格式如下：

在这里插入图片描述

3.开始训练-train.py

在准备好数据集之后，最艰难的步骤就结束了——训练代码非常的简短，如下所示：

from ultralytics import YOLO
from swanlab.integration.ultralytics import add_swanlab_callback
import swanlab

def main():
    swanlab.init(project="Cats_Dogs_Detection", experiment_name="YOLOv8n",)
    model = YOLO("yolov8n.pt")
    add_swanlab_callback(model)
    # 将下面的路径替换成你的绝对路径
    model.train(data="path/to/cats_dogs_dataset/data.yaml", epochs=5, imgsz=320, batch=32)

if __name__ == "__main__":
    main()

请将model.train中data参数的路径替换成你的data.yaml的绝对路径。

这里我们使用了yolov8n模型（6MB左右）训练5个epoch，batchsize为32，并使用swanlab进行训练可视化。

在运行训练脚本的时候，如果你是第一次使用swanlab，那么需要去swanlab官网注册一个账号，然后在用户设置界面复制API Key，然后在命令行输入swanlab login，粘贴API Key即可完成登录。

训练过程（可访问猫狗检测-SwanLab查看）：

在这里插入图片描述

我这里做了个两个实验，分别使用yolov8n和yolov8s两个模型训练100个epoch，可以看到最终的结果，指标非常的不错：

在这里插入图片描述

训练好的模型会存放在Ultralytics自动生成的runs文件夹下。

4.推理代码-predict.py

from ultralytics import YOLO

# 载入训练好的模型
model = YOLO("path/to/model.pt")
# 推理多张图像
results = model(["img1.jpg", "img2,jpg"], device="cpu")

# 保存结果图
for iter, result in enumerate(results):
    result.save(filename=f"result{iter}.jpg")

5.Gradio推理Demo-app.py

import gradio as gr
from ultralytics import YOLO
from PIL import Image

# 加载预训练的 YOLO 模型
model = YOLO('path/to/best.pt')

def predict_image(image, conf_threshold, iou_threshold):
    # 使用模型进行推理
    results = model.predict(
        source=image, 
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,)
    
    # 提取结果
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
    
    return im

# 定义 Gradio 接口
demo = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
    ],
    outputs=gr.Image(type="pil", label="Result"),
    title="猫狗检测Demo",
    description="传一张带有猫狗的图像来进行推理。",
)

# 启动 Gradio 应用
if __name__ == "__main__":
    demo.launch()