本周的总结

论文阅读

U-Net: Convolutional Networks for Biomedical Image Segmentation （U-Net：用于生物医学图像分割的卷积网络）

Abstract（摘要）

因为深度网络的成功训练需要数千个带注释的训练样本，在本文中，我们提出了一种网络和训练策略，该策略依赖于大量使用数据增强来更有效地利用现在已有的可用的标注样本，该体系结构由捕获上下文的收缩路径（a contracting path to capture context ）和支持精确本地化的对称扩展路径组成（a symmetric expanding path that enables precise localization），我们表明，这样的网络可以从很少的图像端到端地训练，并且在ISBI挑战中表现优于先前最好的方法(滑动窗口卷积网络)。

Introduction（介绍）

卷积网络的典型应用是在分类任务中，图像的输出是单个类别标签。然而，在许多视觉任务中，特别是在生物医学图像处理中，期望的输出应该包括定位（localization），即应该为每个像素分配一个类别标签。此外，在生物医学任务中，数以千计的训练图像通常是遥不可及的（beyond reach）。所以，Ciresan等训练了一个神经网络，用滑动窗口来预测每个像素的类标签，提供像素的周围区域（patch）作为输入。首先，这个网络可以定位。第二，输入的是patches，这样训练数据就比图片数据多很多。这个网络大幅度赢得了ISBI2012。但是，明显的，Ciresan的方法有两个缺点。第一，它很慢，因为这个网络必须训练每个patch，并且因为patch间的重叠有很多的冗余，这个权重会被同一些特征训练两次，造成资源的浪费，减慢训练时间和效率，容易导致过拟合，第二，定位准确性和上下文间不可兼得。许多现在的方法使用不同层的特征来同时兼容定位和利用context。

在本文中，我们建立在一种更优雅的体系结构之上，即所谓的“完全卷积网络”（FCN）。我们对该体系结构进行了修改和扩展，使其能够处理非常少的训练图像，并产生更精确的分割。

 主要思想是用连续的层来补充通常的收缩网络，其中用上采样操作符代替池操作符。因此，这些层提高了输出的分辨率。为了进行本地化，将收缩路径中的高分辨率要素与上采样的要素相结合（即跳跃连接skip layer）。

我们架构中的一个重要修改是，在上采样部分，我们也有大量的特征通道，允许网络将上下文信息传播到更高分辨率的层。使得扩张路径与收缩路径对称，并产生U形结构（所以叫Unet）。该网络不具有任何完全连接的层，并且仅使用每个卷积的有效部分，即分割图仅包含在输入图像中可获得其完整上下文的像素。为了预测图像边界区域中的像素，通过镜像输入图像来预测缺失的上下文。这种平铺策略（tiling strategy ）对于将网络应用于大型图像非常重要，否则分辨率将受到GPU内存的限制。

因为我们的任务只能获得很少的训练数据，我们通过弹性变形来数据增强。这在生物医学图像分割是很重要的，因为变形很常见。另一个细胞分割的挑战是同一类别连接物体的分离。如图三。为此，我们使用加权损失函数，连接细胞间的背景标签分离获得较大的权重。

 Network Architecture（网络架构）

网络结构大体分为收缩和扩张路径来组成。因为形似一个字母U，得名Unet。收缩路径（contracting path (left side)）仍然是利用传统卷积神经网络的卷积池化组件，其中经过一次下采样之后，channels变为原来的2倍。扩张路径（expansive path (right side).）由2 * 2的反卷积，反卷机的输出通道为原来通道数的一半，再与原来的feature map（裁剪之后）串联，得到和原来一样多的通道数的feature map，再经过2个尺寸为3 * 3的卷积和ReLU的作用。裁剪特征图是必要的，因为在卷积的过程中会有边界像素的丢失。在最后一层通过卷积核大小为1 * 1的卷积作用得到想要的目标种类。在Unet中一共有23个卷积层。但是这个网络需要谨慎的选择输入图片的尺寸，以保证所有的Max Pooling操作作用于长宽为偶数的feature map。

Training（训练）

输入图像及其对应的分割图用随机梯度下降实现来训练网络，损失函数如下：

P（x）表示每一个像素点预测的分类概率

 这个公式表示的是：给像素分配权重然后进行加权，d1(x)表示图中某一背景像素点到离这个点最近的细胞边界的距离，d2(x)表示离这个像素点第二近的细胞的距离，即在细胞边界附近的像素点给的权重会大一些，离细胞比较远的像素点的权重会小一些，为什么这么做呢？因为，如果同类细胞贴的比较近，可能就会增大训练的难度，减少准确率，毕竟卷积会考虑该像素点周围的一些特征，而两个相同的类的细胞贴在一起，就容易误判，所以对这种两个相同类贴在一起的细胞边界，给予较大的权重，使的训练之后分类分割更准确。

Experiments

 可以看到各种Error（分割评估指标），都比较好

Conclusion

U-Net架构在很多不同的生物医学分割应用上取得了非常好的性能，并且它只需要很少的注释图像，在NVidia Titan GPU (6 GB)上的训练时间也非常合理，只需10小时，我们确信u-net架构可以很容易地应用于更多的任务。

Unet代码实战

对TCGA颅脑MRI进行语义分割

代码

# 导入需要的库
import os
import math
import numpy as np
import glob 
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
import random
import time
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import torchvision

kaggle_3m='../input/tcgamri/kaggle_3m/'
dirs=glob.glob(kaggle_3m+'*')
dirs
len(dirs)

data_img=[]
data_label=[]
for subdir in dirs:
    dirname=subdir.split('\\')[-1]
    for filename in os.listdir(subdir):
        img_path=subdir+'/'+filename
        if 'mask' in img_path:
            data_label.append(img_path)
        else:
            data_img.append(img_path)
data_imgx=[]
for i in range(len(data_label)):
    img_mask=data_label[i]
    img=img_mask[:-9]+'.tif'
    data_imgx.append(img)
data_imgx   

data_newimg=[]
data_newlabel=[]
for i in data_label:
    value=np.max(cv2.imread(i))
    try:
        if value>0:
            data_newlabel.append(i)
            i_img=i[:-9]+'.tif'
            data_newimg.append(i_img)
    except:
        pass

train_transformer=transforms.Compose([
    transforms.Resize((256,256)),
    transforms.ToTensor(),
])
test_transformer=transforms.Compose([
    transforms.Resize((256,256)),
    transforms.ToTensor()
])


# 创建数据集类
class BrainMRIdataset(Dataset):
    def __init__(self,img,mask,transformer):
        self.img=img
        self.mask=mask
        self.transformer=transformer
    def __getitem__(self,index):
        img=self.img[index]
        mask=self.mask[index]
        
        img_open=Image.open(img)
        img_tensor=self.transformer(img_open)
        
        mask_open=Image.open(mask)
        mask_tensor=self.transformer(mask_open)
        
        mask_tensor=torch.squeeze(mask_tensor).type(torch.long)
        
        return img_tensor,mask_tensor
    def __len__(self):
        return len(self.img)
                                                
# 划分数据集和测试集大小
s=1000
train_img=data_newimg[:s]
train_label=data_newlabel[:s]
test_img=data_newimg[s:]
test_label=data_newlabel[s:]

# 创建训练集和验证集的dataset和dataloader
train_data=BrainMRIdataset(train_img,train_label,train_transformer)
test_data=BrainMRIdataset(test_img,test_label,test_transformer)
dl_train=DataLoader(train_data,batch_size=4,shuffle=True)
dl_test=DataLoader(test_data,batch_size=4,shuffle=True)

img,lable=next(iter(dl_train))
img.shape
lable.shape

img,label=next(iter(dl_train))
plt.figure(figsize=(12,8))
for i,(img,label) in enumerate(zip(img[:4],label[:4])):
    img=img.permute(1,2,0).numpy()
    label=label.numpy()
    plt.subplot(2,4,i+1)
    plt.imshow(img)
    plt.subplot(2,4,i+5)
    plt.imshow(label)

# 下采样模块
class Downsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Downsample, self).__init__()
        self.conv_relu = nn.Sequential(
                            nn.Conv2d(in_channels, out_channels, 
                                      kernel_size=3, padding=1),
                            nn.ReLU(inplace=True),
                            nn.Conv2d(out_channels, out_channels, 
                                      kernel_size=3, padding=1),
                            nn.ReLU(inplace=True)
            )
        self.pool = nn.MaxPool2d(kernel_size=2)
    def forward(self, x, is_pool=True):
        if is_pool:
            x = self.pool(x)
        x = self.conv_relu(x)
        return x
# 上采样模块
 class Upsample(nn.Module):
    def __init__(self, channels):
        super(Upsample, self).__init__()
        self.conv_relu = nn.Sequential(
                            nn.Conv2d(2*channels, channels, 
                                      kernel_size=3, padding=1),
                            nn.ReLU(inplace=True),
                            nn.Conv2d(channels, channels,  
                                      kernel_size=3, padding=1),
                            nn.ReLU(inplace=True)
            )
        self.upconv_relu = nn.Sequential(
                               nn.ConvTranspose2d(channels, 
                                                  channels//2, 
                                                  kernel_size=3,
                                                  stride=2,
                                                  padding=1,
                                                  output_padding=1),
                               nn.ReLU(inplace=True)
            )
        
    def forward(self, x):
        x = self.conv_relu(x)
        x = self.upconv_relu(x)
        return x

# Unet模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.down1 = Downsample(3, 64)
        self.down2 = Downsample(64, 128)
        self.down3 = Downsample(128, 256)
        self.down4 = Downsample(256, 512)
        self.down5 = Downsample(512, 1024)
        
        self.up = nn.Sequential(
                               nn.ConvTranspose2d(1024, 
                                                  512, 
                                                  kernel_size=3,
                                                  stride=2,
                                                  padding=1,
                                                  output_padding=1),
                               nn.ReLU(inplace=True)
            )
        
        self.up1 = Upsample(512)
        self.up2 = Upsample(256)
        self.up3 = Upsample(128)
        
        self.conv_2 = Downsample(128, 64)
        self.last = nn.Conv2d(64, 2, kernel_size=1)

    def forward(self, x):
        x1 = self.down1(x, is_pool=False)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)
        
        x5 = self.up(x5)
        
        x5 = torch.cat([x4, x5], dim=1)           # 32*32*1024
        x5 = self.up1(x5)                         # 64*64*256)
        x5 = torch.cat([x3, x5], dim=1)           # 64*64*512  
        x5 = self.up2(x5)                         # 128*128*128
        x5 = torch.cat([x2, x5], dim=1)           # 128*128*256
        x5 = self.up3(x5)                         # 256*256*64
        x5 = torch.cat([x1, x5], dim=1)           # 256*256*128
        
        x5 = self.conv_2(x5, is_pool=False)       # 256*256*64
        
        x5 = self.last(x5)                        # 256*256*3
        return x5

model=Net()

img,label=next(iter(dl_train))
model=model.to('cuda')
img=img.to('cuda')
pred=model(img)
pred=pred
label.shape
label=label.to('cuda')

# 定义损失函数和优化器
loss_fn=nn.CrossEntropyLoss()
loss_fn(pred,label)
optimizer=torch.optim.Adam(model.parameters(),lr=0.0001)


# 定义训练方法
from tqdm import tqdm
def train_epoch(epoch, model, trainloader, testloader):
    print("-------------------------Epoch:{}-------------------------------".format(epoch+1))
    correct = 0
    total = 0
    running_loss = 0
    epoch_iou = []
    
    model.train()
    for x, y in tqdm(testloader):
 
        x, y = x.to('cuda'), y.to('cuda')
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()
            
            intersection = torch.logical_and(y, y_pred)
            union = torch.logical_or(y, y_pred)
            batch_iou = torch.sum(intersection) / torch.sum(union)
            epoch_iou.append(batch_iou.item())
            
            
            
    epoch_loss = running_loss / len(trainloader.dataset)
    epoch_acc = correct / (total*256*256)
        
        
    test_correct = 0
    test_total = 0
    test_running_loss = 0 
    epoch_test_iou = []
    
    model.eval()
    with torch.no_grad():
        for x, y in tqdm(testloader):

            x, y = x.to('cuda'), y.to('cuda')
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()
            
            intersection = torch.logical_and(y, y_pred)
            union = torch.logical_or(y, y_pred)
            batch_iou = torch.sum(intersection) / torch.sum(union)
            epoch_test_iou.append(batch_iou.item())
            
    
    epoch_test_loss = test_running_loss / len(testloader.dataset)
    epoch_test_acc = test_correct / (test_total*256*256)
    
    
    print('epoch: ', epoch, 
          'loss： ', round(epoch_loss, 3),
          'accuracy:', round(epoch_acc, 3),
          'IOU:', round(np.mean(epoch_iou), 3),
          'test_loss： ', round(epoch_test_loss, 3),
          'test_accuracy:', round(epoch_test_acc, 3),
           'test_iou:', round(np.mean(epoch_test_iou), 3)
             )
        
    return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc

# 测试方法并绘图显示
def test():
    image, mask = next(iter(dl_test))
    image=image.to('cuda')
    model.eval()
    pred_mask = model(image)
    pred_mask=pred_mask
    mask=torch.squeeze(mask)
    pred_mask=pred_mask.cpu()
    num=4
    plt.figure(figsize=(10, 10))
    for i in range(num):
        plt.subplot(num, 4, i*num+1)
        plt.imshow(image[i].permute(1,2,0).cpu().numpy())
        plt.subplot(num, 4, i*num+2)
        plt.imshow(mask[i].cpu().numpy())
        plt.subplot(num, 4, i*num+3)
        plt.imshow(torch.argmax(pred_mask[i].permute(1,2,0), axis=-1).detach().numpy())

# 训练50个epochs 
epochs = 50

# 开始训练
for epoch in range(epochs):
    train_epoch(epoch,
        model,
        dl_train,
         dl_test)
# 测试并绘图
test()

效果

在kaggle的免费GPU训练50个epochs，训练集和验证集的损失和IOU分别如下：

测试分割效果：左边第一列是原始图像，第二列是原始图像的mask，第三列是预测的mask，可以看出效果还是蛮好的！