【Yolov5】训练自己的yolov5模型并集成到安卓应用中【上】——模型训练

前言

最近接到一个项目，希望使用手机摄像头对图像数据进行采集，并使用训练好的模型对图像数据进行检测，主要任务包括：

1. 寻找一个轻量的检测模型，以方便集成到安卓应用中
2. 使用自己的检测数据集对模型进行训练
3. 探索模型集成到安卓应用中的方式

由于需要轻量模型，这里很自然的想到最近比较热的yolov5s模型。兔丁哥由于刚刚接触深度学习，比较喜欢简单容易上手的pytorch，而yolov5正好有pytorch版。但众所周知的是，pytorch在集成上并不如tensorflow方便高效，这也导致了pytorch在工业界不如tensorflow，这次项目让我对这一块深有体会。接下来我会从数据采集和处理，模型训练，模型导出，安卓集成等方面来介绍这个项目，由于内容比较多，将分成三篇文章对该项目进行讲解，分别是：

1. 【Yolov5】训练自己的yolov5模型并集成到安卓应用中【上】——模型训练
2. 【Yolov5】训练自己的yolov5模型并集成到安卓应用中【中】——模型转化
3. 【Yolov5】训练自己的yolov5模型并集成到安卓应用中【下】——模型集成

数据采集和处理

采集

由于项目中的数据集包括一些涉密数据，因此本文以检测车辆为例来介绍项目。
首先使用自己的摄像设备到大街上，停车场上采集数据（当然也可以爬取或使用网上的公开数据集），这里就直接下载汽车数据集了，展示如下：

标记

这里使用labelImg来对图像进行标记，官方文档对安装和使用进行了详细的说明，本文在Windows + Anaconda中使用，因此安装过程如下：

git clone https://github.com/tzutalin/labelImg.git
cd labelImg
conda install pyqt=5
conda install -c anaconda lxml
pyrcc5 -o libs/resources.py resources.qrc
python labelImg.py
# or
# python labelImg.py [IMAGE_PATH] [PRE-DEFINED CLASS FILE]

打开后展示如下界面：

按照图片中的指示设置后，可以左手按w启动标记，右手鼠标画框标记，左手按d保存并切换至下一张图片。

标记是个大工程，突然来了灵感，能不能在这种标记软件上集成现成的或者想训练的模型，边标记边训练，边训练边由模型标记下一张图片，然后手工调整模型标记结果，再训练，再模型标记，再调整……随着训练的进行，准确率的提高，将极大减少调整过程，这样不仅提高了标记速度，在标记完成后模型也训练好了，一举两得。先记下灵感，以后有能力了再和小伙伴慢慢实现。

处理

标记完成后，会生成xml文件如下所示：

	tagImage
	00001.jpg
	E:\tagImage\00001.jpg
	
		Unknown
	
	
		33
		33
		1
	
	0
	

为了能在yolov5中训练，还需要对标记结果进行处理，假设图片存储在images目录中，XML存储在tagXml，这里使用python类的形式对数据进行了处理，处理过程包括：

1. 检测是否存在漏标记的异常数据
2. 读取XML，并完成坐标转化
3. 将转化的坐标写入标记文件txt中
4. 拆分训练集和测试集
5. 保存处理结果到文件中

首先先导入必要的包

import os
import sys
import random
import xml.etree.ElementTree as ET
import matplotlib.image as mp

读取XML文件

ReadXML类用于读取xml文件，并提取图片尺寸信息及图片中的标记（方框及类别）。由于图片中的标记可能存在多个，所以返回的是list。

class ReadXML():
    def __init__(self, filename):
        self.filename = filename
        with open(self.filename,'rb') as xmlfile:
            tree = ET.parse(xmlfile)
            root = tree.getroot()
            size = root.find('size')
            self.imageSize = ( int(size.find('width').text), int(size.find('height').text) )

            self.objs = []

            for obj in root.iter("object"):
                cls = obj.find('name').text
                if cls not in classes:
                    continue
                classID = classes.index(cls)
                bndbox = obj.find('bndbox')
                boxCoor = ( float(bndbox.find('xmin').text), float(bndbox.find('xmax').text), float(bndbox.find('ymin').text), float(bndbox.find('ymax').text) )
                self.objs.append({'classID': classID,  'boxCoor' : boxCoor})

    def getImageSize(self):
        return self.imageSize

    def getObjs(self):
        return self.objs

坐标转化

XML标记文件中的方框信息取的是左上角的坐标(x1,y1)，和右小角的坐标(x2, y2)，而yolov5需要的坐标表示形式为方框中心点坐标与原图片的占比(x, y)和方框的宽度w和高度h与原图片的占比，convert函数实现了此功能：

def convert(imageSize, boxCoor):
    x = (boxCoor[0] + boxCoor[1]) / 2.0 / imageSize[0]
    y = (boxCoor[2] + boxCoor[3]) / 2.0 / imageSize[1]
    w = (boxCoor[1] - boxCoor[0]) * 1.0 / imageSize[0]
    h = (boxCoor[3] - boxCoor[2]) * 1.0 / imageSize[1]
    return x, y, w, h

检测异常数据

DataSet类用来对数据集进行处理，其中getOutliers用来检测异常数据，getDatas返回正确的数据集（即既有图片又有XML标签文件的数据）：

class DataSet():
    def __init__(self, imagePath, xmlPath):
        self.train, self.val, self.text = None, None, None

        self.imagePath = imagePath
        self.xmlPath = xmlPath

        imageFiles = os.listdir(self.imagePath)
        self.imageData = [name.split(".")[0] for name in imageFiles if name.endswith("jpg") or name.endswith("jpeg")]
        xmlFiles = os.listdir(self.xmlPath)
        self.xmlData = [name.split(".")[0] for name in xmlFiles if name.endswith("xml")]
        self.dataSet = [data for data in self.imageData if data in self.xmlData]

    def getDatas(self):
        return self.dataSet

    def getOutliers(self):
        return {"NoXML":[data for data in self.imageData if data not in self.xmlData], "NoImage": [data for data in self.xmlData if data not in self.imageData]}
   # ......省略部分内容......

生成txt标记文件

DataSet类的writeAnnotationsToFile函数在调用convert进行坐标转化后，将生成可供yolov5训练的txt标记文件

class DataSet():
    # ......省略部分内容......
    def writeAnnotationsToFile(self, pathname):

        pbar = ProgressBar(50)
        total = len(self.dataSet)
        count = 0

        for data in self.dataSet:
            xml = ReadXML(os.path.join(self.xmlPath, data + '.xml'))
            with open(os.path.join(pathname, data + '.txt'), 'w') as annFile:
                imageSize = xml.getImageSize()
                objs = xml.getObjs()
                for obj in objs:
                    x, y, w, h = convert(imageSize, obj['boxCoor'])
                    cid = obj['classID']
                    annFile.write(str(cid) + " " + str(x) + " " + str(y) + " " + str(w) + " " + str(h) + " \n")

            count += 1
            pbar.log(count/total, data + '.xml')
   # ......省略部分内容......

拆分数据集

DataSet类的splitSet函数用于拆分数据集，其中trainval_per表示训练和验证集占整体数据集的比例，剩余的为测试集，train_per为训练集占训练和验证集的比例，可以根据情况自行调整。__writeSetToFile用于将拆分的数据集写入文件中，方便查看，并且可以指定是否需要后缀名，及是否需要图片路径。

class DataSet():
    # ......省略部分内容......
    def splitSet(self, trainval_per = 0.8, train_per = 0.8):
        total = len(self.dataSet)
        trainval_total = int(trainval_per * total)
        train_total = int(train_per * trainval_total)

        trainval_temp = random.sample(self.dataSet, trainval_total)
        train_temp = random.sample(trainval_temp, train_total)
        self.trainval = [data for data in self.dataSet if data in trainval_temp]
        self.train = [data for data in  self.trainval if data in train_temp]
        self.val = [data for data in  self.trainval if data not in self.train]
        self.text = [data for data in self.dataSet if data not in self.trainval]

        return self.train, self.val, self.text

    # 省略各种get方法

    def __writeSetToFile(self, datas, filename, suffix = False, path = False):
        with open(filename, 'w') as outfile:
            if path:
                for data in datas:
                    if os.path.exists(os.path.join(self.imagePath, data + ".jpg")):
                        outfile.write(self.imagePath + '/' + data + ".jpg\n")
                    elif os.path.exists(os.path.join(self.imagePath, data + ".jpeg")):
                        outfile.write(self.imagePath + '/' + data + ".jpeg\n")

            elif suffix:
                for data in datas:
                    if os.path.exists(os.path.join(self.imagePath, data + ".jpg")):
                        outfile.write(data + ".jpg\n")
                    elif os.path.exists(os.path.join(self.imagePath, data + ".jpeg")):
                        outfile.write( data + ".jpeg\n")
                    
            else:
                for data in datas:
                    outfile.write(data + '\n')

    def writeTrainToFile(self, filename, suffix = False, path = False):
        self.__writeSetToFile(self.train, filename, suffix, path)

    def writeValToFile(self, filename, suffix = False, path = False):
        self.__writeSetToFile(self.val, filename, suffix, path)

    def writeTrainValToFile(self, filename, suffix = False, path = False):
        self.__writeSetToFile(self.trainval, filename, suffix, path)

    def writeTextToFile(self, filename, suffix = False, path = False):
        self.__writeSetToFile(self.text, filename, suffix, path)

    def writeAllToFile(self, filename, suffix = False, path = False):
        self.__writeSetToFile(self.dataSet, filename, suffix, path)

    def writeSetToFile(self, pathname, suffix = False, path = False):
        self.writeAllToFile(os.path.join(pathname, 'all.txt'), suffix, path)
        self.writeTextToFile(os.path.join(pathname, 'test.txt'), suffix, path)
        self.writeValToFile(os.path.join(pathname, 'val.txt'), suffix, path)
        self.writeTrainToFile(os.path.join(pathname, 'train.txt'), suffix, path)
        self.writeTrainValToFile(os.path.join(pathname, 'trainval.txt'), suffix, path)

处理入口

之后调用上诉函数对数据集进行处理。

if __name__ == "__main__":
    datas = DataSet("E:/data/images","E:/data/tagXml")
    print(datas.getOutliers())
    train, val, text = datas.splitSet(0.8, 0.8)
    # print(len(train), len(val), len(text))
    datas.writeSetToFile("E:/data", True, True)
    datas.writeAnnotationsToFile("E:/data/labels")

注意：建议将图片存放在images目录中，生成的txt标记存放在labels目录中，因为yolov5在训练使并不需要指定标签目录，只需指定图片目录即可，yolov5将在图片的当前目录下寻找同名的txt标记文件，或将路径中的images简单替换为labels目录，并在其目录下寻找同名txt标记文件。

模型训练

下载yolov5，并用自己喜欢IDE打开yolov5目录，将上诉images文件夹、labels文件夹、以及生成的训练集train.txt、验证集val.txt、测试集test.txt放在data目录下。兔丁哥对VSCode情有独钟，其文件目录如下所示：

安装依赖包

官方给的安装依赖包的方式是pip install -qr yolov5/requirements.txt，但我对conda情有独钟，打开requirements.txt文件，看到官方也提供了conda安装依赖包的方式

# Conda commands (in place of pip) ---------------------------------------------
# conda update -yn base -c defaults conda
# conda install -yc anaconda numpy opencv matplotlib tqdm pillow ipython
# conda install -yc conda-forge scikit-image pycocotools tensorboard
# conda install -yc spyder-ide spyder-line-profiler
# conda install -yc pytorch pytorch torchvision
# conda install -yc conda-forge protobuf numpy && pip install onnx==1.6.0  # https://github.com/onnx/onnx#linux-and-macos

其实只需安装几个重要的包即可：numpy, opencv, matplotlib, tqdm, pillow, ipython, scikit-image, tensorboard, pytorch, torchvision, protobuf。

更改配置文件

在data目录下新建一个car.yaml文件，参照coco.yaml写入自己的数据集目录，及类别信息，如下所示：

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: data\\train.txt  
val: data\\val.txt  
test: data\\test.txt  

# number of classes
nc: 1

# class names
names: ['car']

打开models/yolov5s.yaml，更改nc为自己的类别数，如下：

# parameters
nc: 1  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# ......省略后面其他内容......

下载Yolov5权重文件

虽然训练的时候会自动下载，但由于是外网数据，下载极其缓慢，因此可以使用一切下载工具到tutorial.ipynb中提到的Google Drive进行下载，并存放在models目录下，本文下载的是yolov5s.pt文件。

模型训练

虽然tutorial.ipynb中使用如下命令对模型进行训练：

python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --cfg yolov5s.yaml --weights yolov5s.pt --nosave --cache

但我更喜欢直接在train.py文件中进行修改，在文件末尾找到主入口if __name__ == '__main__':，修改如下：

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='models/yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path') # 存储模型结构的配置文件
    parser.add_argument('--data', type=str, default='data/car.yaml', help='data.yaml path') # 存储训练、测试数据的文件
    parser.add_argument('--hyp', type=str, default='', help='hyperparameters path, i.e. data/hyp.scratch.yaml')
    parser.add_argument('--epochs', type=int, default=30)       # 指的就是训练过程中整个数据集将被迭代多少次
    parser.add_argument('--batch-size', type=int, default=5, help='total batch size for all GPUs') # 一次看完多少张图片才进行权重更新，梯度下降的mini-batch,
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes') # 输入图片宽高
    parser.add_argument('--rect', action='store_true', help='rectangular training') # 进行矩形训练
    parser.add_argument('--resume', nargs='?', const='get_last', default=False, 
                        help='resume from given path/last.pt, or most recent run if blank')  # 恢复最近保存的模型开始训练
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') # 仅保存最终checkpoint
    parser.add_argument('--notest', action='store_true', help='only test final epoch') # 仅测试最后的epoch
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') # 进化超参数
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') # gsutil bucket
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') # 缓存图像以加快训练速度
    parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied') # 重命名results.txt to results_name.txt
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') # cuda device, i.e. 0 or 0,1,2,3 or cpu
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') # 多尺度训练，img-size +/- 50%
    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') # 单类别的训练集
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') # 使用adam优化
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--logdir', type=str, default='runs/', help='logging directory')

这里主要修改了

参数	default值	备注
–weights	models/yolov5s.pt	权重文件
–cfg	models/yolov5s.yaml	存储模型结构的配置文件
–data	data/car.yaml	存储训练、测试数据的文件
–epochs	30	训练过程中整个数据集的迭代次数
–batch-size	5	处理多少张图片后再进行权重更新
–img-size	[640, 640]	输入图片的宽和高

--epochs、--batch-size、--img-size可能需要根据自己电脑的配置来缩小，否则可能造成训练时间过长，显存溢出的问题。如果提示显存溢出，缩小--batch-size的值即可。

之后便可以在控制台输入以下命令进行训练了：

python train.py

在训练过程中，还可以输入以下命令打开tensorboard可视化训练过程

tensorboard --logdir runs

根据提示在浏览器上打开http://localhost:6006/即可查看，如下图所示：

模型检测

模型训练好后将在runs目录下生成训练过程和结果文件，其中runs\exp0\weights中存放着模型训练好的权重文件，此时就可以使用该权重文件进行检测。将需要检测的图片存放在data/tests/中，打开detect.py文件，在文件末尾找到主入口if __name__ == '__main__':，修改如下：

参数	default值	备注
–weights	runs\exp0\weights\best.pt	权重文件
–source	data/test	存储需要检测的图片
–output	data/output	存储检测的结果
–img-size	640	输入图片的宽和高

然后在控制台输入以下命令进行检测了：

python detect.py

之后就可以在data/output目录下查看检测的结果了。

总结

通过这次项目，感觉Yolov5的检测能力还是挺强的，准确度和速度都比较高，而且官方文档从训练到检测都写得很清楚，想训练自己的模型是非常容易的。由于兔丁哥并没有真正使用车辆数据再跑一次模型，因此文中图片较少，但主要的细节都已经记录下来，如有问题，欢迎留言讨论。接下来我将使用TorchScrip对模型进行转化，敬请期待。。。