【PaddlePaddle】使用Tiny-Yolo做“垃圾图像”目标识别

PaddlePaddle-使用Tiny-Yolo实现“垃圾图像”目标识别

---

代码环境：

在AI Studio平台基于
Python 2.7
PaddlePaddle 1.5.0

在本地基于
windows 10 1803
Python 3.7
PaddlePaddle 1.5.1
cuda 8.0
cudnn 7.5.1

---

写在前面
前段时间，在参加2019年的人工智能创意赛，在初赛的时候，我们团队是利用残差神经网络ResNet-50做了一个垃圾图像分类的项目，之后会把初赛的项目也写入blog里。进入复赛时，老师给我们的建议是：“残差做不了目标识别，换个模型”。

本文主要是讲讲数据的获取、预处理以及Paddle-Reader的构造，最后分享一下在本地实现实时预测的方法。训练的过程就不在本文阐述了，下一篇文章会分享一下我们使用基于mobilenet-v1的SSD来实现垃圾图像的识别。

---

数据集来源与制作

在初赛的时候，我们的数据集来自斯坦福大学Gary Thung，Mindy Yang的学年论文《Classification of Trash for Recyclability Status》，数据集中原本有6类可回收垃圾，“glass”,“cardboard”,“plastic”,“paper”,“trash”,“metal”，共2527张垃圾图像。

• 594 paper
• 501 glass
• 137 trash
• 410 metal
• 482 plastic
• 403 cardboard

经过筛选，我们删除了“cardboard”，“trash”两类图片，原因分别是：“cardboard”类图像占总图片的面积过大，我们担心会影响该类垃圾的识别；“trash”类图像数量过少，仅有137张，数量相较于其他类过少，我们担心该类垃圾会成为模型中的“奇异值”。

斯坦福论文垃圾图像数据集

同时，我们项目组补充了252张多目标图片加入训练数据集。

然后，使用LabelImg对每张垃圾图像进行目标标识和标注，制作符合 Pascal VOC 标准的数据集 。LabelImg的使用方法，我在此简单介绍一下。

LabelImg 是一个可视化的图像标定工具。Faster R-CNN，YOLO，SSD等目标检测网络所需要的数据集，均需要借此工具标定图像中的目标。生成的 XML 文件是遵循 PASCAL VOC 的格式的。

图像预处理

在图片预处理工作上，我们全部使用第三方库Pillow，做图形图像处理。
train 和 test 在图像预处理 (图像增强) 的部分使用的操作不相同：
在训练的过程中，尽可能希望模型能够看到更多更丰富更多样化的输入数据，所以经常会使用类似于

• 随机调整图像明度、色彩、对比度、饱和度
• 随机翻转
• 随机切块

"image_distort_strategy": {
   
        "expand_prob": 0.5,
        "expand_max_ratio": 4,
        "hue_prob": 0.5,
        "hue_delta": 18,
        "contrast_prob": 0.5,
        "contrast_delta": 0.5,
        "saturation_prob": 0.5,
        "saturation_delta": 0.5,
        "brightness_prob": 0.5,
        "brightness_delta": 0.125
    }
    
def random_brightness(img):
    prob = np.random.uniform(0, 1)
    if prob < 0.5:
        brightness_delta = 0.125
        delta = np.random.uniform(-brightness_delta, brightness_delta) + 1
        img = ImageEnhance.Brightness(img).enhance(delta)
    return img


def random_contrast(img):
    prob = np.random.uniform(0, 1)
    if prob < 0.5:
        contrast_delta = 0.5
        delta = np.random.uniform(-contrast_delta, contrast_delta) + 1
        img = ImageEnhance.Contrast(img).enhance(delta)
    return img


def random_saturation(img):
    prob = np.random.uniform(0, 1)
    if prob < 0.5:
        saturation_delta = 0.5
        delta = np.random.uniform(-saturation_delta, saturation_delta) + 1
        img = ImageEnhance.Color(img).enhance(delta)
    return img


def random_hue(img):
    prob = np.random.uniform(0, 1)
    if prob < 0.5:
        hue_delta = 18
        delta = np.random.uniform(-hue_delta, hue_delta)
        img_hsv = np.array(img.convert('HSV'))
        img_hsv[:, :, 0] = img_hsv[:, :, 0] + delta
        img = Image.fromarray(img_hsv, mode='HSV').convert('RGB')
    return img

生成图像列表及Paddle-Reader构造

生成图像列表

for images in all_images:
    trainval = []
    test = []
    if data_num % 10 == 0:
        # 每10张图像取一个做测试集
        name = images.split('.')[0]
        annotation = os.path.join(annotation_path, name + '.xml')
        # 如果该图像的标注文件不存在，就不添加到图像列表中
        if not os.path.exists(annotation):
            continue
        test.append(os.path.join(images_path, images))
        test.append(annotation)
        # 添加到总的测试数据中
        test_list.append(test)
    else:
        # 其他的的图像做训练数据集
        name = images.split('.')[0]
        annotation = os.path.join(annotation_path, name + '.xml')
        # 如果该图像的标注文件不存在，就不添加到图像列表中
        if not os.path.exists(annotation):
            continue
        trainval.append(os.path.join(images_path, images))
        trainval.append(annotation)
        # 添加到总的训练数据中
        trainval_list.append(trainval)
    data_num += 1

生成图像列表train.txt

构造同步读取reader

def custom_reader(file_list, data_dir, input_size, mode):
    def reader():
        np.random.shuffle(file_list)
        for line in file_list:
            if mode == 'train' or mode == 'eval':
                ######################  以下可能是需要自定义修改的部分   ############################
                image_path, label_path = line.split()
                image_path = os.path.join(data_dir, image_path)
                label_path = os.path.join(data_dir, label_path)
                img = Image.open(image_path)
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                im_width, im_height = img.size
                # layout: label | xmin | ymin | xmax | ymax | difficult
                bbox_labels = []
                root = xml.etree.ElementTree.parse(label_path).getroot()
                for object in root.findall('object'):
                    bbox_sample = []
                    # start from 1
                    bbox_sample.append(float(train_parameters['label_dict'][object.find('name').text]))
                    bbox = object.find('bndbox')
                    difficult = float(object.find('difficult').text)
                    bbox_sample.append(float(bbox.find('xmin').text) / im_width)
                    bbox_sample.append(float(bbox.find('ymin').text) / im_height)
                    bbox_sample.append(float(bbox.find('xmax').text) / im_width)
                    bbox_sample.append(float(bbox.find('ymax').text) / im_height)
                    bbox_sample.append(difficult)
                    bbox_labels.append(bbox_sample)
                ######################  可能需要自定义修改部分结束   ############################
                if len(bbox_labels) == 0: continue
                img, sample_labels = preprocess(img, bbox_labels, input_size, mode