全卷积神经网络FCN模型Unet 图像分割项目实践

前言

本文基于下面链接的项目, 实践一次基于Unet模型的图像分割.

个人在实现时遇到很多问题, 许多问题在网上并没有找到解答, 所以写下本文, 记录了在实现过程中遇到的问题和大部分代码的注解, 主要从代码逻辑入手, 分析整个实践过程.

我的实现代码放在文章最后, 供大家参考

参考链接

效果展示

一些声明

• 本文实现的代码仅考虑灰度图, 即单通道图

• 本文中涉及到的数据增强, Unet模型部分, 均为直接调用API, 只讲解具体怎么使用, 涉及具体原理部分需要自己在网上查询

• 我的实现写成了多个版本, 这里讲解第一个和第四个版本

• 代码运行在Google Colab上, 这里放一下网上别人的使用教程 ( 使用Colab需要科学上网, 这样比较稳 )

  Colab使用参考链接

整体概述

整个项目需要实现的是对图像的分割, 为二分类问题, 因此我们主要的代码逻辑如下 :

1. 读入数据集和标签, 进行数据增强(也可不进行数据增强)
2. 预编译一个loss函数为binary_crossentropy的Unet模型
3. 使用数据和标签对Unet进行训练, 将训练好的模型进行保存
4. 加载模型进行预测和图片生成

读入数据

首先明确一下, 我们的图像和标签都是 512 x 512 的图片.

我们首先将图片转换成numpy数组并将多张图片合到一起, 此时变成了三维的数组 : n x 512 x 512

def load_png_files(image_path,start,end): # 加载图片成numpy的list
    im_array = []
    for i in range(start,end+1):
        im = Image.open(os.path.join(image_path,str(i)+'.png'))
        tmp_array = np.array(im) #图片转numpy数组
        tmp_array = tmp_array[np.newaxis,:,:] #numpy数组添加一维,为了把二维图片转成三维图片集
        if len(im_array) == 0:
            im_array = tmp_array
        else:
            im_array = np.concatenate((im_array,tmp_array),axis=0) #将新的图片加入到之前的图片集
    return im_array

由于我们将数据用于训练时还需要一维的通道, 所以使用以下代码增加一维

im_array = im_array[:,:,:,np.newaxis]

注意 : 此时用于训练的数据均是 0~255的 uint 型的数据

预编译模型

model = Unet('resnet34', input_shape = (512, 512, 1), encoder_weights = None)#1代表通道数  model.compile('Adam', loss = 'binary_crossentropy', metrics = ['accuracy'])

这里创建了一个Unet模型并进行了预编译, 具体参数个人感觉比较好理解, 有不懂的话这里提供中文文档供查阅

Keras-model中文文档参考链接

模型训练和保存

model.fit(
    x = im_array,
    y = la_array,
    batch_size = 10, #batch_size代表每次从im_array中的n个图片中选10个来进行训练
    epochs = 8, #操作8次
    validation_split = 0.2, #取训练集中的0.2作为验证集
    shuffle = True
)
model.save("model_v1.h5") #保存模型 

这里是进行模型的训练和保存, 编译部分的参数可从上面的文档中查询到, 模型保存的方法也可以百度到

图像预测和生成

首先我们取5张图片用上述的方法形成 5 x 512 x 512 x 1 的numpy数组测试样本, 进行预测

model = load_model("model_v1.h5") #加载模型
res=model.predict(te_array,5) #对te_array进行预测,5代表图片个数

此时获取到的res数组应该是 float 类型的 5 x 512 x 512 x 1 数组, 此时的res数组代表的是对应像素位置是 0/1 的概率, 因此res是一个概率数组. 由于我们最后生成的图片的每一位是 0~255 的, 所以此时需要转换.

这里提供两种方法进行转换 :

• 将res数组所有元素乘以255, 再转换成 int 型,然后使用 img.save 方法进行保存

  img = Image.fromarray(im_npp)