“添翼杯”人工智能创新应用大赛之垃圾分类

前言：

前一段时间参加了一个关于“添翼杯”人工智能创新应用大赛，题目要求是给你一张图像要你预测这个张图像属于可回收垃圾的概率，本质也就是一个图像二分类问题。简单记录一下比赛之中的经验吧。首先选择这个赛题是因为前一段时间总结了一下几个经典的图像识别网络，所以打算理论和实战结合看看各个模型在实际比赛中的效果。比赛成不成功对于我而言没有啥意义的，就是玩玩而已。

 

首先展示一下比赛数据集中的训练集和测试集图像的样式：

训练集中的不可回收垃圾，命名特点是以00开始，数据没有做任何的遮挡，比较干净。

 

训练集中的可回收垃圾，命名特点是以10开头，数据没有做任何的遮挡，比较干净。

 

 

测试集中数据的样式是：测试集中数据的特点是对图像做了旋转45度以及加了一些遮挡。因为题目是要求我们预测一张图像属于可回收垃圾的概率是多少，所以说将训练集中的图像进行一些预处理，使其和测试集中的分布一致这个很关键。对于这个加上一些个人的理解吧。例如你训练一个神经网络让它能够识别图像是不是猫，但是你给的数据集全部都是猫的全身照，但是在测试的时候，你给的是猫的半身照，这个时候有可能神经网络也能猜出来它是猫，因为神经网络发现了输入图像耳朵啊，眼睛这些特征和训练集的比较相似，所以他就预测这个是猫。但是神经网络输出的概率肯定相对于全身照的猫而言概率会低一些。

 

看完训练集和测试集中图像的样式，我们首先增加一些对原始的训练数据进行一些预处理操作也即是数据增广操作。主要是resize和rotate操作。代码如下。对于模型部分我主要使用两种方案。首先测试单模型效果，然后使用多个模型提取图像的特征，最后使用全连接层的进行图像分类。代码见一章

import torch
import torch.nn.functional as F

import argparse
from datetime import datetime

import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms

import random
import math

from PIL import Image
import matplotlib.pyplot as plt
import glob as gb
from skimage import io
import os
from shutil import rmtree


def cal_pad(w, h, r_w, r_h):
    left = math.floor(0.5 * (w - r_w))
    right = w - _left_ - r_w
    up = math.floor(0.5 * (h - r_h))
    down = h - _up_ - r_h
    return (left, up, right, down)


def my_image_transform(img_path, random_ratio, save_dir):

    a = torch.rand(6,1) * 2 - torch.ones((6,1))
    theta = torch.tensor([[a[0],                      -math.sqrt(1 - a[0]*a[0]) + a[1]/10, a[2]/4],
                          [-math.sqrt(1 - a[0]*a[0]), a[0] + a[4]/10,                      a[5]/4]], dtype=torch.float)

    img = Image.open(img_path).convert('RGB')
    io.imsave(save_dir + '/' + img_path.split('/')[-1].split('.')[0] + '.jpg', img)
    img1 = img.resize(
        (math.floor(np.array(img).shape[1] * random_ratio), math.floor(np.array(img).shape[0] * random_ratio)),
        Image.BILINEAR)

    _pad_ = cal_pad(np.array(img).shape[1],
                    np.array(img).shape[0],
                    math.floor(np.array(img).shape[1] * random_ratio),
                    math.floor(np.array(img).shape[0] * random_ratio))

    img_pil = transforms.ToPILImage()(np.array(img1))
    img_pad = transforms.Pad(padding=_pad_, fill=(0,0,0),
                             padding_mode='constant')(img_pil)

    io.imsave(save_dir + '/resize_' + img_path.split('/')[-1].split('.')[0] + '.jpg', img_pad)

    img_torch = transforms.ToTensor()(img_pad)

    grid = F.affine_grid(theta.unsqueeze(0), img_torch.unsqueeze(0).size())
    output = F.grid_sample(img_torch.unsqueeze(0), grid)

    new_img_torch = output[0]
    new_img = transforms.ToPILImage()(new_img_torch)
    io.imsave(save_dir + '/my_' + img_path.split('/')[-1].split('.')[0] + '.jpg', new_img)


if __name__=='__main__':

    save_dir = 'data/my_trainset'
    if os.path.exists(save_dir):
        rmtree(save_dir)
    os.mkdir(save_dir)

    import glob 
    img_paths = glob.glob('data/Trainset/*.jpg')

    for i, img_path in enumerate(img_paths):
 
        random_expand = [0.5, 0.7]
        random_ratio = random.randint(random_expand[0] * 100, random_expand[1] * 100) / 100

        my_image_transform(img_path, random_ratio, save_dir)