基于U-net的肝脏肿瘤分割实战（Pytorch实现）

这是我去年本科毕业设计的一部分，去年使用tensorflow去实现的，这几天就随手用Pytorch做了一下实现。肝脏肿瘤的分割是医学影像分割里面比较重要的一部分，实现从人的CT或MRI影像中将肝脏区域肿瘤给分割出来，对临床治疗还是有一定的实际意义。在这里网络还是使用最简单的U-net网络，U-net作为医学影像分割中的最经典的方法，其Encoder和Decoder的结构以及对称式的设计对后面的很多网络都有很大的影响。

网络结构

跳跃连接部分就是讲encoder的特征图和decoder的特征图在通道纬度上做拼接。

数据集介绍

因为医学影像的特殊性，大多数人的数据集都不会公开，这里介绍一个公开数据集：3D-IRCADB，里面包含有20个病人的CT影像序列，每个病人包含大概几百张吧，记不清楚了，文件是以DICOM格式存储的，mask也是用的DICOM格式存储。我不是很习惯，所以我先将图像都处理成了png格式。长这样：
image：

label：

转换的过程会用到一个第三方库，不详细介绍，会重新写一篇博客。

整个训练过程以及网络结构跟我的前两片博客完全一致，就是dataloader修改了一下，这里不详细介绍了，完整的代码可以访问我的Github找到，用了两千张图像训练了80个epoch得到的结果从肉眼上看起来很不错。

Dataset.py

# !/usr/bin/python3
# -*- coding:utf-8 -*-
# Author:WeiFeng Liu
# @Time: 2021/12/9 下午1:28
"""
使用的是视网膜血管分割的数据集，训练集就二十张图像
"""
import torch
import cv2
import os
import glob
from torch.utils.data import Dataset
import random

class _3Dircadb_Dataset(Dataset):
    def __init__(self,data_path):
        self.data_path = data_path
        self.image_path = glob.glob(os.path.join(data_path,'image/*.png'))
        self.label_path = glob.glob(os.path.join(data_path,'label/*.png'))
    def augment(self,image,mode):
        """
        :param image:
        :param mode: 1 :水平翻转 0 : 垂直翻转 -1 水平+垂直翻转
        :return:
        """
        file = cv2.flip(image,mode)
        return file
    def __len__(self):
        return len(self.image_path)

    def __getitem__(self,index):
        image_path = self.image_path[index]
        label_path = self.label_path[index]

        #读取
        image = cv2.imread(image_path)
        label = cv2.imread(label_path)

        #转为灰度图
        image = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label,cv2.COLOR_BGR2GRAY)

        # 随机进行数据增强,2时不做数据增强
        mode = random.choice([-1, 0, 1, 2])
        if mode != 2:
            image = self.augment(image, mode)
            label = self.augment(label, mode)

        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])
        # 标签二值化 ，将255 -> 1
        label = label / 255
        return image, label

# isbi = _3Dircadb_Dataset(r"Dataset/train/")
# print(len(isbi))
# train_loader = torch.utils.data.DataLoader(isbi,
#                                            batch_size=2,
#                                            shuffle=True)
# for image ,label in train_loader:
#     print(image.shape)

test

image

label

segment result