CV比赛入门

赛事

本文主要借鉴DataWhale的“如何打一个CV比赛”，项目来源于华为2022全球校园AI精英大赛，赛题二：车道渲染数据智能质检赛题

这次比赛是一个缺陷检测方面的图像分类赛。给出的数据为图片格式，分为问题图片和无问题图片。本题的任务是构建一种模型，根据地图渲染图片数据来预测图片是否存在问题。

模型搭建

配置环境

先配置环境，安装相关依赖库，并导入。环境的配置在之前的博文里有。

#安装相关依赖库 如果是windows系统，cmd命令框中输入pip安装，或在Jupyter notebook中!pip安装，参考上述环境配置
#!pip install pandas numpy cv2 torch torchvision codecs PIL glob
#---------------------------------------------------
#导入库
import os
import glob
from PIL import Image
import csv, time
import numpy as np

# pytorch相关
import torch
import torchvision
import torch.optim as optim
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.utils.data as data

自定义dataset

自定义一个读取文件的dataset·，重写init，getitem，len方法。

# 自定义读取数据集
class ImageSet(data.Dataset):
    def __init__(
            self,
            images,
            labels,
            transform):
        self.transform = transform
        self.images = images
        self.labels = labels

    def __getitem__(self, item):
        imagename = self.images[item]
        
        # 防止文件出错，这里生成一个随机的照片
        try:
            image = Image.open(imagename)
            image = image.convert('RGB')
        except:
            image = Image.fromarray(np.zeros((256, 256), dtype=np.int8))
            image = image.convert('RGB')

        image = self.transform(image)
        return image, torch.tensor(self.labels[item])

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

预处理数据

读取图片，标注数据，将标签二值化，数据扩增
数据扩增使用的是pytorch库的transform里的几个常见函数

• torchvision.transforms.Compose()类，主要作用是串联多个图片变换，就是将多个transform操作组合在一起
• transforms.Resize(256)是按照比例把图像最小的一个边长放缩到256，
• transforms.RandomHorizontalFlip依概率水平翻转
• transforms.ToTensor() 转换为tensor格式，这个格式可以直接输入进神经网络了
• transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])对像素值进行归一化处理。

import pandas as pd
import codecs

# 训练集标注数据
lines = codecs.open('train_label.csv').readlines()
train_label = pd.DataFrame({
    'image': ['train_image/' + x.strip().split(',')[0] for x in lines],
    'label': [x.strip().split(',')[1:] for x in lines],
})

# 将标签进行二值化处理
train_label['new_label'] = train_label['label'].apply(lambda x: int('0' in x))

import cv2, os
def check_image(path):
    try:
        if os.path.exists(path):
            return True
        else:
            return False
    except:
        return False
# 筛选路径存在的训练集
train_is_valid = train_label['image'].apply(lambda x: check_image(x) )
train_label = train_label[train_is_valid]

# 数据扩增方法
trfs = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 训练集dataset和dataloder
# 这里我们使用前1000张图片进行训练，后续可以自行修改
train_dataset = ImageSet(train_label['image'].values[:1000],
                         train_label['new_label'].values[:1000],
                         trfs)
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=1,
    pin_memory=True,
)

# 测试集dataset和dataloder
test_images = glob.glob('./test_images/*')
test_dataset = ImageSet(test_images, [0] * len(test_images), trfs)

test_loader = DataLoader(
    test_dataset,
    batch_size=32,
    shuffle=False,
    num_workers=5,
    pin_memory=True,
)

模型搭建

使用resnet18预训练模型，然后修改最终模型的输出层维度：

# 加载resnet18预训练模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = torch.nn.Linear(512, 2)
model = model.to('cuda') #使用GPU

# 模型优化器
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 模型损失函数
loss = nn.CrossEntropyLoss()

使用pytorch完成模型的训练：

# 设置迭代轮数epochs，可调整，轮数越多，所花时间越久
epochs = 3
for epoch in range(epochs):
    start_t = time.time()
    epoch_l = 0
    epoch_t = 0
    
    # 批量训练
    for batch_idx, batch in enumerate(train_loader):
        optimizer.zero_grad()
        image, label = batch
        image, label = image.to('cuda'), label.to('cuda')
        output = model(image) # 正向传播

        l = loss(output, label) # 计算损失
        l.backward()
        optimizer.step()

        batch_l = l.item()
        epoch_l += batch_l
        batch_t = time.time() - start_t
        epoch_t += batch_t
        start_t = time.time()
        
        # 打印loss
        if batch_idx % 10 == 0:
            print(l.item(), batch_idx, len(train_loader))

    epoch_t = epoch_t / len(train_loader)
    epoch_l = epoch_l / len(train_loader)
    print('...epoch: {:3d}/{:3d}, loss: {:.4f}, average time: {:.2f}.'.format(
        epoch + 1, epochs, epoch_l, epoch_t))

模型验证

对测试集进行预测，并生成提交文件。

model.eval()
to_prob = nn.Softmax(dim=1)
with torch.no_grad():
    imagenames, probs = list(), list()
    for batch_idx, batch in enumerate(test_loader):
        image, _ = batch
        image = image.to('cuda')
        pred = model(image)
        prob = to_prob(pred)
        prob = list(prob.data.cpu().numpy())
        probs += prob

https://developer.huawei.com/consumer/cn/activity/digixActivity/digixdetail/101655281685926449?ha_source=co10&ha_sourceId=89000243