计算机视觉CV-Pytorch 天池计算机视觉入门赛:SVHN数据集实战

计算机视觉CV-Pytorch 天池计算机视觉入门赛:SVHN数据集实战

这里我们以datawhale和天池合作的天池计算机视觉入门赛为例，通过案例实战来进一步巩固本章所介绍的图像分类知识。

比赛简介与赛题分析

该比赛以SVHN街道字符为赛题数据，数据集报名后可见并可下载，该数据来自收集的SVHN街道字符，并进行了匿名采样处理，详细的介绍见赛事官网。

注：以下代码均默认已将比赛数据的根文件夹命名为tianchi_SVHN并放置于Dive-into-CV-PyTorch/dataset/tianchi_SVHN下

我们要做的就是识别图片中的数字串，赛题给定的数据图片中不同图片中包含的字符数量不等，如下图所示。

看起来好像有点棘手，和本章介绍的图像分类还并不一样。这里我们利用一个巧妙的思路，将赛题转化为一个分类问题来解。

赛题数据集大部分图像中字符个数为2-4个，最多的字符个数为6个。因此可以对于所有的图像都抽象为6个字
符的定长字符识别问题，少于6位的部分填充为X。

例如字符23填充为23XXXX，字符231填充为231XXX。

于是相当于将赛题转化为了分别对6个数字进行的分类问题，每个数字预测0-9/X。

注：这种思路显然不是本次比赛以及SVHN数据集的最佳解法，但却十分巧妙，这样设计的实战项目可以巩固对于本章图像分类知识的学习，考察灵活掌握的程度。

环境安装

本节所介绍的实战代码对环境没有特殊依赖，理论上Python2/3，Pytorch1.x版本均可以跑通。此外，由于数据集较小，有无GPU都是可以的，动手开始实战吧~

下面给出 python3.7 + torch1.3.1-gpu版本的环境安装示例

注：假设你已经安装了Anaconda及CUDA10.0,如果你对环境安装不太了解，请阅读第一章的环境安装部分教程。

首先在Anaconda中创建一个专门用于本次天池联系赛的虚拟环境。

$ conda create -n py37_torch131 python=3.7

然后激活环境，并安装gpu版本Pytorch

$ source activate py37_torch131

$ conda install pytorch=1.3.1 torchvision cudatoolkit=10.0

最后通过下面的命令一键完成其他依赖库的安装。

$ pip install tqdm pandas matplotlib opencv-python jupyter

首先导入必要的库

# -*- coding: utf-8 -*-                                                                          
import os, sys, glob, shutil, json
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

import cv2
import numpy as np
from PIL import Image

from tqdm import tqdm, tqdm_notebook

%pylab inline

import torch
torch.manual_seed(0)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset

定义读取数据集

class SVHNDataset(Dataset):
    def __init__(self, img_path, img_label, transform=None):
        self.img_path = img_path
        self.img_label = img_label
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None

    def __getitem__(self, index):
        img = Image.open(self.img_path[index]).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        lbl = np.array(self.img_label[index], dtype=np.int)
        lbl = list(lbl)  + (5 - len(lbl)) * [10]
        return img, torch.from_numpy(np.array(lbl[:5]))

    def __len__(self):
        return len(self.img_path)

定义读取数据dataloader

train_path = glob.glob('../../../dataset/tianchi_SVHN/train/*.png')
train_path.sort()
train_json = json.load(open('../../../dataset/tianchi_SVHN/train.json'))
train_label = [train_json[x]['label'] for x in train_json]
print(len(train_path), len(train_label))

train_loader = torch.utils.data.DataLoader(
    SVHNDataset(train_path, train_label,
                transforms.Compose([
                    transforms.Resize((64, 128)),
                    transforms.RandomCrop((60, 120)),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])),
    batch_size=40,
    shuffle=True,
    num_workers=10,
)

val_