电脑基础 · 2023年4月19日

UNet 网络做图像分割DRIVE数据集

目录

1. 介绍

2. 搭建 UNet 网络

3. dataset 数据加载

4. train 训练网络

5. predict 分割图像

6. show

7. 完整代码


1. 介绍

项目的目录如下所示

  1. DRIVE 存放的是数据集
  2. predict 是待分割的图像
  3. result 里面放分割predict 的结果
  4. dataset 是处理数据的文件、model存放unet网络、predict是预测、train是网络的训练、UNet.pth 是训练好的权重文件

UNet 网络做图像分割DRIVE数据集

之前做了一个图像分割的例子,里面大部分的代码和本篇的内容重合,所以每个脚本的代码只会做简单的介绍。具体的可以参考之前的内容,这里给出链接:

model :  UNet - unet网络

dataset :UNet - 数据加载 Dataset

train : UNet - 训练数据train

predict : UNet - 预测数据predict(多个图像的分割)

DRIVE ( Digital Retinal Images for Vessel Extraction ):用于血管提取的数字视网膜图像

训练样本:灰度图像

UNet 网络做图像分割DRIVE数据集

 对应的标签:二值图像

UNet 网络做图像分割DRIVE数据集

因为这个分割项目完成几周了,最近才整理。所以,原数据集 DRIVE 可能是彩色图像 + mask 掩膜(具体的记不清了)

  • 这里没有使用 mask 
  • 如果是彩色图像的话,在生成unet网络的时候,传入的channel设置成3就行了。或者想用灰度图像的形式,要么用opencv转一下,可以看见灰度化的效果类似于展示的那样;要么在预处理的里面转成灰度图片 transform.Grayscale()

2. 搭建 UNet 网络

和之前unet网络不同的是,这里通过填充size,可以保证任意图像维度的输入

之前的代码需要经过4此下采样,每次维度扩展,size减半,所以需要保证输入图像的大小是 2的4次方

UNet 网络做图像分割DRIVE数据集

具体这块怎么实现我也看不懂,经过测试,可以实现任意输入的size

3. dataset 数据加载

数据加载的时候,将图像的预处理也放到了这里

这里训练的图像要 ToTensor ,归一化+改变通道顺序+转为tensor等等。同时,为了加快训练,对图像正规化,因为训练的图像是灰度图,所以只需要单通道的均值和标准差

UNet 网络做图像分割DRIVE数据集


然后是 数据加载 的初始化

UNet 网络做图像分割DRIVE数据集

这里的imgs里面的内容是,传入路径root下的图像路径,这里是:

['01.png', '02.png', '03.png', '04.png', '05.png', '06.png', '07.png', '08.png', '09.png']

self.imgs 是将root 路径和root 里面每个图像的路径 拼接在一块的路径,这里是:

['./DRIVE/test/image\\01.png', './DRIVE/test/image\\02.png', './DRIVE/test/image\\03.png', './DRIVE/test/image\\04.png', './DRIVE/test/image\\05.png', './DRIVE/test/image\\06.png', './DRIVE/test/image\\07.png', './DRIVE/test/image\\08.png', './DRIVE/test/image\\09.png']

如图:

UNet 网络做图像分割DRIVE数据集


 初始化路径和预处理后,需要对图像进行处理

UNet 网络做图像分割DRIVE数据集

这里训练的样本和对应的二值图像的label文件名要保证一样,否则需要做别的处理。例如,这里只需要将训练样本的图像路径里面的image 替换(replace)成label 就能找到对应的分割图像

UNet 网络做图像分割DRIVE数据集

然后读取图像,预处理之后,在进行返回即可。

这里为了防止label不是严格的二值图像,在归一化(灰度值 / 255)后,将中间的灰度值也映射为前景像素点

UNet 网络做图像分割DRIVE数据集

4. train 训练网络

训练网络的代码基本上没有改变,这里简单介绍

判断网络运行的设备,将网络to到device上

UNet 网络做图像分割DRIVE数据集

加载训练集+测试集

这里传入的是训练的样本,因为Data_loader 会将样本的路径替换成 label找到对应分割的标签图像

因为内存不足,所以这里将batch size 设置成 1

UNet 网络做图像分割DRIVE数据集

然后定义优化器+损失函数,并且保存网络的训练权重文件

有关BCEWithLogitsLoss可以参考这个:聊聊关于图像分割的损失函数 - BCEWithLogitsLoss

UNet 网络做图像分割DRIVE数据集

训练的时候,需要网络在train模式下,然后就是正确的前向传播预测+反向梯度下降的内容

UNet 网络做图像分割DRIVE数据集

最后是计算正确率,需要将网络放到eval模式下

这里将网络的预测转为二值图像,然后计算准确率的方式是预测的二值图像和label进行逐个像素点的比对,最后比上整幅图像的空间分辨率,即图像的大小。

test_label 的通道顺序是:batch、channel、height、width

UNet 网络做图像分割DRIVE数据集

5. predict 分割图像

这里的预处理要和处理样本的预处理一致

UNet 网络做图像分割DRIVE数据集

加载网络+读取网络参数

UNet 网络做图像分割DRIVE数据集

预测的时候,需要扩展维度。保存图像的时候,需要将batch和channel减去

然后将预测的结果转为二值图像就可以了

UNet 网络做图像分割DRIVE数据集

6. show

训练了20个epoch,结果显示如下

UNet 网络做图像分割DRIVE数据集

这里来预测的图像在test数据集里面,predict里面的图像为:

UNet 网络做图像分割DRIVE数据集

UNet 分割的结果:

UNet 网络做图像分割DRIVE数据集

真实的label为:

UNet 网络做图像分割DRIVE数据集

分割了大部分的信息,但是仍有细节没有分割出来

图像的size 是 565*584 的,大概预测的准确率是 0.96 左右

也就是说 还有 565*584*0.04 = 13198 ,这些损失的像素点就是缺少的细节

7. 完整代码

model部分:

import torch.nn as nn
import torch
import torch.nn.functional as F
# 搭建unet 网络
class DoubleConv(nn.Module):    # 连续两次卷积
    def __init__(self,in_channels,out_channels):
        super(DoubleConv,self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1,bias=False),
            nn.BatchNorm2d(out_channels),                           # 用 BN 代替 Dropout
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1,bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self,x):
        x = self.double_conv(x)
        return x
class Down(nn.Module):   # 下采样
    def __init__(self,in_channels,out_channels):
        super(Down, self).__init__()
        self.downsampling = nn.Sequential(
            nn.MaxPool2d(kernel_size=2,stride=2),
            DoubleConv(in_channels,out_channels)
        )
    def forward(self,x):
        x = self.downsampling(x)
        return x
class Up(nn.Module):    # 上采样
    def __init__(self, in_channels, out_channels):
        super(Up,self).__init__()
        self.upsampling = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) # 转置卷积
        self.conv = DoubleConv(in_channels, out_channels)
    def forward(self, x1, x2):
        x1 = self.upsampling(x1)
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])         # 确保任意size的图像输入
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        x = torch.cat([x2, x1], dim=1)  # 从channel 通道拼接
        x = self.conv(x)
        return x
class OutConv(nn.Module):   # 最后一个网络的输出
    def __init__(self, in_channels, num_classes):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, num_classes, kernel_size=1)
    def forward(self, x):
        return self.conv(x)
class UNet(nn.Module):   # unet 网络
    def __init__(self, in_channels = 1, num_classes = 1):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.in_conv = DoubleConv(in_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 1024)
        self.up1 = Up(1024, 512)
        self.up2 = Up(512, 256)
        self.up3 = Up(256, 128)
        self.up4 = Up(128, 64)
        self.out_conv = OutConv(64, num_classes)
    def forward(self, x):
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.out_conv(x)
        return x

dataset 数据处理部分:

import os
from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms
data_transform = {
    "train": transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, ), (0.5, ))]),
    "test": transforms.Compose([transforms.ToTensor()])
}
# 数据处理文件
class Data_Loader(Dataset):     # 加载数据
    def __init__(self, root, transforms_train=data_transform['train'],transforms_test=data_transform['test']):    # 初始化
        imgs = os.listdir(root)                                                         # 读取图像的路径
        self.imgs = [os.path.join(root,img) for img in imgs]                            # 取出路径下所有的图片
        self.transforms_train = transforms_train                                        # 预处理
        self.transforms_test = transforms_test
    def __getitem__(self, index):                      # 获取数据、预处理等等
        image_path = self.imgs[index]                  # 根据index读取图片
        label_path = image_path.replace('image', 'label')   # 根据image_path生成label_path
        image = Image.open(image_path)                      # 读取图片和对应的label图
        label = Image.open(label_path)
        image = self.transforms_train(image)        # 样本预处理
        label = self.transforms_test(label)         # label 预处理
        label[label > 0] = 1
        return image, label
    def __len__(self):  # 返回样本的数量
        return len(self.imgs)

train 网络训练部分:

from model import UNet
from dataset import Data_Loader
from torch import optim
import torch.nn as nn
import torch
# 网络训练模块
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # GPU or CPU
print(device)
net = UNet(in_channels=1, num_classes=1)        # 加载网络
net.to(device)                                  # 将网络加载到device上
# 加载训练集
trainset = Data_Loader("./DRIVE/train/image")
train_loader = torch.utils.data.DataLoader(dataset=trainset,batch_size=1,shuffle=True)
len = len(trainset)                         # 样本总数为 31
# 加载测试集
testset = Data_Loader("./DRIVE/test/image")
test_loader = torch.utils.data.DataLoader(dataset=testset,batch_size=1)
# 加载优化器和损失函数
optimizer = optim.RMSprop(net.parameters(), lr=0.00001,weight_decay=1e-8, momentum=0.9)     # 定义优化器
criterion = nn.BCEWithLogitsLoss()                             # 定义损失函数
# 保存网络参数
save_path = './UNet.pth'       # 网络参数的保存路径
best_acc = 0.0                 # 保存最好的准确率
# 训练
for epoch in range(20):
    net.train()     # 训练模式
    running_loss = 0.0
    for image,label in train_loader:
        optimizer.zero_grad()                          # 梯度清零
        pred = net(image.to(device))                   # 前向传播
        loss = criterion(pred, label.to(device))       # 计算损失
        loss.backward()                                # 反向传播
        optimizer.step()                               # 梯度下降
        running_loss += loss.item()                    # 计算损失和
    net.eval()  # 测试模式
    acc = 0.0   # 正确率
    total = 0
    with torch.no_grad():
        for test_image, test_label in test_loader:
            outputs = net(test_image.to(device))     # 前向传播
            outputs[outputs >= 0] = 1  # 将预测图片转为二值图片
            outputs[outputs < 0] = 0
            # 计算预测图片与真实图片像素点一致的精度:acc = 相同的 / 总个数
            acc += (outputs == test_label.to(device)).sum().item() / (test_label.size(2) * test_label.size(3))
            total += test_label.size(0)
    accurate = acc / total  # 计算整个test上面的正确率
    print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f %%' %
          (epoch + 1, running_loss/len, accurate*100))
    if accurate > best_acc:     # 保留最好的精度
        best_acc = accurate
        torch.save(net.state_dict(), save_path)     # 保存网络参数

predict 预测部分:

import numpy as np
import torch
import cv2
from model import UNet
from torchvision import transforms
from PIL import Image
transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,),(0.5))
    ])
# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = UNet(in_channels=1, num_classes=1)
net.load_state_dict(torch.load('UNet.pth', map_location=device))
net.to(device)
# 测试模式
net.eval()
with torch.no_grad():
    img = Image.open('./predict/img.png')           # 读取预测的图片
    img = transform(img)                            # 预处理
    img = torch.unsqueeze(img,dim = 0)              # 增加batch维度
    pred = net(img.to(device))                      # 网络预测
    pred = torch.squeeze(pred)                      # 将(batch、channel)维度去掉
    pred = np.array(pred.data.cpu())                # 保存图片需要转为cpu处理
    pred[pred >=0 ] =255                            # 转为二值图片
    pred[pred < 0 ] =0
    pred = np.uint8(pred)                           # 转为图片的形式
    cv2.imwrite('./result/res.png', pred)           # 保存图片