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

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

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

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

1.环境安装

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

txt 复制代码
ultralytics
swanlab>=0.3.6
gradio

一键安装命令:

pip install ultralytics swanlab gradio

他们的作用分别是:

  1. Ultralytics:YOLO官方团队推出的CV训练与推理框架,不仅支持目标检测任务,还支持分割、姿态识别、分类等更多任务。本项目用Ultralytics作为训练框架。
  1. swanlab : 一个深度学习实验管理与训练可视化工具,由西安电子科技大学团队打造,官网, 融合了Weights & Biases与Tensorboard的特点,可以记录整个实验的超参数、指标、训练环境、Python版本等,并可视化图表,帮助你分析训练的表现。本项目用swanlab主要用于记录指标和可视化。
  1. gradio: HuggingFace出品的推理Demo构建工具,是深度学习最流行的Demo框架之一,可以用python代码轻松搭建网页。本项目用Gradio作为推理Demo框架。

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

2.准备数据集

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

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

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

首先,Ultralytics推荐的数据集结构是这样的:

txt 复制代码
datasets
├── images
│   ├── train
│      ├── 00001.jpg
│      ├── ...
│   ├── val
│   ├── test
├── labels
│   ├── train
│      ├── 00001.txt
│      ├── ...
│   ├── val
│   ├── test
├── data.yaml

这里面是一个数据集文件夹,包含imageslabels两个文件夹和一个data.yaml配置文件:

  • images文件夹放图像,labels文件夹放标注文件,图像和标注文件的名称要一一对应
  • imageslabels文件夹下分别放train、val、test三个子文件夹,作为训练集、验证集和测试集
  • data.yaml的格式如下:
yaml 复制代码
path:  path/to/datasets # 这里填写你数据集所在的绝对路径
train: images/train
val: images/train
test: images/test

# 标签和对应的类别
names:
  0: cat
  1: dog
  • 标注文件的格式如下:
txt 复制代码
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

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

python 复制代码
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

python 复制代码
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

python 复制代码
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()

运行程序后,会出现以下输出:

点开链接,出现猫狗检测的Demo网页:

用猫和狗的图片试试:

效果很完美!

至此,我们完成了用Ultralytics、SwanLab、Gradio三个开源工具训练1个猫狗检测模型的全部过程,更多想了解的可以参考相关链接或评论此文章。

如果有帮助,请点个赞和收藏吧~

5. 相关链接

相关推荐
Robot2511 分钟前
浅谈,华为切入具身智能赛道
人工智能
只怕自己不够好6 分钟前
OpenCV 图像运算全解析:加法、位运算(与、异或)在图像处理中的奇妙应用
图像处理·人工智能·opencv
果冻人工智能1 小时前
2025 年将颠覆商业的 8 大 AI 应用场景
人工智能·ai员工
代码不行的搬运工1 小时前
神经网络12-Time-Series Transformer (TST)模型
人工智能·神经网络·transformer
石小石Orz1 小时前
Three.js + AI:AI 算法生成 3D 萤火虫飞舞效果~
javascript·人工智能·算法
孤独且没人爱的纸鹤1 小时前
【深度学习】:从人工神经网络的基础原理到循环神经网络的先进技术,跨越智能算法的关键发展阶段及其未来趋势,探索技术进步与应用挑战
人工智能·python·深度学习·机器学习·ai
阿_旭1 小时前
TensorFlow构建CNN卷积神经网络模型的基本步骤:数据处理、模型构建、模型训练
人工智能·深度学习·cnn·tensorflow
羊小猪~~1 小时前
tensorflow案例7--数据增强与测试集, 训练集, 验证集的构建
人工智能·python·深度学习·机器学习·cnn·tensorflow·neo4j
极客代码2 小时前
【Python TensorFlow】进阶指南(续篇三)
开发语言·人工智能·python·深度学习·tensorflow
zhangfeng11332 小时前
pytorch 的交叉熵函数,多分类,二分类
人工智能·pytorch·分类