训练自己的yolov5数据集并部署到android全套流程(学习记录)

一、先决条件

Python 3.8
Pytorch 1.10.0
CUDA         11.3
Tensorflow 2.13.0
Torchaudio 0.10.0
Torchvision 0.11.1
Android Studio Gradle Verrsion 7.5
Android Gradle Plugin Version         7.4.1
Tensorflow-lite                

2.8.0

首先安装CUDA和cudnn,参考:CUDA安装教程(超详细)

然后安装pytorch。打开anaconda promt终端,创建虚拟环境:

conda create --name envname python=3.8

这里envname要替换成你自己的虚拟环境的名字,python版本也可以选择其他对应的版本。创建好之后激活环境:
 

conda activate envname

安装pytorch:

打开https://pytorch.org/

点击Previous versions of PyTorch,打开过往的版本,找到所需要下载的torch版本号,例如我这里是torch=1.10.0,cuda=11.3,所以选择如图,复制指令到终端。也可以根据自己所需的版本来选择

这里torchversion=0.11.0,但其实我所需要的版本为0.10.1,因此需要做一些修改:

conda install pytorch==1.10.0 torchvision==0.11.1 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge

然后复制到终端运行(前提是在激活的虚拟环境中)

尽量不要使用其他源下载,可能会出现HTTP 000 CONNECTION FAILED之类的错误,使用科学上网可能会好一点,耐心等待下载完成 

安装Android studio可以参考:Android Studio2022安装教程

打开AS,会有相关依赖的下载,耐心等待下载完成

二、数据集

1.打标签

在虚拟环境中输入

pip install labelImg

安装labelimg。

然后在命令行输入labelimg来运行软件

labelimg

 打开要打标签的图片,右键点击图片,点击创建rectbox,在图中对应目标画出框覆盖它,然后在弹出的标签名称中写上对应的名称,然后保存,记得选择YOLO格式,对应的是txt文件。

当然也可以打开图像文件夹进行操作会更快介意的,这里就不展示了

标注完成之后会得到一系列的txt文件,txt文件和图片名称是一一对应的

2.数据集整理

在当前项目中新建一个文件夹,例如dataset

在文件下新建两个文件夹分别为images和labels

在images中新建三个文件夹train,val,test,在labels中同样也新建三个文件夹train,val,test

其中train用来存放训练集,val用来存放验证集,test用来存放预测集。images中是图像,而labels中是标签。

将之前打了标签的图像和对应导出的txt标签进行数据集划分,数据集划分参考:对yolov5的数据集进行划分【训练集、验证集、测试集】7:2:1和【训练集、验证集】8:2
然后将训练集,验证集,测试集的图像和标签分别保存到刚刚创建的文件夹中,文件结构如下:


————dataset
       ├─ images
       │    ├─ test # 测试集图片
       │    ├─ train # 训练集图片
       │    └─ val # 验证集图片
       └─ labels
              ├─ test # 测试集标签
              ├─ train # 训练集标签
              ├─ val # 验证集标签

三、训练自己的yolov5模型

1.官方下载yolov5-master

YOLOV5源码下载

下载好之后pycharm打开

2.添加数据集配置文件

在data目录下创建一个mydata.yaml文件,用于配置我们的训练数据

mydata.yaml 代码:

#数据路径
train: C:/Users/11645/Desktop/dataset/images/train
val: C:/Users/11645/Desktop/dataset/images/val

# 类别数
nc: 3

# 类别名称
names: ["bus", "person","stone"]

根据自己的数据集路径和类别数进行调整。

3.修改模型配置文件

在model文件夹下找到自己所需要模型的配置文件,我这里用yolov5s,打开修改类别数即可。

注意:也可以复制一份yolov5s.yaml代码到该目录下,修改名称为mymodel.yaml,修改类别(为了不破坏原始代码)

4.训练

在pycharm终端运行

python train.py --data mydata.yaml --cfg yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device 0

如果上一步 复制一份yolov5s.yaml代码到该目录下修改名称为mymodel.yaml,那么这里就要把代码改为:

python train.py --data mydata.yaml --cfg mymodel.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device 0

注意:我这里是使用GPU训练的,device后面的0代表GPU,如果要用cpu请把0改为cpu

训练过程中报错no moudle named ***,就在当前环境下pip install ***安装即可

其中No module named ‘cv2‘是指opencv,执行

pip install opencv-python

 5.预测

训练完成后,权重文件会保存在runs/train/exp/weights中,有一个best.pt和last.pt,我们用best.pt。

如果要预测某一张图片,可以运行:

python detect.py --weights runs/train/exp/weights/best.pt --source data/images/test.jpg

路径替换成自己的

6.模型转换

这里就要将我们的权重文件转换到android端

之前利用官方demoObject Detection with YOLOv5 on Android,在这个基础上修改,将best.pt转换为torchscript格式,即best.torchscript.ptl,替换了官方提供的yolov5s.pt并修改了classes中的类别,但是遇到了各种奇怪的错误,demoAPP会闪退,找了很多办法也没有解决,于是换了种思路,将best.pt转换为best-fp16.tflite格式,成功运行

在pycharm中运行

python export.py --weights runs/train/exp/weights/best.pt --include tflite

然后在weights文件夹下会生成一个.tflite格式文件,这个就是我们所需要的文件了,然后检测是否可用,并查看输入和输出大小,执行下列代码:
 

python
import tensorflow as tf
model = "runs/train/exp/weights/best-fp16.tflite" 
interpreter = tf.lite.Interpreter(model_path = model) 
print(interpreter.get_input_details())  #打印输入大小
print(interpreter.get_output_details())  #打印输出大小

然后运行结果如图 ,输入大小和输出大小后面会用到

四、部署到Android Studio

参考demo:https://github.com/yyccR/yolov5-tflite-android

将自己上一步得到的best-fp16.tflite放在assets文件夹下面,并新增一个分类文件coco_labels2.txt

分类文件.txt中的类别数要和训练模型时保持一致

打开Yolov5TFLiteDetector,查找将其中的模型替换为自己的best-fp16.tflite

查找

将coco_label.txt替换为自己的coco_labels2.txt

然后查看3.6中输入和输出大小。

private final Size INPUT_SIZE=new Size(320,230)

private final int[] OUTPUT_SIZE=new int[]{1,6300,85}

更改为自己输入和输出的大小即可

private final Size INPUT_SIZE=new Size(640,640)
private final int[] OUTPUT_SIZE=new int[]{1,25200,8}

参考博客:
Android嵌入自己训练的yolov5模型(tfLite)交通标志

手把手教你使用YOLOV5训练自己的目标检测模型-口罩检测-视频教程

猜你喜欢

转载自blog.csdn.net/qwerdf6660/article/details/134983459