一、网络模型

常见分割网络

1.U-Net 3d

// U-Net3d
class UNet3D(nn.Module):
    def __init__(self, params):
        super(UNet3D, self).__init__()
        self.params = params
        self.in_chns = self.params['in_chns']
        self.ft_chns = self.params['feature_chns']
        self.n_class = self.params['class_num']
        self.trilinear = self.params['trilinear']
        self.dropout = self.params['dropout']
        assert (len(self.ft_chns) == 5 or len(self.ft_chns) == 4)

        self.in_conv = ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
        self.down1 = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])
        self.down2 = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])
        self.down3 = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])
        if (len(self.ft_chns) == 5):
            self.down4 = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])
            self.up1 = UpBlock(self.ft_chns[4], self.ft_chns[3], self.ft_chns[3],
                               dropout_p=0.0, trilinear=self.trilinear)
        self.up2 = UpBlock(self.ft_chns[3], self.ft_chns[2], self.ft_chns[2],
                           dropout_p=0.0, trilinear=self.trilinear)
        self.up3 = UpBlock(self.ft_chns[2], self.ft_chns[1], self.ft_chns[1],
                           dropout_p=0.0, trilinear=self.trilinear)
        self.up4 = UpBlock(self.ft_chns[1], self.ft_chns[0], self.ft_chns[0],
                           dropout_p=0.0, trilinear=self.trilinear)

        self.out_conv = nn.Conv3d(self.ft_chns[0], self.n_class,
                                  kernel_size=3, padding=1) #bias = -2.19
        if self.params['activation'] == 'sigmoid':
            self.activation =  nn.Sigmoid()
        elif self.params['activation'] == 'softmax':
            self.activation = nn.Softmax()
        else:
            self.activation = 'None'

    def forward(self, x):
        x0 = self.in_conv(x)
        x1 = self.down1(x0)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        if (len(self.ft_chns) == 5):
            x4 = self.down4(x3)
            x = self.up1(x4, x3)
        else:
            x = x3
        x = self.up2(x, x2)
        x = self.up3(x, x1)
        x = self.up4(x, x0)
        output = self.out_conv(x)

        if self.activation == 'None':
            return output
        output_sigmoid = self.activation(output)

        return output_sigmoid

2.U-Net 2d

// parts of unet
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F

""" Parts of the U-Net model """
"""https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_parts.py"""


class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        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)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

// full unet
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from model.parts_unet import *

""" Full assembly of the parts to form the complete network """
"""Refer https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py"""


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(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)
        logits = self.outc(x)
        return logits


if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)

二、损失函数

1.Focal loss

def FocalLoss(predict, soft_y, softmax):
    alpha_num = 0.01
    gamma = 2
    with torch.no_grad():
        alpha = torch.empty_like(predict).fill_(1 - alpha_num)
        alpha[soft_y == 1] = alpha_num

    crit = nn.BCELoss(reduction='none')
    ce_loss = crit(predict.float(), soft_y.float())
    pt = torch.exp(-ce_loss)

    loss = (alpha * torch.pow(1 - pt, gamma) * ce_loss)

    loss = loss.mean()

    return loss

2.Hausdoff loss

def hd_loss(seg, gt, seg_dtm, gt_dtm):
    """
    compute Hausdorff distance loss for binary segmentation based on distance transform
    :param seg: seg results, shape=(b,c,d,h,w) ,for binary c = 1
    :param gt: ground truth, shape=(b,c,d,h,w), for binary c = 1
    :param seg_dtm: segmentation distance transform, shape=(b,c,d,h,w)
    :param gt_dtm: ground truth distance transform, shape=(b,c,d,h,w)
    :return: boundary Hausdorff distance
    """
    delta = (seg - gt) ** 2

    seg_dtm_alpha = seg_dtm ** 2
    gt_dtm_alpha = gt_dtm ** 2

    dtm = seg_dtm_alpha + gt_dtm_alpha
    multiple_d = torch.einsum('bcxyz, bcxyz->bcxyz', dtm, delta)
    hd_loss_value = multiple_d.mean()

    return hd_loss_value

3.Dice loss

def dice_loss(predict, soft_y, softmax=False):

    smooth = 1e-5
    num = predict.size(0)
    p_vol = predict.view(num, -1)
    y_vol = soft_y.view(num, -1)
    intersection = (p_vol * y_vol).sum(1)

    dice_score = (2. * intersection + smooth) / (p_vol.sum(1) + y_vol.sum(1) + smooth)
    return 1 - dice_score.sum() / num

三、训练一个U-Net

1.数据加载
2.模型选择，这里为U-Net
3.损失函数
4.训练
5.预测

//dataset
# !/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
from torch.utils.data import Dataset
import glob
import cv2
import random
import os



'''data loader'''
# inherit class Dataset():
class DataLoader(Dataset):

    def __init__(self, data_path):
        # initialize file paths or a list of file names
        self.data_path = data_path
        self.imgs_path = glob.glob(os.path.join(data_path, 'image/*png'))

    def augment(self, image, flipCode):
        # horizontal flip : flipCode = 1; vertical:=0; h & v : =-1
        flip = cv2.flip(image, flipCode)
        return flip

    def __getitem__(self, index):
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        image_path = self.imgs_path[index]
        label_path = image_path.replace('image', 'label')

        image = cv2.imread(image_path)
        label = cv2.imread(label_path)

        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])

        if label.max() > 1:
            label = label / 255

        # random flip the image & label if flipCode != 2
        flipCode = random.chioce([-1, 0, 1, 2])
        if flipCode != 2:
            image = self.augment(image, flipCode)
            label = self.augment(label, flipCode)
        return image, label

    def __len__(self):
        # get the len of train set
        return len(self.imgs_path)

if __name__ == '__main__':
    liunj_dataset = DataLoader('data/train/')
    print('number of train data:', len(liunj_dataset))

    train_loader = torch.utils.data.DataLoader(dataset=DataLoader,
                                           batch_size=2,
                                           shuffle=True)
    for image, label in train_loader:
        print(image.shape)

//train_model
# -*- coding: utf-8 -*-
from model.full_unet import UNet
from dataset import DataLoader
from torch import optim
import torch.nn as nn
import torch
import time
import matplotlib.pyplot as plt


def train(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    start = time.time()

    liunj_dataset = DataLoader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=liunj_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    # define the optimizer
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8,momentum=0.9)

    # define the loss function
    criterion = nn.BCEWithLogitsLoss()


    # initial the loss +wuqiong
    best_loss = float('inf')


    for epoch in range(epochs):
        print('Epoch {}/{}'.format(epoch + 1, epochs))
        print('-' * 50)

        # train model
        net.train()

        for image, label in train_loader:

            optimizer.zero_grad()

            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)

            pred = net(image)

            loss = criterion(pred, label)
            print('Loss/train', loss.item())

            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')

            loss.backward()
            optimizer.step()



        #     if step % 10 == 0:
        #         print('Epoch {}/{} | Current step: {} | Loss: {} | Acc: {} | AllocMem (Mb): {}' \
        #               .format(epoch + 1, epochs, step, loss, acc, torch.cuda.memory_allocated()/1024/1024)
        #               )
        #     # current step: {} / {}
        #     epoch_loss = running_loss / len(dataloader.dataset)
        #     epoch_acc = running_acc / len(dataloader.dataset)
        #
        #     train_loss.append(epoch_loss) if phase == 'Train' else valid_loss.append(epoch_loss)
        #
        # time_elapsed = time.time() - start
        # print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))




# def acc_metric(predb, yb):
#     return (predb.argmax(dim=1) == yb.cuba()).float().mean()


if __name__ == "__main_":

    device = torch.device('cuda' if torch.cuda.is_availabel() else 'cpu')

    net = UNet(n_channels=3, n_classes=3)
    net.to(device=device)

    data_path = 'data/train/'

    train(net, device, data_path)
	# 此段存在问题，后续可以使用tensorboard可视化训练过程
    plt.figure(figsize=(12, 4))
    plt.subplot(121)
    plt.plot(loss[:])
    plt.title("train_loss")
    plt.subplot(122)
    plt.plot(train_epochs_loss[1:], '-o', label="train_loss")
    plt.plot(valid_epochs_loss[1:], '-o', label="valid_loss")
    plt.title("epochs_loss")
    plt.legend()
    plt.show()

# -*- coding: utf-8 -*-
import glob
import numpy as np
import torch
import cv2
from model.full_unet import UNet


if __name__ == '__main__':

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = UNet(n_channels=3, n_classes=2)
    net.to(device=device)

    net.load_state_dict(torch.load('best_model.pth', map_location=device))

    # test model
    net.eval()

    # get all the images: '../../...png'
    test_path = glob.glob('data/test/*.png')

    for test_path in test_path:

        save_res_path = test_path.split('.')[0] + '_res.png'

        img = cv2.imread(test_path)
        img = cv2.reshape(1, 1, img.shape[0], img.shape[1])
        # change to tensor
        img_tensor = torch.from_numpy(img)
        # copy the img_tensor to device
        img_tensor = img_tensor.to(device=device, dtype=torch.float32)

        pred = net(img_tensor)

        pred = np.array(pred.data.cpu()[0][0])

        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0

        cv2.imwrite(save_res_path, pred)