随着当前网络科技的发展，我们可以从多种渠道获得蔬果的图像，甚至可以直接下载现有的蔬果数据集。比如在Kaggle网站中就有很多的水果数据集可供人们下载，其中最典型的有Fruit360数据集等，或者从Open Images数据集中选择需要的蔬果种类直接下载。甚至，你也可以自己拍摄采集图像形成自己的图像数据集。通常情况下，直接从网站中下载的现有数据集都有一个共同的特点——这些数据集都只能识别种类。于是，当你想要使用这些数据集想要去识别出更深层的意义时，现有的数据集便显得捉襟见肘。

基于上述的原因，我们可以采用基于现有图像，再次进行图像标记的方法来达到识别出更深层意义的目的。

2. 数据集划分——训练集、测试集和验证集

在训练模型的过程中，通常需要将准备训练的数据集划分为训练集、测试集和验证集。下面可以使用这个方法进行快速的划分：


import os  
import random  
import time  
import shutil
 
#设置图像文档和分类后的图像存放地址
imagefile_path = r"D:/COONEO/data"
#保存文件的地址
savefile_path = r"D:/COONEO/Annotations"
 
#设置训练集、验证集、测试集的比例
trainval_rate = 0.8   #训练验证集
train_rate = 0.8      #训练集
 
 
#读取图像数据集
total_images = os.listdir(imagefile_path)
#记录当前图像数据集的总数量
total_num = len(total_images)
image_list = range(total_num)
 
#计算验证集和测试集的数量
validation_set_num = int(total_num * trainval_rate)   #训练验证集的数量
train_set_num = int(validation_set_num * train_rate)  #训练集的数量
print("Validation set size:", validation_set_num)
print("Train set size:", train_set_num)
 
 
#随机进行采样作为训练集、测试集和验证集
random.seed(1)  
validation_sample = random.sample(image_list, validation_set_num) 
train_sample = random.sample(validation_sample, train_set_num)  
 
#数据集划分
start = time.time()    
test_num = 0   #测试集
train_num = 0   #训练集
val_num = 0    #训练验证集
 
for i in image_list:
    name = total_images[i]
    if i in validation_sample:
        if i in train_sample:
            directory = "train"
            train_num += 1
            images_path = os.path.join(os.getcwd(), "D:/COONEO/Annotations/{}".format(directory))
            if(not os.path.exists(images_path)):  
                #如果不存在directory所指示的文件夹，就用下列的命令进行创建
                os.mkdir(images_path)          
 
            file_path = os.path.join(imagefile_path, name)  
            newfile = os.path.join(savefile_path, os.path.join(directory,name))
            shutil.copyfile(file_path, newfile)           
 
        else:
            directory = "validation"
            images_path = os.path.join(os.getcwd(), "D:/COONEO/Annotations/{}".format(directory))
            if(not os.path.exists(images_path)):
                os.mkdir(images_path)                
 
            val_num += 1
            file_path = os.path.join(imagefile_path, name)
            newfile = os.path.join(savefile_path, os.path.join(directory,name))
            shutil.copyfile(file_path, newfile)
           
    else:
        directory = "test"
        images_path = os.path.join(os.getcwd(), "D:/COONEO/Annotations/{}".format(directory))
        if(not os.path.exists(images_path)):
            os.mkdir(images_path)
        
        test_num += 1 
        file_path = os.path.join(imagefile_path, name)
        newfile=os.path.join(savefile_path, os.path.join(directory,name))
        shutil.copyfile(file_path, newfile)
       
end = time.time() 
finish = end-start
print("Time taken:{0:.2f} seconds".format(finish))

完成数据集的划分工作后，我们需要依次针对训练集、验证集和测试集进行较为繁琐和枯燥的数据标注工作。

3.蔬果数据集标记

这里使用苹果成熟度的检测来作为方法分享的例子。首先我们将苹果分为了四种类别，分别是：

高成熟度：High Rigeness
中成熟度：Medium Rigeness
低成熟度：Low Rigeness
坏苹果：Bad Apple

然后，使用labelImg图像标注工具分别对训练集、测试集、验证集的苹果图像进行特征标注。

接下来，将和大家分享一下labelImg工具的具体安装和使用方法：

这里安装的环境主要是基于Windows10，然后使用Anaconda3对labelImg图像标注工具进行管理。具体的安装步骤如下：

(1)打开Anaconda Prompt。

(2)在Anaconda中创建一个虚拟环境，专门用来使用labelImg目标检测图形标记工具。这里需要注意，在安装的过程中一定要附加上python的安装版本，要不容易发生python版本与labelImg不匹配的问题，导致安装失败。

conda  create -n labelimg python=3.6

(3) 创建虚拟环境之后，使用 conda activate labelimg 激活当前虚拟环境。

运行语句后，可以看到当前的环境已经切换到了labelimg环境下。接下来我们的安装操作都在这个环境下进行。

(4)进入labelimg虚拟环境之后，使用 pip install labelimg语句来安装labelImg图像标注工具。

上图显示当前labelImg已经安装成功。

(5)在当前环境下，输入labelImg，便可以成功启动labelImg图像标记工具啦。

4. 创建VOC数据集

在使用自有数据集训练目标检测模型的过程中主要使用了目标检测领域的常用算法模型——SSD-MobileNet-V2。这个模型使用到的数据集格式是VOC。因此，我们需要基于现有的图像数据，自行创建VOC图像数据集。

那么怎么创建VOC数据集呢？

回到之前我们准备对划分好的图像训练集、验证集和测试集进行数据标注的工作。在已经安装完成的labelImg图像标记工具上分别对训练集、验证集和测试集进行数据标注。基于前文所提到的原因，我们在使用labelImg工具进行数据标注的过程中需要注意标注文件的保存格式。

在labelImg中，有三种类型的数据保存格式，并且在图像标记的过程中可以随时进行调换。这三种标注保存格式分别为：PescalVOC、YOLO和CreateML。这里由于选择的目标检测模型匹配的数据集是VOC数据集，所以在labelImg中选择PescalVOC格式进行保存。保存的数据标注文件后缀为.xml

下图中的红框部分即可切换保存标记的格式，这里默认使用PescalVOC格式。

苹果标注示意：

当全部图像全部标注完成时，我们需要将训练集、测试集和验证集生成的.xml标注文件各自保存到单独的文件夹，以便于接下来各自单独开展.xml文件格式的转换为.csv格式的工作。

下面的代码可以批量把生成的.xml文件快速转化为一个.csv文件。


import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET
 
def xml_to_csv(path):  
    xml_list = []  
    for xml_file in glob.glob(path + '/*.xml'):  
        tree = ET.parse(xml_file)  
        root = tree.getroot()
       
        print(root.find('filename').text)  
        for member in root.findall('object'): 
            value = (root.find('filename').text[:-4],  
                int(root.find('size')[0].text),   #Width  
                int(root.find('size')[1].text),   #Height  
                member[0].text,   
                int(member[4][0].text), 
                int(float(member[4][1].text)), 
                int(member[4][2].text),  
                int(member[4][3].text)
                )  
            xml_list.append(value)
    column_name = ['ImageID', 'width', 'height', 'ClassName', 'XMin', 'YMin', 'XMax', 'YMax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)  
    return xml_df

随后，直接调用下面的 store() 函数便可以直接进行上述方法的调用和生成的文件.csv文件的储存。


def store():
    for directory in ['train','test','validation']:
        xml_path = os.path.join(os.getcwd(), 'D:/COONEO/cooneo_xml/{}'.format(directory)) #存放.xml文件夹地址
        xml_df = xml_to_csv(xml_path)
        xml_df.to_csv('D:/COONEO/Annotations/sub-{}-annotations-bbox.csv'.format(directory),
                      index=None) 
        print("Successfully converted xml to csv.")

需要注意的是，上述代码中的文件命名格式最好不要进行更改，因为.csv文件的名称是按照模型训练的主文件中的读取方式进行命名的，这样可以便于模型训练，减少阅读代码的工作量。


#函数调用实现
store()

将.xml文件转化成为.csv文件之后，我们的数据集文件夹是如下的形式，下图文件夹中存放的是训练集、测试集和验证集的源图像。

到这里，我们使用自有数据集构建VOC数据集的工作算是告一个段落。准备好VOC数据集之后，我们便可以开始目标检测模型的训练。

03 模型训练

项目中的目标检测模型主要基于Pytorch(GPU)+Python3环境，为了减少模型训练环境搭建的繁琐工作，提高效率，我们可以直接云环境中进行。当然，下面的模型训练方法在本地环境上也同样适用。

1. 上载自建VOC数据集

首先，将之前准备好的VOC数据集打包压缩为zip文件夹，并上传到云存储中。然后在云环境中输入如下的语句解压压缩包。这样进行操作可以大幅度减少上传文件所消耗的时间，非常快速的就可将数据集上载到云环境中。


unzip 压缩包所在地址 -d 解压后文件的存放地址
 
unzip /content/drive/MyDrive/Cooneo/Apple-Rigeness/Annotations.zip -d /content/drive/MyDrive/Cooneo/Apple-Rigeness/Apple-Rigeness

运行后，图像数据集就已经存在于当前的云环境中。

2. 运行环境调节：GPU

随后，将云环境中的代码运行环境调整为GPU模式。


#查看当前GPU环境配置
nvidia-smi

上图显示了当前云环境所使用的显卡以及一些环境参数。随后，加载云存储至代码运行环境中，便于数据集的调取。


from google.colab import drive
drive.mount('/content/drive')

加载成功后，便可以下载预训练模型。

3. 下载预训练模型和训练脚本

使用下面的方法可以在线下载预训练模型至指定的文件夹中。

wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O /content/drive/MyDrive/Cooneo/models/mobilenet-v1-ssd-mp-0_675.pth

预训练模型下载完成后，需要继续将模型训练的代码克隆过来。

git clone https://github.com/dusty-nv/pytorch-ssd

4. 模型训练

将前面所提到的所有工作准备好之后，便可以正式开始模型的训练。


#在云环境下训练
/usr/bin/python3 /content/pytorch-ssd/train_ssd.py --data=/content/drive/MyDrive/Cooneo/Apple-R
#在本地环境训练
python3 train_ssd.py -data=数据集存放地址 --model-dir=训练后的模型存放地址 --batch-size=10 --epochs=100

在上述程序运行之前，需要打开 train_ssd.py文件，修改其中预训练模型的地址。如果不进行地址修改，将无法进行模型训练。因此，这一个细节至关重要。

同时，下面提供一些训练模型脚本的常用选项：

在训练模型的过程中，如果出现了因为内存不足或者是在训练过程中出现"killed"的情况，可以尝试减少--batch-size 或者--workers的参数。

模型训练结束后，将会在模型保存的地址中出现很多以 .pth为结尾的模型（说明训练完成后的模型格式为PyTorch格式）以及一个名为 labels.txt 的文件。

5. 模型格式转换

接下来，我们需要将训练好的模型从PyTorch转换为ONNX，以便于我们可以使用TensorRT进行加载。因为我们最终会使用搭载了NVIDIA Jetson的农业机器人来进行蔬果成熟度识别测试，同时Jetson上已经预先安装了TensorRT。因此，将模型的格式进行转换后，可以为后续的工作提供便利。


#在云环境中
/usr/bin/python3 /content/pytorch-ssd/onnx_export.py --input=/content/drive/MyDrive/Cooneo/Apple-Ripeness_model/mb1-ssd-Epoch-99-Loss-32774.71484375.pth --model-dir=/content/drive/MyDrive/Cooneo/Apple-Ripeness_model
#在本地运行
python3 模型脚本地址/onnx_export.py --input=最后训练好的模型 --model-dir=转换后的模型存放地址

上述语句运行后，PyTorch格式的模型就正式转变成了.onnx结尾的文件。后续只需将模型训练中生成的 labels.txt文件以及.onnx模型一起拷贝到NVIDIA Jetson中，便可以开始机器人进行目标检测的工作。

6. 使用TensorRT进行实时目标检测

将模型训练中生成的labels.txt文件以及.onnx模型一起拷贝到机器人的Jetson中，此时的操作环境已经变成了Jetson中内置的LINUX操作系统。这里会用到Jetson-inference，可以提前下载并配置其环境。之后就可以在Jetson-inference/

python/training/detection/ssd下创建一个新的文件夹 apple-rigeness 来存放前面所提到的两个文件。

随后，在 Jetson-inference/python/training/

detection/ssd 下打开终端，运行

$ detectnet.py --model=apple-rigeness/ssd-mobilenet_apple-Rigeness.onnx --labels=apple-rigeness//labels.txt --input-blob=input_0 --output-cvg=scores --output-bbox=boxes csi://0

第一次运行时，将会花费一些时间编译模型，但是一会儿便会看到机器人上的摄像机被成功调用，此时便正式开始了实时的目标检测。

下面附图一张~

以上便是如何使用自建数据集进行蔬果成熟度检测的全过程，不知道你学会了吗？

欢迎加入我们的交流群，该群面向热爱机器人研发的朋友们，方便大家一起学习、分享、交流智能机器人创造，结识更多志同道合的小伙伴。更有不定期的社区专属福利哦！关注公众号（COONEO）即可获取入群方式。

创作不易，如果喜欢这篇内容，请您也转发给您的朋友，一起分享和交流创造的乐趣，也激励我们为大家创作更多的机器人研发攻略，让我们一起learning by doing!

本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/凡人多烦事01/article/detail/510909