tensorflow版使用uNet进行医学图像分割(Skin数据集)

程序员文章站 2022-04-08 08:41:11

...

tensorflow版使用uNet进行医学图像分割(Skin数据集)

计算机视觉、深度学习uNet医学图像分割实验笔记

文章目录

tensorflow版使用uNet进行医学图像分割(Skin数据集)
1. 数据集处理
2. 定义uNet网络模型
3. 主函数
4. 预测
5. 结果展示

1. 数据集处理

在这个实例中用到了Skin皮肤数据集，数据集的百度网盘下载地址： skin皮肤数据集提取码: vb9e
数据集的格式如图：
tensorflow版使用uNet进行医学图像分割(Skin数据集)

1.1 加载数据集并保存

分别加载训练集和测试集并保存为jpg格式：

import numpy as np
import scipy.misc
# 加载训练集，以jpg格式保存到train文件夹下
imgs_train = np.load('/skin_dataset/data_train.npy')  # 读入.npy文件
print(imgs_train.shape)
for i in range(imgs_train.shape[0]):
    B = imgs_train[i, :, :, 0]
    scipy.misc.imsave('./train/' + str(i) + ".jpg", B)  # 保存为jpg格式，也可将jpg换为png等其他格式

保存结果如图：

tensorflow版使用uNet进行医学图像分割(Skin数据集)

1.2 修改图像格式

查看图像格式后发现image的位深度为24，而label的位深度为8，所以需要将将其变为同一值

import numpy as np
from PIL import Image
import os

path = './SKIN_data/train/labels/'	# 需要改变格式的图像路径
newpath = './trainLabels/'	# 改变格式后的存储路径


def turnto24(path):
    files = os.listdir(path)
    files = np.sort(files)
    i = 0
    for f in files:
        imgpath = path + f
        img = Image.open(imgpath).convert('RGB')
        dirpath = newpath
        file_name, file_extend = os.path.splitext(f)
        dst = os.path.join(os.path.abspath(dirpath), file_name + '.jpg')
        img.save(dst)


turnto24(path)

1.3 创建csv文件来保存训练集和标签的路径

我在实验中用了csv来指定训练集和标签的一对一关系，所以需要把训练集的路径和标签的路径保存在csv文件里：

# coding:utf-8
import os
import csv

# 创建一个csv文件保存训练集和其标签的路径
def create_csv():
    path = './SKIN_data/train/images/'  # 修改为你自己的图片路径即可
    name = os.listdir(path)
    print(name)
    with open('train.csv', 'w') as csvfile:		# 保存为train.cvs文件
        writer = csv.writer(csvfile)
        # writer.writerow(['images_path', 'labels_path'])	# 先添加一行为images_path, labels_path
        for n in name:		# 循环添加路径
            if n[-4:] == '.jpg':
                print(n)
                writer.writerow(
                    ['./SKIN_data/train/images/' + str(n), './SKIN_data/train/labels/' + str(n[:-4] + '.jpg')])
            else:
                pass


if __name__ == "__main__":
    create_csv()

保存后的train.csv和test.csv文件内容：

1.4 定义读取数据的函数

现在还是只是图像的路径，训练器需要的是已经读取的数据，所以需要遍历刚刚存储的csv文件，这里利用cv_imread函数进行读取：

import cv2

# 创建一个读取图像的函数 cv_imread
def cv_imread(filePath):
    cv_img = cv2.imdecode(np.fromfile(filePath, dtype=np.uint8), -1)
    return cv_img

最后打包到image_data_array和label_data_array并返回：

import csv
import json
import os
import numpy as np
import cv2

# 创建加载训练数据的函数
# def load_train(csvDir, width, height, batch_size):
def load_train(csvDir, width, height):
    fx = 0.0
    fy = 0.0
    # 处理得到的数组
    images_path = []
    labels_path = []
    # 利用csv.reader读取csv文件，然后将返回的值转化为列表
    # 然后就可以得到x(训练集), y(标签)的地址
    csvFile = open(csvDir, "r")
    reader = csv.reader(csvFile)
    content = list(reader)
    for item in content:
        images_path.append(item[0])
        labels_path.append(item[1])
    # 进入循环读取图像
    # while True:
    # 下面定义两个数组来装每个批次(batch_size)的数据
    image_data_array = []
    label_data_array = []
    # 随机选一组数据
    # index_group = np.random.randint(0, len(images_path), batch_size)
    index_group = len(images_path)
    # print("batch_size:", str(index_group))
    for index in range(index_group):
        image = images_path[index]
        label = labels_path[index]

        image_data = cv_imread(image)
        # 这里需要resize一下图片的长宽，让长宽与模型接收的长宽一致
        image_data = cv2.resize(image_data, (width, height), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
        image_data = image_data.astype(np.float32)
        image_data = np.multiply(image_data, 1.0 / 255.0)
        image_data_array.append(image_data)

        label_data = cv_imread(label)
        # label_data = cv2.cvtColor(label_data, cv2.COLOR_GRAY2BGR) # 颜色转化
        label_data = cv2.resize(label_data, (width, height), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
        label_data = label_data.astype(np.float32)
        label_data = np.multiply(label_data, 1.0 / 255.0)
        # # 做一个阈值处理，输出的概率值大于0.5的就认为是对象，否则认为是背景
        # label_data[label_data > 0.5] = 1
        # label_data[label_data <= 0.5] = 0
        label_data_array.append(label_data)

    image_data_r = np.array(image_data_array)
    label_data_r = np.array(label_data_array)

    return image_data_r, label_data_r

和上面的思路一样，我们还需要写一个加载测试集的函数（当然这一步可以不要，但是这样你就看不到val_loss和val_accuracy两个数据了）：

# 创建加载测试集的函数
# def load_test(csvDir, width, height, batch_size):
def load_test(csvDir, width, height):
    fx = 0.0
    fy = 0.0
    # 处理列表得到数组
    images_path = []
    labels_path = []
    csvFile = open(csvDir, "r")
    reader = csv.reader(csvFile)
    content = list(reader)
    for item in content:
        images_path.append(item[0])
        labels_path.append(item[1])
    # 进入循环读取照片

    # for image, label in zip(images_path, labels_path):
    image_data_array = []
    label_data_array = []
    # index_group = np.random.randint(0, len(images_path), batch_size)
    index_group = len(images_path)
    # print("batch_size:", str(index_group))
    for index in range(index_group):
        image = images_path[index]
        label = labels_path[index]

        image_data = cv_imread(image)
        image_data = cv2.resize(image_data, (width, height), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
        image_data = image_data.astype(np.float32)
        image_data = np.multiply(image_data, 1.0 / 255.0)
        image_data_array.append(image_data)

        label_data = cv_imread(label)
        # label_data = cv2.cvtColor(label_data, cv2.COLOR_GRAY2BGR) # 颜色转化
        label_data = cv2.resize(label_data, (width, height), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
        label_data = label_data.astype(np.float32)
        label_data = np.multiply(label_data, 1.0 / 255.0)
        label_data_array.append(label_data)

    image_data_r = np.array(image_data_array)
    label_data_r = np.array(label_data_array)

    return image_data_r, label_data_r

2. 定义uNet网络模型

创建一个文件(model.py)用来写uNet网络模型

from keras.models import *
from keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Dropout, concatenate
from keras.optimizers import *


def unet(pretrained_weights=None, input_size=(256, 256, 3)):
    inputs = Input(input_size)
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
    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)  # 将部分隐藏层神经元丢弃，防止过于细化而引起的过拟合情况
    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)

    up6 = Conv2D(512, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
        UpSampling2D(size=(2, 2))(drop5))
    merge6 = concatenate([drop4, up6], axis=3)  # axis=1代表列合并, axis=2代表行合并，axis=3代表层合并
    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 = 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 = 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 = 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)
    conv10 = Conv2D(3, 1, activation='sigmoid')(conv9)

    print(inputs.shape, conv10.shape)
    model = Model(inputs=inputs, outputs=conv10)

    model.compile(optimizer=Adam(lr=1e-6), loss='binary_crossentropy', metrics=['accuracy'])


    # model.summary()

    if (pretrained_weights):
        model.load_weights(pretrained_weights)

    return model

数据的读入和uNet网络模型解决后，接下来就是写主函数跑数据了。

3. 主函数

创建一个文件（train.py）用来先主函数跑数据

from model import *
import matplotlib.pyplot as plt
from keras.callbacks import ModelCheckpoint
from myGenerator import load_train, load_test

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = "-1"  # 选择哪一块gpu，-1就是调用cpu，0就是调用gpu
config = tf.ConfigProto()  # 对session进行参数配置
config.allow_soft_placement = True  # 如果你指定的设备不存在，允许TF自动分配设备
config.gpu_options.per_process_gpu_memory_fraction = 0.7  # 分配百分之七十的显存给程序使用，避免内存溢出，可以自己调整
config.gpu_options.allow_growth = True  # 按需分配显存，这个比较重要

session = tf.Session(config=config)

model = unet()
# 这里调用load_train(csvDir, width, height, batch_size)产生数据
image_data_r, label_data_r = load_train('./train.csv', 256, 256)
print(image_data_r.shape)
print(label_data_r.shape)
model_checkpoint = ModelCheckpoint('uNet_Skin.h5', monitor='loss', verbose=1, save_best_only=True)

print('Fitting Model....')
# 如果内存小batch_size就设为1吧
model = model.fit(image_data_r, label_data_r, batch_size=4, epochs=30, verbose=1, validation_split=0.2, shuffle=True,
                  callbacks=[model_checkpoint])

# 展示一下精确度随训练的变化图
plt.plot(model.history['acc'])
plt.plot(model.history['val_acc'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

# 展示一下loss随训练的变化图
plt.plot(model.history['loss'])
plt.plot(model.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()