【CV实战】年轻人的第一个深度学习图像分割项目应该是什么样的（Pytorch框架）？...

我们上次给新手们介绍了第一个合适入门的深度学习CV项目，可阅读【CV实战】年轻人的第一个深度学习CV项目应该是什么样的？（支持13大深度学习开源框架），本次我们再给大家介绍一个新的任务，图像分割，包括数据的处理，模型的训练与测试。

作者&编辑 | 言有三

本文资源与图像分割结果展示

本文篇幅：4800字

背景要求：会使用Python和Pytorch深度学习开源框架

附带资料：开源代码一份，支持Pytorch

数据一份：文末有获取方法

1 项目背景

图像处理中，研究者往往只对图像中的某些区域感兴趣，在此基础上才有可能对目标进行后续的处理与分析。图像分割技术就是把图像中属于目标区域的感兴趣区域进行半自动或者自动地提取分离出来，属于计算机视觉领域中最基础的任务之一。

为了让新手们能够一次性体验一个图像分割任务的完整流程，本次我们选择带领大家完成一个嘴唇图像分割任务，包括数据集的处理，模型的训练和测试，同时也将这次的实验与上一期内容结合起来。

2 数据处理

2.1 数据获取

我们上次已经介绍过，如果没有开源的数据集，我们首先要学会使用爬虫爬取图像，然后对获得的图片数据进行整理，包括重命名，格式统一，如果不清楚整个流程，可以参考我们上一次【CV实战】年轻人的第一个深度学习CV项目应该是什么样的？（支持13大深度学习开源框架），获取后整理的图像如下：

2.2 数据标注

接下来，我们需要对数据进行标注。图像分割任务要求对每一个像素进行预测，所以需要像素级别的标注结果，当然我们实际标注的时候往往是通过画轮廓形成闭合区域，开源的标注工具有很多，我们可以使用LabelMe等工具进行标注，当然你也可以使用其他工具进行标注，这里我们就不再展开讲了，因为本文面向的读者已经不是纯新手了。

标注完之后的样本和结果如下：

需要注意的是，标注的结果并不是我们用于训练的标签，因为图像分割本身是对每一个图像像素进行分类，在当前的开源框架中，每一个像素的类别也是从0，1，2，3这样的顺序依次增加。

所以在这里，我们一定需要注意训练时候的标签处理，这个大家根据自己的实际情况进行调整。

3 数据读取

得到了数据之后，接下来咱们使用Pytorch框架来进行模型的训练，首先需要实现的就是数据读取。Pytorch本身并没有图像分割任务的数据接口，所以我们需要自己定义，读取图像和掩膜，做一些简单的数据增强操作，我们定义一个类为SegDataset如下。

class SegDataset(Dataset):

    def __init__(self,filetxt,imagesize,cropsize,transform=None):

        lines = open(filetxt,'r').readlines()

        self.samples = []

        self.imagesize = imagesize

        self.cropsize = cropsize

        self.transform  = transform

        if self.transform is None:

            transform = transforms.Compose([

                   transforms.ToTensor(),

                   transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))

                   ]) 

        for line in lines:

            line = line.strip()

            imagepath,labelpath = line.split(' ')

            self.samples.append((imagepath,labelpath))

    def __getitem__(self,index):##获取数据

        imagepath,labelpath = self.samples[index]

        image = cv2.imread(imagepath)

        label = cv2.imread(labelpath,0) ##读成1通道

       

        ## 添加基本的数据增强，对图片和标签保持一致

        ## 添加固定尺度的随机裁剪，使用最近邻缩放（不产生新的灰度值）+裁剪

        image = cv2.resize(image,(self.imagesize,self.imagesize),interpolation=cv2.INTER_NEAREST)

        label = cv2.resize(label,(self.imagesize,self.imagesize),interpolation=cv2.INTER_NEAREST)

        offsetx = np.random.randint(self.imagesize-self.cropsize)

        offsety = np.random.randint(self.imagesize-self.cropsize)

        image = image[offsety:offsety+self.cropsize,offsetx:offsetx+self.cropsize]

        label = label[offsety:offsety+self.cropsize,offsetx:offsetx+self.cropsize]

       

        return self.transform(image),label ##只对image做预处理操作

    def __len__(self): ##统计数据集大小

        return len(self.samples)

上述的SegDataset类实现了图像分割任务数据的读取，有以下几个需要说明的地方：

(1) 输入的filetxt是我们预先准备好的文件，其中每一行按照[图片 标签]的对应格式存储着数据。

(2) 我们这里自己添加了一个随机裁剪的数据增强操作，对于裁剪类操作，标签图也需要在同样的裁剪参数下进行变换，对于颜色类操作则不需要。更多常见的数据增强操作可以查看下文：

【技术综述】深度学习中的数据增强方法都有哪些？

接下来我们就可以来测试一下数据集的读取，分别使用简单的方法和Dataloader接口。

## 简单方法，根据index来读取某一条数据

filetxt = "data/train.txt"

imagesize = 256

cropsize = 224

mydataset = SegDataset(filetxt,imagesize,cropsize)

print(mydataset.__length__())

image,label = mydataset.__getitem__(0) 

## 使用DataLoader来遍历，相关接口我们上期已经讲过，不再赘述

from torch.utils.data import DataLoader

batchsize = 64

train_dataset = SegDataset(train_data_path,imagesize,cropsize,data_transform)

train_dataloader = DataLoader(train_dataset, batch_size=batchsize, shuffle=True)

确认数据读取没有问题后，我们就可以开始来准备模型以及训练了。

 

4 模型训练

得到了数据之后，接下来咱们使用Pytorch框架来进行模型的训练，包括模型定义、结果保存与分析。

4.1 模型定义

接下来我们定义分割模型，首先是若干卷积层，然后是若干反卷积层。

import torch

from torch import nn

class simpleNet5(nn.Module):

    def __init__(self):

        super(simpleNet5, self).__init__()

        self.conv1 = nn.Sequential(

            nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1),

            nn.BatchNorm2d(32),

            nn.ReLU(True),

        )

        self.conv2 = nn.Sequential(

            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),

            nn.BatchNorm2d(64),

            nn.ReLU(True),

        )

        self.conv3 = nn.Sequential(

            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),

            nn.BatchNorm2d(128),

            nn.ReLU(True),

        )

        self.conv4 = nn.Sequential(

            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),

            nn.BatchNorm2d(256),

            nn.ReLU(True),

        )

        self.conv5 = nn.Sequential(

            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),

            nn.BatchNorm2d(512),

            nn.ReLU(True),

        )

        self.deconv1 = nn.Sequential(

            nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=3,

                               stride=2, padding=1, output_padding=1),

            nn.BatchNorm2d(256),

            nn.ReLU(True)

        )

        self.deconv2 = nn.Sequential(

            nn.ConvTranspose2d(256, 128, 3, 2, 1, 1),

            nn.BatchNorm2d(128),

            nn.ReLU(True)

        )

        self.deconv3 = nn.Sequential(

            nn.ConvTranspose2d(128, 64, 3, 2, 1, 1),

            nn.BatchNorm2d(64),

            nn.ReLU(True)

        )

        self.deconv4 = nn.Sequential(

            nn.ConvTranspose2d(64, 32, 3, 2, 1, 1),

            nn.BatchNorm2d(32),

            nn.ReLU(True)

        )

        self.deconv5 = nn.Sequential(

            nn.ConvTranspose2d(32, 16, 3, 2, 1, 1),

            nn.BatchNorm2d(16),

            nn.ReLU(True)

        )

        self.classifier = nn.Conv2d(16, 3, kernel_size=1)

    def forward(self, x):       

        out = self.conv1(x)     

        out = self.conv2(out)     

        out = self.conv3(out)   

        out = self.conv4(out)

        out = self.conv5(out)   

        out = self.deconv1(out) 

        out = self.deconv2(out)

        out = self.deconv3(out)

        out = self.deconv4(out)

        out = self.deconv5(out)

        out = self.classifier(out)

        return out

if __name__ == '__main__':

    img = torch.randn(2, 3, 224, 224)

    net = simpleNet5()

    sample = net(img)

    print(sample.shape)

在这里我们需要知道的是转置卷积的API为nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=1, padding=0, output_padding=0,groups=1, bias=True, dilation=1)，其中各个参数大家可以去查API说明。

卷积的具体配置如下：

反卷积层的具体配置如下：

4.2 训练

数据集接口准备好之后我们进行训练，完整代码如下。

import torch

import torch.nn as nn

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.autograd import Variable

import torchvision

from torchvision import datasets, models, transforms

import time

import os

from net import simpleNet5

from dataset import SegDataset

from tensorboardX import SummaryWriter

from torch.utils.data import DataLoader

import numpy as np

writer = SummaryWriter() #可视化

batchsize = 64

epochs = 200

imagesize = 256 #缩放图片大小

cropsize = 224 #训练图片大小

train_data_path = 'data/train.txt' #训练数据集

val_data_path = 'data/val.txt' #验证数据集

# 数据预处理

data_transform = transforms.Compose([

            transforms.ToTensor(),

            transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])

# 图像分割数据集

train_dataset = SegDataset(train_data_path,imagesize,cropsize,data_transform)

train_dataloader = DataLoader(train_dataset, batch_size=batchsize, shuffle=True)

val_dataset = SegDataset(val_data_path,imagesize,cropsize,data_transform)

val_dataloader = DataLoader(val_dataset, batch_size=val_dataset.__len__(), shuffle=True)

image_datasets = {}

image_datasets['train'] = train_dataset

image_datasets['val'] = val_dataset

dataloaders = {}

dataloaders['train'] = train_dataloader

dataloaders['val'] = val_dataloader

# 定义网络，优化目标，优化方法

device = torch.device('cpu')

net = simpleNet5().to(device)

criterion = nn.CrossEntropyLoss() #使用softmax loss损失，输入label是图片

optimizer = optim.SGD(net.parameters(), lr=1e-1, momentum=0.9)

scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1) #每50个epoch，学习率衰减

if not os.path.exists('checkpoints'):

    os.mkdir('checkpoints')

for epoch in range(1, epochs+1):

    print('Epoch {}/{}'.format(epoch, epochs - 1))

    for phase in ['train', 'val']:

        if phase == 'train':

            scheduler.step()

            net.train(True)  # Set model to training mode

        else:

            net.train(False)  # Set model to evaluate mode

        running_loss = 0.0

        running_accs = 0.0

        n = 0

        for data in dataloaders[phase]:

            imgs, labels = data

            img, label = imgs.to(device).float(), labels.to(device).float()

            output = net(img)

            loss = criterion(output, label.long()) #得到损失

            output_mask = output.cpu().data.numpy().copy()

            output_mask = np.argmax(output_mask, axis=1)

            y_mask = label.cpu().data.numpy().copy()

            acc = (output_mask == y_mask) #计算精度

            acc = acc.mean()

            optimizer.zero_grad()

            if phase == 'train':

            # 梯度置0，反向传播，参数更新

                loss.backward()

                optimizer.step()

            running_loss += loss.data.item()

            running_accs += acc

            n += 1

        epoch_loss = running_loss / n

        epoch_acc = running_accs / n

        if phase == 'train':

            writer.add_scalar('data/trainloss', epoch_loss, epoch)

            writer.add_scalar('data/trainacc', epoch_acc, epoch)

            print('train epoch_{} loss='+str(epoch_loss).format(epoch))

            print('train epoch_{} acc='+str(epoch_acc).format(epoch))

        else:

            writer.add_scalar('data/valloss', epoch_loss, epoch)

            writer.add_scalar('data/valacc', epoch_acc, epoch)

            print('val epoch_{} loss='+str(epoch_loss).format(epoch))

            print('val epoch_{} acc='+str(epoch_acc).format(epoch))

    if epoch % 10 == 0:

        torch.save(net, 'checkpoints/model_epoch_{}.pth'.format(epoch))

        print('checkpoints/model_epoch_{}.pth saved!'.format(epoch))

writer.export_scalars_to_json("./all_scalars.json")

writer.close()

在上面代码中，我们使用了带动量项的SGD作为优化方法，配置学习率为0.1，使用了CrossEntropy损失作为优化目标，相关代码如下：

criterion = nn.CrossEntropyLoss() #使用softmax loss损失，输入label是图片

optimizer = optim.SGD(net.parameters(), lr=1e-1, momentum=0.9)

scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1) #每50个epoch，学习率衰减

另外我们使用了tensorboardX工具来进行可视化，工具使用我们之前已经介绍过了，大家可以去查看之前图像分类的内容，训练的结果如下：

5 模型测试

上面已经训练好了模型，我们接下来的目标，就是要用它来做推理，真正把模型用起来，下面我们载入一个图片，用模型进行测试。

import torch

import torch.nn as nn

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.autograd import Variable

import torchvision

from torchvision import datasets, models, transforms

import time

import os

import cv2

import sys

import torch.nn.functional as F

import numpy as np

data_transforms =  transforms.Compose([

            transforms.ToTensor(),

            transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])

modelpath = sys.argv[1] #模型目录

net = torch.load(modelpath,map_location='cpu')

net.eval() #设置为推理模式，不会更新模型的k，b参数

imagepaths = os.listdir(sys.argv[2]) #测试图片目录

torch.no_grad() #停止autograd模块的工作，加速和节省显存

for imagepath in imagepaths:

    image = cv2.imread(os.path.join(sys.argv[2],imagepath)) #读取图像

    image = cv2.resize(image,(224,224),interpolation=cv2.INTER_NEAREST)

    imgblob = data_transforms(image).unsqueeze(0) #填充维度，从3维到4维

    predict = F.softmax(net(imgblob)).cpu().data.numpy().copy() #获得原始网络输出，多通道

    predict = np.argmax(predict, axis=1) #得到单通道label

    result = np.squeeze(predict) #降低维度，从4维到3维

    print(np.max(result)) 

    result = (result*127).astype(np.uint8) #灰度拉伸，方便可视化

    resultimage = image.copy()

    for y in range(0,result.shape[0]): 

        for x in range(0,result.shape[1]):

            if result[y][x] == 127:

                resultimage[y][x] = (0,0,255) 

            elif result[y][x] == 254:

                resultimage[y][x] = (0,255,255) 

    combineresult = np.concatenate([image,resultimage],axis=1)

    cv2.imwrite(os.path.join(sys.argv[3],imagepath),combineresult) #写入新的目录

从上面的代码可知，使用torch.load函数载入模型，然后读取图像，进行与训练相同的预处理操作，就可以得到网络输出，再进行一些维度变换和softmax操作就得到最终的结果。

下面展示了一些分割结果。

从图中我们可以看到，总体的分割结果还是不错的，不过本次的任务还有许多可以提升的空间，包括但不限于：(1) 做更多的数据增强。(2) 改进模型。这些就留给读者去进行实验。

然后就可以自己输入图片得到推理结果，index就是预测的类别。

 6 资源获取和拓展学习

本文的完整代码，可以在我们的开源项目中获取，项目地址如下。

https://github.com/longpeng2008/yousan.ai

由于数据集较大，如果想要获得数据集，请到有三AI知识星球中下载，当然你也完全可以将其替换成自己的数据集：

有三AI知识星球链接与介绍如下：

【杂谈】有三AI知识星球指导手册出炉！和公众号相比又有哪些内容？

如果想要学习更多图像分割的内容，请关注下面的视频课程。

【视频课】图像分割最新内容来了（言有三新录制6大理论部分+1个案例实践讲解）

总结

本次我们完成了一个图像分割任务项目的全部流程，如果小伙伴还记得上一次的表情分类任务，那么就可以接着在上一次任务的基础上将其做成分割任务，得到本文开头的gif展示结果。

转载文章请后台联系

侵权必究

往期文章

• 【CV夏季划】2021年有三AI-CV夏季划出炉，冲刺秋招，从CV基础到模型优化彻底掌握
  

• 【CV秋季划】模型优化很重要，如何循序渐进地学习好？

• 【CV秋季划】人脸算法那么多，如何循序渐进地学习好？

• 【CV秋季划】图像质量提升与编辑有哪些研究和应用，如何循序渐进地学习好？

• 【CV秋季划】生成对抗网络GAN有哪些研究和应用，如何循序渐进地学习好？