利用Yolo模型实现图像分类的迅速上手Blog！

置顶

精

活动

最近要做一个有关于图像分类的项目，大体需要区分下面三种图像，分别是模糊的PCB孔，外孔清晰照片，内孔清晰照片。

为了实现这一目的，于是我想到了利用Yolo（You only Look Once）模型来实现图像实时识别。

    本期我们就介绍如何快速上手使用Yolo训练自己的实时分类模型。

YOLO模型对硬件的要求相对灵活，最低配置下，CPU需要4核，内存至少16GB。

    除此之外，推荐安装Python 3.6及以上的版本的来使用。

下载和安装Python可以前往Python官网下载与安装，本文使用的Python环境是Python 3.9

编译器使用的PyCharm 2023,也可自行搜索PyCharm进行安装。

除了环境要求之外，准备好我们需要检测的图片，接下来我们将对图片进行标签分类。

标注工具我们使用开源的 LabelImg，在CMD窗口输入以下代码，即可安装Labelimg（这里也可以用虚拟环境）

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

完成后，即可在Python的根目录中找到LabelImag的启动程序了。

接下来在图片中，选中待识别图片的路径，打开图片，挨个对其进行标注了。

保存完后我们就可以得到图片标注信息了，包括“类别索引”，标注框信息的归一化数据等。

 创建一个训练文件夹和预测文件夹（注意，标签也要创建），分别放入图片和标注结果（txt文件）,这里我分别放置了70张图片用作训练，10张图片用作预测。

 接下来利用Pycharm创建一个新的工程，在终端中可以输入如下代码安装Yolo。

pip install ultralytics

之后创建一个Yolo的配置文件，描述一下我们要训练数据情况。

py文件中的代码非常非常的简单，只需要下面一段即可。

import torch
from ultralytics import YOLO
# 加载模型
model = YOLO('yolov8n.yaml')  # 使用 YoloV8 的预训练模型

# 训练模型
results = model.train(
    data='./dataset.yaml',  # 数据集配置文件路径
    epochs=100,  # 训练轮数
    imgsz=800,#图片大小 800*800像素
    batch=16,#一次训练的图片批次
    name='yolov8_custom',
    save_dir="ModeSave",    #模型保存路径
    save_period=10  #每十轮保存一次
)

# 评估模型
results = model.val()

    接着开始运行，下面是训练过程部分截图

 一百次训练大概花了一个多小时，十四张用于预测的图片全部识别成功。

我这里的标签是：blur是模糊的照片，clear1是外孔清晰照片，clear2是内孔清晰照片，数字是置信度，除了第五张图片的置信度为0.3之外，其他的都非常准确。