目录
上一篇 yolov8的完整部署 讲解了Yolov8模型的部署和运行
实际生活中,我们需要的检索目标,很多并不在官方给出的模型之中。那我们就需要训练自己的数据模型。
1. 训练前准备
1.1 收集训练数据,完成分类标记
首先,我们在根目录datasets/文件夹下像存放coco128数据的格式创建fire文件,在fire/文件下创建Myimages和Annotations文件夹,这个后面我们会用到。
这里我们要用到一个工具 labelimg,通过下面命令安装这个工具:
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple
安装好以后,我们在fire文件夹下创建一个predefined_classed.txt文件, 在这个文件中,我们可以写入你自己定义的类别名称。每一个类别换一行。这里我指定以了三个分类 smog、fire、person
在命令管理器里输入下面命令,进入labelimg界面
labelimg predefined_classed.txt
把收集的训练图片全部输入在fire目录下创建的Myimages文件夹
待标注图片数据的路径文件夹,选择Myimages文件夹
保存类别标签的路径文件夹,选择Annotations 文件夹
这个按键可以说明我们标注的标签为voc格式,点击可以换成yolo或者createML格式。
点击View,会出现如图红色框框中的选项。最好和我一样把勾勾勾上。
点击 Create RectBox,框选图片中的目标。保存名称为分类名。这里是smog、fire,然后点击Next Image切换写一个图片进行标记。
关闭labelimg,我们可以在Annotations 文件夹里,看到新生成的.xml文件。
1.2 划分训练集、验证集和测试集
在datasets目录下创建程序 xml2txt.py并运行,将XML格式转yolo_txt格式,代码如下(地址可修改)
import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join
classes = []
def convert(size, box):
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def convert_annotation(xmlpath, xmlname):
with open(xmlpath, "r", encoding='utf-8') as in_file:
txtname = xmlname[:-4] + '.txt'
txtfile = os.path.join(txtpath, txtname)
tree = ET.parse(in_file)
root = tree.getroot()
filename = root.find('filename')
img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
h, w = img.shape[:2]
res = []
for obj in root.iter('object'):
cls = obj.find('name').text
if cls not in classes:
classes.append(cls)
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
bb = convert((w, h), b)
res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
if len(res) != 0:
with open(txtfile, 'w+') as f:
f.write('\n'.join(res))
if __name__ == "__main__":
postfix = 'jpg'
imgpath = 'fire/Myimages'
xmlpath = 'fire/Annotations'
txtpath = 'fire/Mylabels'
if not os.path.exists(txtpath):
os.makedirs(txtpath, exist_ok=True)
list = os.listdir(xmlpath)
error_file_list = []
for i in range(0, len(list)):
try:
path = os.path.join(xmlpath, list[i])
if ('.xml' in path) or ('.XML' in path):
convert_annotation(path, list[i])
print(f'file {list[i]} convert success.')
else:
print(f'file {list[i]} is not xml format.')
except Exception as e:
print(f'file {list[i]} convert error.')
print(f'error message:\n{e}')
error_file_list.append(list[i])
print(f'this file convert failure\n{error_file_list}')
print(f'Dataset Classes:{classes}')
运行完成后再Mylabels文件下会生成xml文件对应的txt文件,保存不同图像的标注文件,每个图像对应一个txt文件,文件每一行为一个目标的信息,分别为class, x_center, y_center, width, height,这种为 yolo_txt格式。
将Myimages和Mylabels中的文件划分训练、验证、测试集,创建split_data.py并运行可完成此操作,代码如下(更改val_size和test_size控制验证和测试集比例)
import os, shutil
from sklearn.model_selection import train_test_split
val_size = 0.1
test_size = 0.2
postfix = 'jpg'
imgpath = 'fire/Myimages'
txtpath = 'fire/Mylabels'
os.makedirs('fire/images/train', exist_ok=True)
os.makedirs('fire/images/val', exist_ok=True)
os.makedirs('fire/images/test', exist_ok=True)
os.makedirs('fire/labels/train', exist_ok=True)
os.makedirs('fire/labels/val', exist_ok=True)
os.makedirs('fire/labels/test', exist_ok=True)
listdir = [i for i in os.listdir(txtpath) if 'txt' in i]
train, test = train_test_split(listdir, test_size=test_size, shuffle=True, random_state=0)
train, val = train_test_split(train, test_size=val_size, shuffle=True, random_state=0)
print(f'train set size:{len(train)} val set size:{len(val)} test set size:{len(test)}')
for i in train:
shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Myimages/train/{}.{}'.format(i[:-4], postfix))
shutil.copy('{}/{}'.format(txtpath, i), 'Mylabels/train/{}'.format(i))
for i in val:
shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Myimages/val/{}.{}'.format(i[:-4], postfix))
shutil.copy('{}/{}'.format(txtpath, i), 'Mylabels/val/{}'.format(i))
for i in test:
shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Myimages/test/{}.{}'.format(i[:-4], postfix))
shutil.copy('{}/{}'.format(txtpath, i), 'Mylabels/test/{}'.format(i))
运行完目录
1.3 定义训练参数文件
在 fire 文件夹下 新建一个 fire.yaml文件,具体内容如下
- train: 训练图片路径
- val:验证图片路径
- test: 测试图片路径
- names: 训练的类名称集合
修改yolov8.yaml,nc后面的参数改为训练的类数量
2. 训练模型
2.1训练数据
输入训练命令
yolo task=detect mode=train model=yolov8s.pt data=datasets/fire/fire.yaml epochs=100 batch=4
训练好的模型会被保存在 yolov8 目录下的 runs/detect/train/weights/ 下,生成 best.pt,last.pt文件。
2.2 验证数据
输入验证命令,用训练好的模型去验证
yolo task=detect mode=val model=runs/detect/train4/weights/best.pt data=datasets/fire/fire.yaml device=cpu
验证数据会被保存在 yolov8 目录下的 runs/detect/val/下,
2.3 预测数据
yolo task=detect mode=predict model=runs/detect/train4/weights/best.pt source=datasets/fire/images/val device=cpu
测试输出的数据会被保存在 yolov8 目录下的 runs/detect/predict/下
3.安卓端部署
上一篇yolov8模型安卓端部署详细介绍了部署步骤,现在只需将yolov8模型替换成自己训练的模型并修改相关参数即可
3.1 模型转换
转换自己训练的pt权重为ncnn格式
我们采用.pt ->onnx->ncnn的路线来转换自己训练的模型
(1)修改ultralytics/ultralytics/nn/modules/block.py中的
class C2f(nn.Module)如下:
def forward(self, x):
"""Forward pass through C2f layer."""
# y = list(self.cv1(x).chunk(2, 1)) #注释
# y.extend(m(y[-1]) for m in self.m) #注释
# return self.cv2(torch.cat(y, 1)) #注释
x = self.cv1(x)
x = [x, x[:, self.c:, ...]]
x.extend(m(x[-1]) for m in self.m)
x.pop(1)
return self.cv2(torch.cat(x, 1))
(2)修改ultralytics/ultralytics/nn/modules/head.py中的
class Detect(nn.Module)改动如下:
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
# if self.end2end:
# return self.forward_end2end(x)
#
# for i in range(self.nl):
# x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
# if self.training: # Training path
# return x
# y = self._inference(x)
# return y if self.export else (y, x)
shape = x[0].shape # BCHW
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training:
return x
elif self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# 中间部分注释掉,return语句替换为
return torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).permute(0, 2, 1)
(3)安装onnx包,在终端运行下面命令
pip install onnx coremltools onnx-simplifier
(4)根据官方文档,创建demo.py,将pt权重文件转换成onnx格式,代码如下
# 将模型导出为 ONNX 格式
from ultralytics import YOLO
model = YOLO("runs/detect/train4/weights/best.pt")
success = model.export(format="onnx", simplify=True, opset=11)
opset的参数(11、12、13皆可),使用其他会造成手机上运行出现很多框重叠错误识别
(5)对onnx文件进行压缩,进入到onnx文件所在目录,运行下面命令
python -m onnxsim best.onnx best-sim.onnx
压缩完之后会生成一个best-sim.onnx的文件,这一步是必须的,如果这一步不做,后面ONNX转NCNN可能会报错
完成上面步骤包含pt文件的weights目录下会生成对应文件
3.2 将onnx文件转换成param、bin文件
将best.onnx转成param和bin文件,可以下载以下文件之一
地址:Releases · Tencent/ncnn · GitHub
解压后,打开文件夹,并将best.onnx复制到该文件夹的对应位置,地址栏输入cmd后,在打开的命令行窗口输入onnx2ncnn.exe best-sim.onnx best.param best.bin回车即可,原文件夹生成了需要的bin文件和param文件
3.3 修改相关配置参数
(1)替换自己的模型到assets下
(2)修改strings.xml文件
把item替换为自己模型的名字
(3)修改yolov8ncnn.cpp文件
修改modeltypes变量为自己模型的名字,target_sizes为320,具体代码如下:
(4)修改yolo.cpp文件
1)Yolo::load方法中,模型名字修改,代码如下图:
2)Yolo::detect方法中,输入输出的名字修改。
可以使用onnx,通过网站link进行输入输出名字的查看。
具体修改代码内容如下:
这里修改了输入为Myimages,输出为output0
3)Yolo::draw方法中,修改class_names为你的类别名。这个模型中,只有一个类别,所以这里设置为smog、fire、person
4)修改yolo.cpp中的方法generate_proposals
修改num_class为你对应的类别数量,我们这里是一类,所以是3,具体代码如下:
完成上述步骤后,直接运行
3.4 运行测试
连上手机,点击run按钮,编译安装调试,运行成功,手机端已经安装好这个APP