结合ModelArts打造基于Yolov5的垃圾分类识别模型

2020年6月参加了华为云AI算法大赛--垃圾分类识别大赛，虽然没有进入决赛，但是还是想将ModelArts在构建的模型经验分享出来，现在的开发的主流思想就是开源为王。yolov5是yolo系列的最新版本的目标检测模型，较之前的yolov3、yolov4的版本有很大的提升，官方给出的效果提升如下：

一、数据预处理

本次采用的数据集是官方提供的VOC2007格式的垃圾分类数据集。

由于yolov5需要的训练数据格式如下：

每个图像的标签文件应该可以通过在其路径名中简单地替换/images/*.jpg为来定位/labels/*.txt。示例图像和标签对为：

因此数据处理脚本如下：

# 运行成功后会生成如下目录结构的文件夹：
# trainval/
#    -images
#        -0001.jpg
#        -0002.jpg
#        -0003.jpg
#    -labels
#        -0001.txt
#        -0002.txt
#        -0003.txt
# 将trainval文件夹打包并命名为trainval.zip, 上传到OBS中以备使用。
import os
import codecs
import xml.etree.ElementTree as ET
from tqdm import tqdm
import shutil
import argparse


def get_classes(classes_path):
    '''loads the classes'''
    with codecs.open(classes_path, 'r', 'utf-8') as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

def convert(size, box):
    w = size[0]
    h = size[1]
    w_center = round((box[2] - box[0]) / w, 6)
    h_center = round((box[3] - box[1]) / h, 6)
    x_center = round(((box[2] + box[0]) / 2) / w, 6)
    y_center = round(((box[3] + box[1]) / 2) / h, 6)
#     dw = 1./(size[0])
#     dh = 1./(size[1])
#     x = (box[0] + box[2])/2.0 -1
#     y = (box[1] + box[3])/2.0 -1
#     w = box[2] - box[0]
#     h = box[3] - box[1]
#     x = x*dw
#     w = w*dw
#     y = y*dh
#     h = h*dh
    return x_center, y_center, w_center, h_center
    #return x_center, y_center, w_center, h



def creat_label_txt(soure_datasets, new_datasets):
    annotations = os.path.join(soure_datasets, 'VOC2007/Annotations')
    txt_path = os.path.join(new_datasets, 'labels')
    class_names = get_classes(os.path.join(soure_datasets, 'train_classes.txt'))

    xmls = os.listdir(annotations)
    for xml in tqdm(xmls):
        txt_anno_path = os.path.join(txt_path, xml.replace('xml', 'txt'))
        xml = os.path.join(annotations, xml)
        tree = ET.parse(xml)
        root = tree.getroot()

        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        line = ''
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in class_names:
                print('name error', xml)
                continue
            cls_id = class_names.index(cls)
            xmlbox = obj.find('bndbox')
            box = [int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text),
                   int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text)]
            if box[2] > w or box[3] > h:
                print('Image with annotation error:', xml)
            if box[0] < 0 or box[1] < 0:
                print('Image with annotation error:', xml)
            # width = round((box[2] - box[0]) / w, 6)
            # height = round((box[3] - box[1]) / h, 6)
            # x_center = round(((box[2] + box[0]) / 2) / w, 6)
            # y_center = round(((box[3] + box[1]) / 2) / h, 6)
            bbox = convert((w, h), box)
            line = line + str(cls_id) + ' ' + " ".join([str(a) for a in bbox])+'\n'
        with open(txt_anno_path, 'w') as f:
            f.writelines(line)

sets = ["train", "valid"]

def split_trainval():
    for image_set in sets:
        image_indexs = open(f"./source_data/trainval/VOC2007/ImageSets/Main/{image_set}.txt").read().strip().split()
        list_file = open(f"{image_set}.txt", "w")
        for image_index in image_indexs:
            list_file.write(f"datasets/trainval/images/{image_index}.jpg\n")
        list_file.close()

def creat_new_datasets(source_datasets, new_datasets):
    if not os.path.exists(source_datasets):
        print('could find source datasets, please make sure if it is exist')
        return

    if new_datasets.endswith('trainval'):
        if not os.path.exists(new_datasets):
            os.makedirs(new_datasets)
        os.makedirs(new_datasets + '/labels')
        print('copying images......')
        shutil.copytree(source_datasets + '/VOC2007/JPEGImages', new_datasets + '/images')
    else:
        print('最后一级目录必须为trainval,且为空文件夹')
        return
    print('creating txt labels:')
    creat_label_txt(source_datasets, new_datasets)
    split_trainval()
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--soure_datasets", "-sd", type=str, help="SODiC官方原始数据集解压后目录")
    parser.add_argument("--new_datasets", "-nd", type=str, help="新数据集路径，以trainval结尾且为空文件夹")
    opt = parser.parse_args()
    # creat_new_datasets(opt.soure_datasets, opt.new_datasets)

    soure_datasets = r'./source_data/trainval'
    new_datasets = r'./trainval'
    creat_new_datasets(soure_datasets, new_datasets)

数据处理完后，就符合yolov5的训练输入格式了。

2、创建数据data/coco128.yaml Dataset.yaml

三、选择模型和模型参数

由于yolov5提供5个不同的参数模型选择，我直接上效果最好的yolov5x模型参数。

4、开始训练模型

直接利用train.py开始训练，训练参数如下：

这里选择了预训练模型yolov5x.pt,如果不想从预训练，从头开始训练，去掉这个参数即可，开启multi_scale可以增加准确度。

python train.py  --data ./data/coco128.yaml --cfg ./models/yolov5x.yaml --weights weights/yolov5x.pt --multi-scale

训练过程如下：

我训练了120batch,最后能有

mAP@0.5最后能达到0.78.

训练过程的图片的识别效果如下：

5、通过modelarts推理代码部署。推理代码如下:

根据ModelArts的模型要求，导入模型

在测试集上的mAP@0.5 能达到 0.76，yolov5的表现没有让人失望。

以上就是基于yolov5的modelarts垃圾分类模型。欢迎大家拍砖。

参考：

https://github.com/ultralytics/yolov5