★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

项目背景介绍

本项使用注意力机制融合的NA-Unet完成眼底血管语义分割，本项目使用自荷兰的糖尿病视网膜病变筛查项目数据集，自搭建Unet模型的encoder与decoder部分。在实现过程中，加入NA注意力模块，经实际测试对比，对眼底语义分割效果有较为明显的提升。

项目使用包的引入

import os
import shutil
import numpy as np
from paddle.io import Dataset,DataLoader
from paddle.vision import transforms as T
from paddle.nn import functional as F
import cv2
import paddle
import matplotlib.pyplot as plt
import paddle.nn as nn
from tqdm import tqdm

/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/__init__.py:107: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import MutableMapping
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/rcsetup.py:20: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import Iterable, Mapping
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/matplotlib/colors.py:53: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
  from collections import Sized

数据介绍

本项目使用自荷兰的糖尿病视网膜病变筛查项目数据集,这个数据集在原数据集的基础上做了数据增强（旋转、缩放）。

显示原数据集标注格式

• 原数据集的一组image和annotation如下：

np.set_printoptions(threshold=np.inf)
import matplotlib
path = "FundusVessels/Annotations/0.png"
# 标注数据RGB各颜色通道标记为1或0
# 各通道相同
im = cv2.imread(path)
im = im.transpose((2,0,1))
im_0 = im[0,:,:]
im_1 = im[1,:,:]
im_2 = im[2,:,:]

由下面两个cell可知，数据集Annotation中的图像RGB三通道每个通道的都使用0,1进行标记，各通道标记情况相同。

np.count_nonzero(im_0)
im_0

np.count_nonzero(im_1)
im_1

np.count_nonzero(im_2)

#验证集的数量
eval_num=99
#所有图像的大小
image_size=(224,224)
#训练图片路径
train_images_path=r"FundusVessels\JPEGImages"
#标签图像路径
label_images_path=r"FundusVessels\Annotations"
#测试图片路径
test_images_path=r"FundusVessels\Prediction"
#批量大小
batch_size=4

数据集的获取与预处理

这部分从路径以列表形式获取图像与测试图像，同时以元组形式合并一对image和label。

def get_path(image_path):
    files=[]
    for image_name in os.listdir(image_path):
        files.append(os.path.join(image_path,image_name))
    return sorted(files)

def get_test_data(test_images_path):
    test_data=[]
    for name in os.listdir(test_images_path):
        img_path=os.path.join(test_images_path,name)
        test_data.append(img_path)
    test_data=np.expand_dims(np.array(test_data),axis=1)
    return test_data

数据集的划分

train_transform=T.Compose([
            T.Resize(image_size),  #裁剪
            T.ColorJitter(0.1,0.1,0.1,0.1), #亮度，对比度，饱和度和色调
            T.Transpose(), #CHW
            T.Normalize(mean=0.,std=255.) #归一化
        ])
eval_transform=T.Compose([
            T.Resize(image_size),
            T.Transpose(),
            T.Normalize(mean=0.,std=255.)
        ])

创建数据读取器

class ImageDataset(Dataset):
    def __init__(self,path,transform):
        super(ImageDataset, self).__init__()
        self.path=path #[]
        self.transform=transform

    def _load_image(self,path):
        '''
        该方法作用为通过路径获取图像
        '''
        img=cv2.imread(path)
        img=cv2.resize(img,image_size)
        return img

    def __getitem__(self,index):
        
        # 是否transform
        path=self.path[index]
        if len(path)==2:# train,eval
            #print("len2:",path)
            data_path,label_path=path
            data,label=self._load_image(data_path),self._load_image(label_path)
            data,label=self.transform(data),label
            label = label.transpose((2, 0, 1))
            # print(label.shape) (3,224,224)
            label = label[0, :, :] # 取出一个通道
            # print("label") (224,224)
            label = np.expand_dims(label, axis=0)
            label = label.astype("int64")
            return data,label
            
        
        if len(path)==1:# test
            #print("len1:",path)
            data=self._load_image(path[0])
            data=self.transform(data)
            return data

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

#获取数据读取器
train_dataset=ImageDataset(train_data,train_transform)
eval_dataset=ImageDataset(eval_data,eval_transform)
test_dataset=ImageDataset(test_data,eval_transform)
train_dataloader=DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
eval_dataloader=DataLoader(eval_dataset,batch_size=batch_size,shuffle=True,drop_last=True)

查看paddle中gpu是否存在。

print(paddle.device.get_device())

cpu

模型的建立

下图为Unet模型架构图，

Encoder下采样部分

Encoder下采样部分两层卷积，两层归一化，最后池化。

class Encoder(nn.Layer):#下采样：两层卷积，两层归一化，最后池化。
    def __init__(self, num_channels, num_filters):
        super(Encoder,self).__init__()#继承父类的初始化
        self.features = nn.Sequential(
            nn.Conv2D(in_channels=num_channels,
                              out_channels=num_filters,
                              kernel_size=3,#3x3卷积核，步长为1，填充为1，不改变图片尺寸[H W]
                              stride=1,
                              padding=1),
            nn.BatchNorm(num_filters,act="relu"),
            nn.Conv2D(in_channels=num_filters,
                              out_channels=num_filters,
                              kernel_size=3,
                              stride=1,
                              padding=1),
            NA(),
            nn.BatchNorm(num_filters,act="relu"),
        )
        
        self.pool  = nn.MaxPool2D(kernel_size=2,stride=2,padding="SAME")#池化层，图片尺寸减半[H/2 W/2]
    def forward(self,inputs):
        x = self.features(inputs)
        x_conv = x           #两个输出，灰色 ->
        x_pool = self.pool(x)#两个输出，红色 | 
        return x_conv, x_pool
    

Decoder上采样部分

Decoder部分一层反卷积，两层卷积层，两层归一化

class Decoder(nn.Layer):#上采样：一层反卷积，两层卷积层，两层归一化
    def __init__(self, num_channels, num_filters):
        super(Decoder,self).__init__()
        self.features = nn.Sequential(
            nn.Conv2D(in_channels=num_filters*2,
                              out_channels=num_filters,
                              kernel_size=3,
                              stride=1,
                              padding=1),
            nn.BatchNorm(num_filters,act="relu"),
            nn.Conv2D(in_channels=num_filters,
                              out_channels=num_filters,
                              kernel_size=3,
                              stride=1,
                              padding=1),
            NA(),
            nn.BatchNorm(num_filters,act="relu")
        )
        self.up = nn.Conv2DTranspose(in_channels=num_channels,
                                    out_channels=num_filters,
                                    kernel_size=2,
                                    stride=2,
                                    padding=0)#图片尺寸变大一倍[2*H 2*W]

       
        
    def forward(self,input_conv,input_pool):
        x = self.up(input_pool)
        h_diff = (input_conv.shape[2]-x.shape[2])
        w_diff = (input_conv.shape[3]-x.shape[3])
        pad = nn.Pad2D(padding=[h_diff//2, h_diff-h_diff//2, w_diff//2, w_diff-w_diff//2])
        x = pad(x)                                #以下采样保存的feature map为基准，填充上采样的feature map尺寸
        x = paddle.concat(x=[input_conv,x],axis=1)#考虑上下文信息，in_channels扩大两倍
        x = self.features(x)
        return x

NA注意力模块

NA即归一化注意模块，学习了channel之间的相关性，筛选出了针对通道的注意力，并进行归一化。

NA注意力模块类

class NA(nn.Layer):
    def __init__(self):
        super().__init__()
        self.gamma = self.create_parameter([1], default_initializer=nn.initializer.Assign(paddle.ones([1])))
        self.beta = self.create_parameter([1], default_initializer=nn.initializer.Assign(paddle.zeros([1])))
        self.softmax = nn.Softmax(axis=2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        C = paddle.mean(x, axis=[2, 3], keepdim=True)
        
        g = -paddle.abs(x - C)
        b, c, h, w = g.shape
        g = paddle.sum(self.softmax(g.reshape((b, c, -1))), axis=2, keepdim=True).reshape((b, c, 1, 1))

        mean = paddle.mean(g, axis=1, keepdim=True)
        std = paddle.std(g, axis=1, keepdim=True) + 1e-5

        g = self.gamma * (g - mean) / std + self.beta

        out = x * self.sigmoid(g)

        return out

NA注意力模块的效果比较

当模型训练参数均设置为epochs=2,verbose=1，log_freq=1，learning_rate=1e-2,weight_decay=1e-3时

在模块构建中不添加NA模块，

示例训练结果如下图

在模块构建中添加NA模块，

示例训练结果如下图

由上面两张图可以看出，NA模块对眼底血管语义分割结果有较为明显的改进。

Unet模型搭建（利用Encoder，Decoder部分）

class UNet(nn.Layer):
    def __init__(self,num_classes=59):
        super(UNet,self).__init__()
        self.down1 = Encoder(num_channels=  3, num_filters=64) #下采样
        self.down2 = Encoder(num_channels= 64, num_filters=128)
        self.down3 = Encoder(num_channels=128, num_filters=256)
        self.down4 = Encoder(num_channels=256, num_filters=512)
        
        self.mid_conv1 = nn.Conv2D(512,1024,1)                 #中间层
        self.mid_bn1   = nn.BatchNorm(1024,act="relu")
        self.mid_conv2 = nn.Conv2D(1024,1024,1)
        self.mid_bn2   = nn.BatchNorm(1024,act="relu")

#         self.mid_conv1 = nn.Conv2D(256,512,1)                 #中间层
#         self.mid_bn1   = nn.BatchNorm(512,act="relu")
#         self.mid_conv2 = nn.Conv2D(512,512,1)
#         self.mid_bn2   = nn.BatchNorm(512,act="relu")

        self.up4 = Decoder(1024,512)                           #上采样
        self.up3 = Decoder(512,256)
        self.up2 = Decoder(256,128)
        self.up1 = Decoder(128,64)
        
        self.last_conv = nn.Conv2D(64,num_classes,1)           #1x1卷积，softmax做分类
        
    def forward(self,inputs):
        x1, x = self.down1(inputs)
        x2, x = self.down2(x)
        x3, x = self.down3(x)
        x4, x = self.down4(x)
        
        x = self.mid_conv1(x)
        x = self.mid_bn1(x)
        x = self.mid_conv2(x)
        x = self.mid_bn2(x)
        
        x = self.up4(x4, x)
        x = self.up3(x3, x)
        x = self.up2(x2, x)
        x = self.up1(x1, x)
        
        x = self.last_conv(x)
        
        return x

模型训练

定义优化器和损失函数，启动模型训练。

model = paddle.Model(UNet(2))    
print('model finish')
#定义优化器和损失函数
opt=paddle.optimizer.Momentum(learning_rate=1e-2,parameters=model.parameters(),weight_decay=1e-3)
model.prepare(opt, paddle.nn.CrossEntropyLoss(axis=1))
print('definition finish')
#启动模型训练
model.fit(train_dataloader, eval_dataloader, epochs=4,verbose=1,save_dir="./net_params",log_freq=1)
print('train')

模型的测试

读取训练好的模行参数文件，对Unet模型进行实例化，将预测结果写入"predict_result/"中。

from paddle.static import InputSpec

#预测图像保存路径
work_path = "predict_result/"
if os.path.exists(work_path):
    shutil.rmtree(work_path)
os.mkdir(work_path)

#读取模型参数文件路径
save_dir=work_path
checkpoint_path="./net_params/final"
                  
#实例化
model = paddle.Model(UNet(2))       #U-Net
model.load(checkpoint_path)

for i,img in tqdm(enumerate(test_dataset)):
    img=paddle.to_tensor(img).unsqueeze(0)
    predict=np.array(model.predict_batch(img)).squeeze(0).squeeze(0)
    predict=predict.argmax(axis=0)
    image_path=test_dataset.path[i]
    path_lst=image_path[0].split("/")
    lst = path_lst[-1][-13:-4]
    print(lst)
    lst = lst.replace("\\","/")
    print(lst)
    save_path=os.path.join(save_dir,lst)+".jpg"
    print(save_path)
  image_path=test_dataset.path[i]
    path_lst=image_path[0].split("/")
    lst = path_lst[-1][-13:-4]
    print(lst)
    lst = lst.replace("\\","/")
    print(lst)
    save_path=os.path.join(save_dir,lst)+".jpg"
    print(save_path)
    cv2.imwrite(save_path,predict*255)

测试结果如下：

参考文章

图像分割的打怪升级之路——UNet、PSPNet、Deeplab

U-Net: Convolutional Networks for Biomedical Image Segmentation

NA：归一化注意力机制

请点击此处查看本环境基本用法.

Please click here for more detailed instructions.

此文章为搬运
原项目链接