1 项目需求描述
博主自己开了一家T宝店,有一个客户有这个需求:有一大堆图像,大概有10多万张图,都是比较小尺寸的图。各种类型都有,我们想要通过将不同类型发图像进行分开,如何实现呢?
2 思路
可以通过训练模型,将图像通过模型预测,来进行筛选图像。比如:我们先训练第一个图像分类模型:卡通/猫。将每一个图像都使用模型预测下,如果预测为卡通或者猫的得分比较大,那么,我们就认为这个图像就是卡通或者猫,就把它从原始文件夹中拿出来,放到预测结果的这个类别的文件夹中。
3 步骤
3.1 图像准备
先从这么多图像中,筛选出卡通图,猫图各自100多张。用于训练。整理数据格式为:
3.2 环境准备
要进行模型训练,我们就需要在windows上配置训练的环境。这里,我们仍使用yolo11进行训练。环境准备过程如下:
3.2.1 环境创建
之前,我们都是使用conda创建环境,这次我们换一种方案,uv。
首先,在桌面上新建一个文件夹TB_Pro-Clas-Cat-Cartoon,然后打开Pycharm选择项目路径。配置环境如下(如果不会配置,可以参考Pycharm配置uv环境):
然后点击创建,进入如下界面:
3.2.2 安装ultralytics
使用如下命令安装ultralytics:
uv pip install ultralytics
3.2.3 克隆仓库
使用如下命令,注意,最好提前打开科学上网,这样下载速度会更快!
git clone https://github.com/ultralytics/ultralytics.git
现在,我们查看一下目录,就会看到如下结果:
3.3 划分数据集
然后,为了训练模型,我们需要将数据集整理为下面的格式:
包含训练集、测试集和验证集,按类别分文件夹存储。
📁 data-train/
├── 📁 train/
│ ├── 📁 catroon/
│ └── 📁 cat/
├── 📁 test/
│ ├── 📁 catroon/
│ └── 📁 cat/
└── 📁 val/ (可选,本次项目,我们不划分)
3.4 开始训练
3.4.1 下载模型
从官网直接下载到的模型:
模型 | 尺寸 (像素) | acc top1 | acc top5 | 速度 CPU ONNX (毫秒) | 速度 T4TensorRT10 (ms) | params (M) | FLOPs (B) at 224 |
---|---|---|---|---|---|---|---|
YOLO11n-cls | 224 | 70.0 | 89.4 | 5.0 ± 0.3 | 1.1 ± 0.0 | 1.6 | 0.5 |
YOLO11s-cls | 224 | 75.4 | 92.7 | 7.9 ± 0.2 | 1.3 ± 0.0 | 5.5 | 1.6 |
YOLO11m-cls | 224 | 77.3 | 93.9 | 17.2 ± 0.4 | 2.0 ± 0.0 | 10.4 | 5.0 |
YOLO11l-cls | 224 | 78.3 | 94.3 | 23.2 ± 0.3 | 2.8 ± 0.0 | 12.9 | 6.2 |
YOLO11x-cls | 224 | 79.5 | 94.9 | 41.4 ± 0.9 | 3.8 ± 0.0 | 28.4 | 13.7 |
为了得到更高的准确率,我们使用YOLO11x-cls。下载后就会得到:
我们可以将这个文件挪到ultralytics文件夹内!
3.4.2 编写训练代码
新建一个train-code.py文件。
并编写训练代码如下:
from ultralytics import YOLO# Load a model
## yolo11x-cls.pt这个模型是从官网直接下载到的模型
model = YOLO("yolo11x-cls.pt") # Train the model
results = model.train(data="data-train", epochs=100, imgsz=224)
然后,在pycharm终端,cd到ultralytics文件夹,运行py文件,就开始训练模型啦!
经过漫长的等待(等待时间取决于你的电脑显卡配置),终于训练完成啦!
然后,我们在路径下,就会看到模型文件:
至此,模型训练就结束啦!