Unet实现图像分割

目的：实现细胞分割
数据集：isbi挑战赛的数据，只有30张512x512的image和label
思路:

1. 读取数据，将数据转换成 30x512x512x1格式；
2. 由于数据太少，所以进行数据增强；(注意：对训练数据进行增强的时候image和label增强方式要一致)
3. 搭建模型。这里直接调用segmentation-models模块的Unet；(提前pip install segmentation-models)
4. 训练。推荐用Google的Colaboratory，免费GPU亲测真的好用；
5. 保存模型；
6. 预测。这里预测暂时5张图片看看效果；

数据处理+训练+保存模型 代码

#coding=utf-8
from segmentation_models import Unet
from keras.preprocessing.image import ImageDataGenerator
from PIL import Image
from keras.optimizers import *
import numpy as np
import os
import skimage.io as io
import skimage.transform as trans
import matplotlib.pyplot as plt
import random

#给定路径,起点和终点,进行图像增强,每次迭代返回一个batch_size的训练图片集和标签集 -> 这里的起点和终点对应下标都是包含在内的!
def train_image_generator(image_path,label_path,st,ed,batch_size,aug = None): 
    nowinx = st #设定初始图片
    while True:
        im_array = []
        lb_array = []
        for i in range(batch_size):
            im = Image.open(os.path.join(image_path,str(nowinx)+'.png'))
            tmp_im_array = np.array(im) #图片转numpy数组
            tmp_im_array = tmp_im_array / 255 #对数据进行归一化
            tmp_im_array = tmp_im_array[np.newaxis,:,:] #numpy数组添加一维,为了把二维图片转成三维图片集
            
            lb = Image.open(os.path.join(label_path,str(nowinx)+'.png'))
            tmp_lb_array = np.array(lb) #图片转numpy数组
            tmp_lb_array = tmp_lb_array / 255
            tmp_lb_array[tmp_lb_array > 0.5] = 1
            tmp_lb_array[tmp_lb_array <= 0.5] = 0 #对数组进行压缩并且变成01分布的数组
            tmp_lb_array = tmp_lb_array[np.newaxis,:,:] #numpy数组添加一维,为了把二维图片转成三维图片集
            
            if len(im_array) == 0:
                im_array = tmp_im_array
                lb_array = tmp_lb_array
            else:
                im_array = np.concatenate((im_array,tmp_im_array),axis=0) #将新的图片加入到之前的图片集
                lb_array = np.concatenate((lb_array,tmp_lb_array),axis=0) #将新的图片加入到之前的图片集
            
            nowinx = st if nowinx==ed else nowinx+1 #如果遍历到了超出最后一个时,就返回到第一个

        im_array = im_array[:,:,:,np.newaxis] #最后给图片集加一个通道,形成四维.
        lb_array = lb_array[:,:,:,np.newaxis]

        if (aug is not None) and (random.random() > 0.4) : #如果传入了数据增强的生成器,就进行数据增强
            new_array = im_array
            new_array = np.concatenate((new_array,lb_array),axis=3) # 把图像和标签合成一张图片进行增强
            new_array = next(aug.flow(new_array,batch_size = batch_size))
            im_array = new_array[:,:,:,0] # 将图像和标签分离
            lb_array = new_array[:,:,:,1]
            im_array = im_array[:,:,:,np.newaxis] #最后给图片集加一个通道,形成四维.
            lb_array = lb_array[:,:,:,np.newaxis]

        yield(im_array,lb_array) #按批次返回数据

aug = ImageDataGenerator( #定义一个数据增强生成器
    rotation_range = 0.05, # 定义旋转范围
    zoom_range = 0.05, # 按比例随机缩放图像尺寸
	width_shift_range = 0.05, # 图片水平偏移幅度
    height_shift_range = 0.05, # 图片竖直偏移幅度
    shear_range = 0.05, # 水平或垂直投影变换
	horizontal_flip = True, # 水平翻转图像
    fill_mode = "reflect" # 填充像素,出现在旋转或平移之后
)
if __name__ == '__main__':

    train_gen = train_image_generator("image","label",0,20,4,aug) # 获取一个训练数据生成器
    validate_gen = train_image_generator("image","label",21,29,3,None) # 获取一个验证数据生成器

    #定义并编译一个模型
    model = Unet('resnet34', input_shape = (512, 512, 1), encoder_weights = None) #1代表通道数
    model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])

    #进行模型的训练,并用his来记录训练过程中的参数变化,方便最后生成图像
    his = model.fit_generator(
        generator = train_gen, #训练集生成器
        steps_per_epoch = 300, #训练集每次epoch的数量
        validation_data = validate_gen, #验证集生成器
        validation_steps = 3, #验证集每次epoch的数量
        #validation_split = 0.2
        epochs = 10 #进行epoch的次数
    )
    model.save("model_v4.h5") #保存模型
    print("Saved model to disk")
    print("done!!!!!")

注意：
读取图像文件时，PIL和Opencv都可以，但是保存的shape不一样，如下：

预测代码

from keras.models import load_model
from PIL import Image
import numpy as np
import os
import skimage.io as io
#给定路径,起点和终点,进行图像增强,每次迭代返回一个batch_size的测试图片集和标签集 -> 这里的起点和终点对应下标都是包含在内的!
def test_image_generator(image_path,st,ed,batch_size): 
    nowinx = st #设定初始图片
    while True:
        im_array = []
        for i in range(batch_size):
            im = Image.open(os.path.join(image_path,str(nowinx)+'.png'))
            tmp_im_array = np.array(im) #图片转numpy数组
            tmp_im_array = tmp_im_array / 255 #对数据进行归一化
            tmp_im_array = tmp_im_array[np.newaxis,:,:] #numpy数组添加一维,为了把二维图片转成三维图片集

            if len(im_array) == 0:
                im_array = tmp_im_array
            else:
                im_array = np.concatenate((im_array,tmp_im_array),axis=0) #将新的图片加入到之前的图片集
            
            nowinx = st if nowinx==ed else nowinx+1 #如果遍历到了超出最后一个时,就返回到第一个

        im_array = im_array[:,:,:,np.newaxis] #最后给图片集加一个通道,形成四维.

        yield(im_array) #按批次返回数据

def saveResult(save_path,npyfile):
    for i,item in enumerate(npyfile): #取出预测结果中的每一个
        im_np = item[:,:,0]
        im_np[im_np > 0.5] = 1
        im_np[im_np <= 0.5] = 0
        io.imsave(os.path.join(save_path,"%d_predict_v4_ly.png"%i),im_np)

model = load_model("model_v4.h5") #加载模型

test_gen = test_image_generator("test",0,4,1) # 获取一个测试数据生成器,一次生成一个
res = model.predict_generator(test_gen,5) # 进行结果预测,5代表预测的个数
saveResult("res",res)

效果图：左边是label，右边是分割结果

补充一下：Google的免费GPU真的很好用，用法总结：

1. 注册谷歌邮箱，然后将整个项目文件上传到谷歌云盘；
2. 右键项目文件夹，点击更多打开方式，搜索谷歌应用Colaboratory，并关联；
3. 关联时，有个Colaboratory的主页网址，点击进入；
4. 将Colaboratory挂载到云盘；
5. 安装必须的包：
   pip install segmentation-models
6. 移动目录
   cd /content/drive/My Drive/Colab Notebooks/your directory
7. #运行代码
   ！python xxx.py

参考文章：
https://blog.csdn.net/y_marcel_h/article/details/90220624
https://www.jianshu.com/p/000d2a9d36a0