admin 管理员组

文章数量: 887007

目录

1. 训练前准备

1.1 收集训练数据,完成分类标记

1.2 划分训练集、验证集和测试集

1.3 定义训练参数文件

2. 训练模型

2.1训练数据

​2.2 验证数据

2.3 预测数据

3.安卓端部署

3.1 模型转换

3.2 将onnx文件转换成param、bin文件

3.3 修改相关配置参数

3.4 运行测试


上一篇 yolov8的完整部署 讲解了Yolov8模型的部署和运行

实际生活中,我们需要的检索目标,很多并不在官方给出的模型之中。那我们就需要训练自己的数据模型。

1. 训练前准备

1.1 收集训练数据,完成分类标记

首先,我们在根目录datasets/文件夹下像存放coco128数据的格式创建fire文件,在fire/文件下创建Myimages和Annotations文件夹,这个后面我们会用到。

这里我们要用到一个工具 labelimg,通过下面命令安装这个工具:

pip install labelimg -i https://pypi.tuna.tsinghua.edu/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

本文标签: 模型 到安卓端