Keras 实现 FCN 语义分割并训练自己的数据之 FCN-32s
语义分割说起来高大上, 其实就是一个抠图大法, 要让电脑学会自动抠图, 那要先教其如何抠图
相信你已经把 FCN 的教程看得够多了, 现在需要的是手动实践一下, 但是又不知道从哪里下手, 那你看这篇文章就看对了
这里我以二分类为例, 多分类的以后再讲, 其实也没有多大区别, 万事先从简单的开始循序渐进才能降低学习难度, 下面要讲的网络结构也是从 FCN 最简单的形式开始
在 语义分割之 加载训练数据 中已经讲过了如何用 Generator 的方式加载数据, 现在就只差训练了
一. Keras Sequential 模型定义全卷积神经网络模型(FCN)
相对于 Sequential 模型, Functional API 可以实现更复杂的模型, 但是今天我们要从简单的入手, 所以先选择 Sequential 模型来实现. 整个过程我用的是 Jupyter Notebook 中完成的, 需要的库如下
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import os.path as osp
import cv2 as cv
import matplotlib.pyplot as plt
from random import shuffle
from PIL import Image # 如果自己写代码的话可以不用, 用 OpenCV 就可以了
import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras.callbacks import ModelCheckpoint
这里的模型旨在说明问题, 先不管性能和效果如何, 把网络搭建好并跑起来才是初学者应该关心的问题. 而不是一开始就要实现一个完美的网络模型, 所以不要想一下子就搞得非常好, 等把这个简单的跑起来之后, 再分析一下哪里不足改哪里. 这样改的时候才有了改的基础. 这才是动手能力强的人的逻辑. 不要坐在那里空想, 一定要动手敲代码. 模型定义如下
# 定义模型
project_name = "fcn_segment"
channels = 3
std_shape = (320, 320, channels) # 输入尺寸, std_shape[0]: img_rows, std_shape[1]: img_cols
# 可以设置为 (None, None channels) 适应变化的尺寸
model = keras.Sequential(name = project_name)
# 第一组
model.add(keras.layers.Conv2D(32, kernel_size = (3, 3), activation = "relu",
padding = "same", input_shape = std_shape,
name = "conv_1"))
model.add(keras.layers.MaxPool2D(pool_size = (2, 2), strides = (2, 2), name = "max_pool_1"))
# 第二组
model.add(keras.layers.Conv2D(64, kernel_size = (3, 3), activation = "relu",
padding = "same", name = "conv_2"))
model.add(keras.layers.MaxPool2D(pool_size = (2, 2), strides = (2, 2), name = "max_pool_2"))
# 第三组
model.add(keras.layers.Conv2D(128, kernel_size = (3, 3), activation = "relu",
padding = "same", name = "conv_3"))
model.add(keras.layers.MaxPool2D(pool_size = (2, 2), strides = (2, 2), name = "max_pool_3"))
# 第四组
model.add(keras.layers.Conv2D(256, kernel_size = (3, 3), activation = "relu",
padding = "same", name = "conv_4"))
model.add(keras.layers.MaxPool2D(pool_size = (2, 2), strides = (2, 2), name = "max_pool_4"))
# 第五组
model.add(keras.layers.Conv2D(512, kernel_size = (3, 3), activation = "relu",
padding = "same", name = "conv_5"))
model.add(keras.layers.MaxPool2D(pool_size = (2, 2), strides = (2, 2), name = "max_pool_5"))
这里要说明一点是, 虽然网络不那么严格的限制输入图像的大小, 但是一个 batch 中的图像大小要相同才行, 当不能保证一个 batch 中的图像尺寸相同时, batch 只能设置成 1. 也有办法处理大小不一样且 batch_size > 1 的情况, 比如使用 语义分割之 加载训练数据 中的输入方式
注: 以下的输出 shape 计算是依据输入尺寸为 (320, 320) 计算的
相信你已经看出来了, 上面的模型并没有全部完成, 只表示了卷积和池化的部分, 一共有 5 组 10 层, 每组包括一个 Conv2D 和一个 MaxPool2D. 最后的输出是 10 * 10 * 512, 表示 大小为 10 * 10, 通道数为 512. 这就是和传统的 CNN 网络相同的部分. 这部分用来检测图像中是否有目标存在, 并且在哪里, 也就是 Feature map. 在结构上和 VGG 之类的并不相同, 因为为了说明问题, 先把所有的问题都简化. 图像经过上面的 5 组操作之后, 尺寸缩小到原来的 1 / 32
接下来的网络就是要把上面的输出还原到与输入相同的尺寸. 直接来一个简单粗暴的 32 倍上采样还原到原来的输入尺寸, 这里使用不需要训练的 UpSampling2D 来进行, 默认使用 nearest 插值方式, 换成 bilinear 有惊喜. 接上面的模型继续
# 上采样
model.add(keras.layers.UpSampling2D(size = (32, 32), interpolation = "nearest", name = "upsamping_6"))
好了, 网络已经 差不多 定义完成, 是不是比你想象的更简单. 当然这样的网络性能可能好不到哪里去, 至少我们按 FCN 的原理迈出了一大步
到这里应该有同学提问才对, 输出是有了, 是不是差了点什么东西但是又说不上来. Yes, 缺少了评价输出结果的损失函数
想一下之前学过的全连接神经网络, 在二分类问题上, 输出层只有一个神经元, 使用 Sigmoid 激活函数, 损失函数为 binary_crossentropy. 那问题来了, 现在的输出可不只一个神经元, 而是和输入相同尺寸的且有 512 个通道的图像, 乖乖不得了了, 怎么弄?
回忆一下, FCN 语义分割最后是对像素进行分类, 有多少类最后的输出图像就有多少个通道, 每个通道的像素值代表了这个通道的像素应划分到哪一个类别的概率, 如果某一个像素位置在第 3 通道的值最大, 那这个位置的像素就属性第 3 个分类
好了, 是不是有一点眉目了, 二分类也是分类的一种吧, 所以我们需要把上面输出的 512 个通道变成两个通道. 这里我们再简化一下, 二分类只需要一个通道, 值小于 0.5 就是背景, 大于 0.5 就是目标. 这样是不是就可以用 Sigmoid 激活了呢, 通道数是定下来了, 但是要怎么变成 1 个通道呢?
你再看一下上面的卷积层, 通道数是不是随卷积核数目在增加? 因为通道数目取决于卷积核的数目, 所以我们只需要用 1 个 512 通道的卷积核再卷积一下就可以了
# 上采样
model.add(keras.layers.UpSampling2D(size = (32, 32), interpolation = "nearest", name = "upsamping_6"))
# 这里只有一个卷积核, 可以把 kernel_size 改成 1 * 1, 也可以是其他的, 只是要注意 padding 的尺寸
# 也可以放到 upsamping_6 的前面, 试着改一下尺寸和顺序看一下效果
# 这里只是说明问题, 尺寸和顺序不一定是最好的
model.add(keras.layers.Conv2D(1, kernel_size = (3, 3), activation = "sigmoid",
padding = "same", name = "conv_7"))
model.summary()
到了这里, 输出的通道数是 1 了, Sigmoid 激活函数也出现了. 调用 model.summary() 可以输出模型结构与参数的数量
Model: "fcn_segment"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
conv_1 (Conv2D) (None, 320, 320, 32) 896
_________________________________________________________________
max_pool_1 (MaxPooling2D) (None, 160, 160, 32) 0
_________________________________________________________________
conv_2 (Conv2D) (None, 160, 160, 64) 18496
_________________________________________________________________
max_pool_2 (MaxPooling2D) (None, 80, 80, 64) 0
_________________________________________________________________
conv_3 (Conv2D) (None, 80, 80, 128) 73856
_________________________________________________________________
max_pool_3 (MaxPooling2D) (None, 40, 40, 128) 0
_________________________________________________________________
conv_4 (Conv2D) (None, 40, 40, 256) 295168
_________________________________________________________________
max_pool_4 (MaxPooling2D) (None, 20, 20, 256) 0
_________________________________________________________________
conv_5 (Conv2D) (None, 20, 20, 512) 1180160
_________________________________________________________________
max_pool_5 (MaxPooling2D) (None, 10, 10, 512) 0
_________________________________________________________________
upsamping_6 (UpSampling2D) (None, 320, 320, 512) 0
_________________________________________________________________
conv_7 (Conv2D) (None, 320, 320, 1) 4609
=================================================================
Total params: 1,573,185
Trainable params: 1,573,185
Non-trainable params: 0
_________________________________________________________________
最后 compile, 其中 optimizer, loss, metrics 这些都算是基本操作了
model.compile(optimizer = "adam",
loss = "binary_crossentropy",
metrics = ["accuracy"])
二. 训练
语义分割之 加载训练数据 中已经讲过训练的套路了. 这里就按套路来, 接代码如下
# 训练模型
epochs = 16
batch_size = 1 # 这里用了 segment_reader, batch_size 就可以大于 1 了, 不过太大显卡可能装不下
augmented = 4 # 默认增强数据只有翻转, 一张图像变成 4 张
# train_path 和 valid_path 由前面的 get_data_path 得到
# 这两个 reader 源源不断地提供训练所需要的数据, 这就是 yield 神奇的地方
train_reader = segment_reader(train_path, batch_size)
valid_reader = segment_reader(valid_path, batch_size)
history = model.fit(x = train_reader, # 训练数据
steps_per_epoch = len(train_path) * augmented // batch_size,
epochs = epochs,
verbose = 1,
validation_data = valid_reader, # 验证数据
validation_steps = max(1, len(valid_path) * augmented // batch_size),
max_queue_size = 8,
workers = 1)
加载数据时 data_set_path 返回的是参与训练的图像和标签的路径, 这个路径可以统计出有多少样本. 在上面代码中我们定义了一个变量 augmented 来表示一张图像做数据增强后变成的数量, 默认只是将图像进行翻转, 所以是 4. 这个是用来计算 steps_per_epoch 的, steps_per_epoch 表示在 1 个 epoch 中的步数, 也就是样本数量除以 batch_size
随便训练一下就有很高的正确率, 不过不能说明什么, 因为背景的像素点比目标的像素点多得多
三. 预测
预测很简单, 直接喂数据就可以, 只是要将图像变成网络可以接受的格式. 也可以用 Generator 的方式进行
# segment_reader 方式预测
# test_path 可以是 get_data_path 返回的 test 目录, 要分割成三部分, 可以像下面这样
# train_path, valid_path, test_path = get_data_path(data_path, (0.7, 0.2, 0.1))
# test_path 也可以是你手动列出来的目录
#
# train_mode = False 可以返回 roi, 运行一次可以预测一个 batch_size
test_path = r"D:\raccoon_test"
test_path = [osp.join(test_path, x) for x in os.listdir(test_path)]
test_reader = segment_reader(test_path, batch_size = 1, train_mode = False)
上面代码演示的是直接指定测试图像所在路径, 有了上面的 test_reader 就可以预测了
# 运行一次测试一个 batch_size
batch_x, batch_roi = next(test_reader)
batch_y = model.predict(batch_x)
那怎么显示预测结果呢? segment_reader 在测试模式也就是 train_mode 为 False 时, 返回的是测试图像和在扩展后图像中的位置, 根据 roi 就可以把测试的图像裁出来
# 显示预测结果
show_index = 0 # 显示 batch_size 中的序号, 这里显示第 0 个
roi = batch_roi[show_index]
# 从扩展后的图像中裁出来
x = batch_x[show_index][roi[1]: roi[1] + roi[3], roi[0]: roi[0] + roi[2]]
y = batch_y[show_index][roi[1]: roi[1] + roi[3], roi[0]: roi[0] + roi[2]]
plt.figure("segment", figsize = (8, 4))
plt.subplot(1, 2, 1)
plt.axis("on")
plt.title("test", color = "orange")
plt.imshow(x[..., : : -1])
plt.subplot(1, 2, 2)
plt.axis("on")
plt.title("predict", color = "orange")
plt.imshow(np.squeeze(y), cmap = "gray")
plt.show()
直接 32 倍上采样出来的结果很粗糙, 贴一个我用 nearest 插值和 groudn_truth 比较的图
此时 predict 图像并不是二值化的结果, 因为没有做二值化的操作, 各像素值是 y 的输出值. 要二值化的话, 自己选择一个合适的阈值, 不一定是 0.5
四. 输出标记
预测结果是有了, 如何标记到原始图像上? 选择一个合适的阈值, 预测值超过阈值就在原图上增加一个颜色值就可以了, 看起来标记也是透明的, 还可以单独标记一个 Mask 以作他用
# 标记预测结果
img_marked = x.copy(); # 标记后的图像
# 单独使用的 Mask
img_mask = np.zeros((img_marked.shape[0], img_marked.shape[1], 1), dtype = np.uint8)
for r in range(img_marked.shape[0]):
for c in range(img_marked.shape[1]):
if y[r][c] >= 0.5: # 阈值
img_marked[r][c] += [0, 0, 0.3] # 在 img_marked 上标记为红色
# 三个值分别是 BGR 颜色, 值越小越透明
img_mask[r][c] = 255
plt.figure("mark_image", figsize = (6, 3))
plt.subplot(1, 1, 1)
plt.axis("on")
plt.title("img_marked", color = "orange")
plt.imshow(img_marked[..., : : -1])
标记后的效果是这样的, 只是在原图上画出相应的颜色
五. 大图训练与预测
当输入尺寸一般大, 硬件能支撑的话, 可以直接进行训练, 但是像 2596 * 1944 这样的大图的话, 一般硬件是吃不消的, 肿么办? 不要慌, 我们不是还有 Generator 吗, 我可以在其中将大图裁切成能训练的小图. 修改 语义分割之 加载训练数据 中的代码如下
# 读图像和标签 segment_reader
# data_set: 就是上面 get_data_path 返回的路径
# batch_size: 一次加载图像的数量
# augment_fun: 数据增强函数
# padding_size: 扩展尺寸, 如果为 0 表示不扩展
# crop_size: 裁切尺寸
# zero_based: True, 右边和底边扩展, False, 四周扩展
# train_mode: 训练模式, 对应的是测试模式, 测试模式会返回 roi 矩形, 方便还原原始的尺寸
# shuffle_enable: 是否要打乱数据, 这个是上面 get_data_path 打乱有什么不一样呢,
# 这个打乱是每一个 epoch 会打乱一次
def segment_reader(data_set, batch_size = 1, augment_fun = None,
padding_size = 32, zero_based = False, crop_size = None,
train_mode = True, shuffle_enable = True):
assert(isinstance(data_set, tuple) or isinstance(data_set, list))
stop_now = False
data_nums = len(data_set)
index_list = [x for x in range(data_nums)] # 用这个列表序号来打乱 data_set 排序
image = [] # 图像
label = [] # 标签
max_rows = 0 # 记录一个 batch 中图像的最大行数
max_cols = 0 # 记录一个 batch 中图像的最大列数
while False == stop_now:
if train_mode and shuffle_enable:
shuffle(index_list)
for i in index_list:
is_with_label = 2 == len(data_set[i]) # 如果 2 == data_set[i], 表示带标签输入, 否则只有图像
data_list = [] # 暂时存放用
if is_with_label:
x = cv.imread(data_set[i][0], cv.IMREAD_UNCHANGED) # 读出图像
# 如果标签图像像素值就是类别的话, 可以这样读
y = cv.imread(data_set[i][1], cv.IMREAD_GRAYSCALE) # 读出标签, 注意是 IMREAD_GRAYSCALE
# 这里如果标签图像是索引图像的话, 可以这样读
# y = np.array(Image.open(data_set[i][1], 'r'))
data_list.append([x, y])
if train_mode:
if augment_fun: # 如果提供了数据增强函数且为训练模式
# 将增强后的数据放到 data_list 中, 注意这里要用 extend
data_list.extend(augment_fun(x, y))
else: # 否则用默认的增强方式, 就是上下左右翻转
for axis in (-1, 0, 1):
x_flip = cv.flip(x, axis)
y_flip = cv.flip(y, axis)
data_list.append([x_flip, y_flip])
else:
train_mode = False
x = cv.imread(data_set[i], cv.IMREAD_UNCHANGED)
data_list.append([x, (0, 0, 0, 0)]) # 这里只有 x, 没有 y, 先用 0 占位, 下面会修改
for data in data_list:
# 按需要的尺寸裁切
if crop_size:
overlap = 16
image_shape = data[0].shape
# 扩展宽度与高度
# 扩展边界成裁切尺寸的倍数 + overlap * 2
# 加 overlap * 2 是为了各块之间有重叠, 这样效果会好了一点
w = crop_size[1] - image_shape[1] % crop_size[1] + overlap * 2
h = crop_size[0] - image_shape[0] % crop_size[0] + overlap * 2
x = w // 2
y = h // 2
data[0] = cv.copyMakeBorder(data[0], y, h - y, x, w - x,
cv.BORDER_CONSTANT, (0, 0, 0))
if is_with_label and train_mode:
data[1] = cv.copyMakeBorder(data[1], y, h - y, x, w - x,
cv.BORDER_CONSTANT, (0, 0, 0))
max_rows = image_shape[0] + crop_size[0] - image_shape[0] % crop_size[0]
max_cols = image_shape[1] + crop_size[1] - image_shape[1] % crop_size[1]
for r in range(max_rows // crop_size[0]):
top = r * crop_size[0]
for c in range(max_cols // crop_size[1]):
left = c * crop_size[1]
# 多出来的像素主要是分块时有重叠部分, 这样效果会好一点
# 因为输入的尺寸不限制大小, 这样并不影响
bottom = top + crop_size[0] + overlap * 2
right = left + crop_size[1] + overlap * 2
image_crop = data[0][top: bottom, left: right]
image.append(np.array(image_crop).astype(np.float32) / 255.0)
if is_with_label and train_mode:
image_crop = data[1][top: bottom, left: right]
label.append(np.array(image_crop).astype(np.float32))
else:
label.append((c + overlap, r + overlap, crop_size[1], crop_size[0]))
if len(image) == batch_size:
yield (np.array(image), np.array(label))
image = []
label = []
else:
max_rows = max(64, max_rows, data[0].shape[0])
max_cols = max(64, max_cols, data[0].shape[1])
image.append(data[0])
label.append(data[1])
if len(image) >= batch_size:
# 一个 batch 中图像的尺寸不一样是不能一起训练的, 所以要将其统一到相同的尺寸
# 加 padding_size 是为了扩展边缘, 不然边缘的像素可能分割不好
if padding_size > 0:
max_rows = max_rows // padding_size * padding_size + padding_size
max_cols = max_cols // padding_size * padding_size + padding_size
for i in range(batch_size):
image_shape = image[i].shape
# 扩展宽度与高度
w = max_cols - image_shape[1]
h = max_rows - image_shape[0]
# 判断是四周扩展还是右面下面扩展
x = 0 if zero_based else w // 2
y = 0 if zero_based else h // 2
# 扩展边界
image[i] = cv.copyMakeBorder(image[i], y, h - y, x, w - x,
cv.BORDER_CONSTANT, (0, 0, 0))
# 转换成 0~1 的范围
image[i] = np.array(image[i]).astype(np.float32) / 255.0
if is_with_label and train_mode:
label[i] = cv.copyMakeBorder(label[i], y, h - y, x, w - x,
cv.BORDER_CONSTANT, (0, 0, 0))
label[i] = np.array(label[i]).astype(np.float32) # 注意这里不用除以 255
else: # 如果不带标签或者是测试模式, 则返回原图像在扩展图像中的位置
label[i] = (x, y, image_shape[1], image_shape[0])
yield (np.array(image), np.array(label))
image = []
label = []
max_rows = 0
max_cols = 0
if False == train_mode:
stop_now = True
代码中增加了一个 crop_size 参数, 如果设置了此参数就说明需要将大的图裁切成小的图训练. 这个裁切的尺寸和裁切的代码可能需要根据你的网络进行适当修改. 调用方式如下
# segment_reader 读入图像
# 如果是分块的话, 需要指定分块大小, 比如 crop_size = (320, 320)
show_reader = segment_reader(train_path, 1, crop_size = (320, 320)) # 读一张图, 当然你也可以多读几张
如果分块的话, 比如分成 4 块, 那让面训练时定义的 augmented 就变成了 4 × 4 = 16
训练的时候用小图, 但是预测的时候不需要那么多的显存, 也许可以直接用大图预测
六. 代码下载
完整的代码可下载 Jupyter Notebook 代码示例
上一篇: 语义分割之 加载训练数据
下一篇: Keras 实现 FCN 语义分割并训练自己的数据之 FCN-16s、FCN-8s