U-net语义分割

最近参加大学生设计大赛，我们选的是大类是大数据应用类，然后小类是生物与医疗大数据

为什么要做U-net语义分割

随着医学图像可视化技术的发展和各种医学成像模式的出现，医学图像自动分析和处理已成为图像工程领域和生物医学工程领域一个重要的研究方向。作为医学图象处理中的一个热点问题，细胞图像的自动分析和识别一直受到人们的普遍重视。目前的图像诊断系统，大多数已采用形态学、灰度特征和色度学，并结合专家系统，对人体细胞进行分析和诊断。而由于我们缺少人体内部细胞的数据集，所有采用了果蝇-龄幼虫腹侧神经索的连续切片透射电镜作为数据集来代替人体内部细胞。本系统针对实际诊断中，识别某些病理细胞需要经验丰富的医生且耗费大量时间这一问题，提出了一种基于深度学习的语义分割管理系统，该系统从全局到局部，分割出医生想要的部分，辅助医生进行更快捷速,可重复,更一致的图像分析诊断，减轻科研人员和临床医生的重复性劳动，减少诊断的主观性和差异性，防止由于劳累产生的误诊，为患者的个性化治疗提供可能。

代码

main.py

from model import *
from data import *


# os.environ["CUDA_VISIBLE_DEVICES"] = "0"

#定义了一个字典，参数是数据增强的相关参数
data_gen_args = dict(rotation_range=0.2,  # 旋转
                     width_shift_range=0.05,  # 宽度变化
                     height_shift_range=0.05,  # 高度变化
                     shear_range=0.05,  # 错切变换
                     zoom_range=0.05,  # 缩放
                     horizontal_flip=True,  # 水平翻转
                     fill_mode='nearest')  # 填充模式
#生成训练所需要的图片和标签，trainGenerator函数在data.py
myGene = trainGenerator(2, 'data/membrane/train', 'image', 'label', data_gen_args, save_to_dir=None)

model_use = unet()#初始化unet模型
#保存模型
model_checkpoint = ModelCheckpoint('unet_membrane.hdf5', monitor='loss', verbose=1, save_best_only=True)
# 在每个epoch后保存模型到filepath
# filename:保存模型的路径
# mointor:需要监视的值
# verbose:表示信息展示模式，1展示，0不展示
# save_best_only:保存在验证集上性能最好的模型


model_use.fit_generator(myGene, steps_per_epoch=300, epochs=1, callbacks=[model_checkpoint])
# 训练函数
# generator:生成器(image,mask)
# steps_per_epoch:训练steps_per_epoch个数据时记一个epoch结束
# epoch:数据迭代轮数
# callbacks:回调函数

#测试集数据增强生成，testGenerator位于data.py文件
testGene = testGenerator("data/membrane/test")
# 为测试图片生成预测
results = model_use.predict_generator(testGene, 30, verbose=1)
# generator:生成器
# steps:在声明一个epoch完成,并开始下一个epoch之前从生成器产生的总步数
# verbose:信息展示模式，1展示，0不展示

#保存测试结果
saveResult("data/membrane/test", results)

model.py

import numpy as np
import os
import skimage.io as io
import skimage.transform as trans
import numpy as np
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
from tensorflow.keras.optimizers import *
from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler
from tensorflow.keras import backend as keras


def unet(pretrained_weights=None, input_size=(256, 256, 1)):
    inputs = Input(input_size)  # 初始化keras张量
    
    #第一层卷积
    #实际上从unet的结构来看每一次卷积的padding应该是valid，也就是每次卷积后图片尺寸减少2，
    #但在这里为了避免裁剪，方便拼接，把padding设成了same，即每次卷积不会改变图片的尺寸。
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
    # filters:输出的维度
    # kernel_size:卷积核的尺寸
    # activation:激活函数
    # padding:边缘填充，实际上在该实验中并没有严格按照unet网络结构进行卷积，same填充在卷积完毕之后图片大小并不会改变
    # kernel_initializer:kernel权值初始化
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)#采用2*2的最大池化
    
    #第二层卷积
    #参数类似于第一层卷积，只是输出的通道数翻倍
    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    
    #第三层卷积
    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
    
    #第四层卷积
    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
    drop4 = Dropout(0.5)(conv4)  # 每次训练时随机忽略50%的神经元，减少过拟合
    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
    
    #第五层卷积
    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4)
    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
    drop5 = Dropout(0.5)(conv5)# 每次训练时随机忽略50%的神经元，减少过拟合
    
    #第一次反卷积
    up6 = Conv2D(512, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
        UpSampling2D(size=(2, 2))(drop5))  # 先上采样放大，在进行卷积操作，相当于转置卷积
    # merge6 = merge([drop4, up6], mode='concat', concat_axis=3)
    #将第四层卷积完毕并进行Dropout操作后的结果drop4与反卷积后的up6进行拼接
    merge6 = concatenate([drop4, up6], axis=3)  # （width,heigth,channels）拼接通道数
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge6)
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)

    #第二次反卷积
    up7 = Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
        UpSampling2D(size=(2, 2))(conv6))
    # merge7 = merge([conv3, up7], mode='concat', concat_axis=3)
    #将第三层卷积完毕后的结果conv3与反卷积后的up7进行拼接
    merge7 = concatenate([conv3, up7], axis=3)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge7)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)
    
    #第三次反卷积
    up8 = Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
        UpSampling2D(size=(2, 2))(conv7))
    # merge8 = merge([conv2, up8], mode='concat', concat_axis=3)
    #将第二层卷积完毕后的结果conv2与反卷积后的up8进行拼接
    merge8 = concatenate([conv2, up8], axis=3)#拼接通道数
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge8)
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)
    
    
    #第四次反卷积
    up9 = Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
        UpSampling2D(size=(2, 2))(conv8))
    # merge9 = merge([conv1, up9], mode='concat', concat_axis=3)
    #将第一层卷积完毕后的结果conv1与反卷积后的up9进行拼接
    merge9 = concatenate([conv1, up9], axis=3)#拼接通道数
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    conv9 = Conv2D(2, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    
    #进行一次卷积核为1*1的卷积操作，卷积完毕后通道数变为1，作为输出结果
    conv10 = Conv2D(1, 1, activation='sigmoid')(conv9)

    model = Model(inputs=inputs, outputs=conv10)
    #keras内置函数，对模型进行编译
    model.compile(optimizer=Adam(lr=1e-4), loss='binary_crossentropy', metrics=['accuracy'])
    # optimizer:优化器
    # binary_crossentropy:与sigmoid相对应的损失函数
    # metrics:评估模型在训练和测试时的性能的指标

    if pretrained_weights:
        model.load_weights(pretrained_weights)

    return model

data.py

from __future__ import print_function
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import numpy as np
import os
import glob
import skimage.io as io
import skimage.transform as trans


def adjustData(img, mask, flag_multi_class, num_class):  # 将图片归一化
    if np.max(img) > 1:
        img = img / 255.0
        mask = mask / 255#将像素值映射到0-1范围内
        mask[mask > 0.5] = 1  # 将mask二值化
        mask[mask <= 0.5] = 0
    return img, mask


# 生成训练所需要的图片和标签
def trainGenerator(batch_size, train_path, image_folder, mask_folder, aug_dict, image_color_mode="grayscale",
                   mask_color_mode="grayscale", image_save_prefix="image", mask_save_prefix="mask",
                   flag_multi_class=False, num_class=2, save_to_dir=None, target_size=(256, 256), seed=1):
    # 图片生成器对数据进行增强，扩充数据集大小，增强模型的泛化能力。比如进行旋转，变形，归一化等等
    #参数就是main里面定义的字典
    image_datagen = ImageDataGenerator(**aug_dict)  
    #对image和mask分别定义图片生成器
    mask_datagen = ImageDataGenerator(**aug_dict)
    
    #image生成器
    #flow_from_directory函数用于根据文件路径，增强数据
    image_generator = image_datagen.flow_from_directory(  
        train_path,#目标文件夹路径，从该路径提取图片来生成数据增强后的图片
        classes=[image_folder],#子文件夹列表，对于image生成器来说也就是train下的image对应的30张图片
        class_mode=None,#确定返回的标签数组的类型，none表示不返回标签数组的类型，只返回标签
        color_mode=image_color_mode,#颜色模式,grayscale为单通道，也就是灰度图像
        target_size=target_size,#目标图像的尺寸，缺省定义为256*256
        batch_size=batch_size,#表示每批数据的大小
        save_to_dir=save_to_dir,#表示保存图片，缺省值为none，也就是不保存数据增强后的图片
        save_prefix=image_save_prefix,#保存提升后图片时使用的前缀, 仅当设置了save_to_dir时生效
        seed=seed)#表示随机数种子，缺省值为true，即表示每次数据增强都要打乱数据


    #mask生成器（label数据增强的结果）
    #参数与image生成器的参数一致，只需要把mask对应的值赋值给这些参数即可。
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes=[mask_folder],
        class_mode=None,
        color_mode=mask_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=mask_save_prefix,
        seed=seed)
    
    # 将image和mask打包成元组的列表[(image1,mask1),(image2,mask2),...]
    train_generator = zip(image_generator, mask_generator)  
    for (img, mask) in train_generator:
        #数据调整函数，adjustData函数在data.py文件中
        img, mask = adjustData(img, mask, flag_multi_class, num_class)
        yield (img, mask)

#测试集数据生成
def testGenerator(test_path, num_image=30, target_size=(256, 256), flag_multi_class=False, as_gray=True):
    for i in range(num_image):
        #读取test原图
        img = io.imread(os.path.join(test_path, "%d.png" % i), as_gray=as_gray)
        img = img / 255.0  # 归一化
        img = trans.resize(img, target_size)
        img = np.reshape(img, img.shape + (1,)) if (not flag_multi_class) else img
        img = np.reshape(img, (1,) + img.shape)  # (1,width,height)
        yield img


def geneTrainNpy(image_path, mask_path, flag_multi_class=False, num_class=2, image_prefix="image", mask_prefix="mask",
                 image_as_gray=True, mask_as_gray=True):
    image_name_arr = glob.glob(os.path.join(image_path, "%s*.png" % image_prefix))
    image_arr = []
    mask_arr = []
    for index, item in enumerate(image_name_arr):
        img = io.imread(item, as_gray=image_as_gray)
        img = np.reshape(img, img.shape + (1,)) if image_as_gray else img
        mask = io.imread(item.replace(image_path, mask_path).replace(image_prefix, mask_prefix), as_gray=mask_as_gray)
        mask = np.reshape(mask, mask.shape + (1,)) if mask_as_gray else mask
        img, mask = adjustData(img, mask, flag_multi_class, num_class)
        image_arr.append(img)
        mask_arr.append(mask)
    image_arr = np.array(image_arr)
    mask_arr = np.array(mask_arr)
    return image_arr, mask_arr


def labelVisualize(num_class, color_dict, img):
    img = img[:, :, 0] if len(img.shape) == 3 else img
    img_out = np.zeros(img.shape + (3,))
    for i in range(num_class):
        img_out[img == i, :] = color_dict[i]
    return img_out / 255.0


def saveResult(save_path, npyfile, flag_multi_class=False, num_class=2):
    for i, item in enumerate(npyfile):
        img = labelVisualize(num_class, COLOR_DICT, item) if flag_multi_class else item[:, :, 0]
        io.imsave(os.path.join(save_path, "%d_predict.png" % i), img)

这个是数据集

这个是训练后的二值图

训练过程

accuracy的值不是非常高，后期会继续完善的。