我們上次給新手們介紹了第一個合適入門的深度學習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有哪些研究和應用，如何循序漸進地學習好？

總結

以上是生活随笔為你收集整理的【CV实战】年轻人的第一个深度学习图像分割项目应该是什么样的（Pytorch框架）？...的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。