语义分割之 加载训练数据

当你想使用自己的数据集训练一个语义分割神经网络时, 可能会遇到下面的问题

• 很多教程在开始时都是使用 mnist 这样的已经处理好的数据, 一个加载函数就实现了数据的加载. 但是要使用自己的数据集时, 很多新手就蒙了, 不知道如何操作
• 训练的数据比较大, 一次性载入内存又不够
• 训练的图像尺寸不一样

对于以上的问题, 就跟大家一起捋一捋
假设标签都是由 语义分割之 标签生成 中讲的那样生成的, 且图像都放到了 D:\raccoon. 一张图像对应一个 json 文件和一个文件夹, 文件夹里面有需要的 img.png, label.png和 label_viz.png

需要使用到的库

# !/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import os.path as osp

import cv2 as cv
import numpy as np
from random import shuffle
from PIL import Image # 如果自己写代码的话可以不用, 用 OpenCV 就可以了

import matplotlib.pyplot as plt

一. 列出数据集目录中与 json 文件同名的全部文件夹

在 语义分割之 标签生成 中已经生成了所需要的标签图像, 所以只需要将其列出来加载就可以了

这里我们只列出与 json 文件同名的全部文件夹, 也就是训练图像和标签图像的路径, 但是没有读其中的图像. 这样方便以后数据生成器(generator)的操作. 有了这个路径列表, 以后图片即可以一次性全部读入内存, 也可以批量读入

# 取得 data_set 目录中的文件
# data_set_path: 数据集所在文件夹, 可以是文件夹列表, 因为你有可能将不同类别数据放到不同文件中
# split_rate: 这些文件中用于训练, 验证, 测试所占的比例
#             如果为 None, 则不区分, 直接返回全部
#             如果只写一个小数, 如 0.8, 则表示 80% 为训练集, 20% 为验证集, 没有测试集
#             如果是一个 tuple 或 list, 只有一个元素的话, 同上面的一个小数的情况
# shuffle_enable: 是否要打乱顺序
# 返回训练集, 验证集和验证集路径列表
def get_data_path(data_set_path, split_rate = (0.7, 0.2, 0.1), shuffle_enable = True):
    data_path = []
    
    if not isinstance(data_set_path, list):	# 变成列表可以以统一的方式进行操作
        data_set_path = [data_set_path]
   
    for dsi in data_set_path:
        folders = os.listdir(dsi)
        
        for folder in folders:
            if osp.isdir(osp.join(dsi, folder)):
                data_path.append(osp.join(dsi, folder));
                
    if shuffle_enable:
        shuffle(data_path)
        
    if None == split_rate:
        return data_path

    total_num = len(data_path)

    if isinstance(split_rate, float) or 1 == len(split_rate):
        if isinstance(split_rate, float):
            split_rate = [split_rate]
        train_pos = int(total_num * split_rate[0])
        train_set = data_path[: train_pos]
        valid_set = data_path[train_pos: ]

        return train_set, valid_set

    elif isinstance(split_rate, tuple) or isinstance(split_rate, list):
        list_len = len(split_rate)
        assert(list_len > 1)

        train_pos = int(total_num * split_rate[0])
        valid_pos = int(total_num * (split_rate[0] + split_rate[1]))

        train_set = data_path[0: train_pos]
        valid_set = data_path[train_pos: valid_pos]
        test_set = data_path[valid_pos: ]

        return train_set, valid_set, test_set

调用 get_data_path 就可以得到训练集和验证集的目录

data_path = "D:\\raccoon"
train_path, valid_path = get_data_path(data_path, split_rate = 0.8)

print(" Total number: ", len(train_path) + len(valid_path),
      " Train number: ", len(train_path),
      " Valid number: ", len(valid_path))

如果你有多个路径的时候可以这样调用

data_path = ["D:\\raccoon", "D:\\dog", "D:\\cat"]
train_path, valid_path = get_data_path(data_path, split_rate = 0.8)

print(" Total number: ", len(train_path) + len(valid_path),
      " Train number: ", len(train_path),
      " Valid number: ", len(valid_path))

二. 根据目录加载图像

上面已经生成了路径列表, 就可以根据列表中的路径加载图像了. 一般我们希望在 pooling 操作时图像的尺寸是 2 的倍数, 这样在融合或者拼接的时候, 两个 layer 中的尺寸能对应上, 所以下面的代码中增加了 padding 的操作, 让图像的尺寸可以被 32 整除. 为什么是 32 呢? 这个要看你的网络是怎么设计的, FCN32s 下采样 5 次, 就变成原来的 1 / 32, 这里也就用了 32. 当然这一步不是必须的, 你也可以在网络中处理尺寸不匹配的问题

# 读图像和标签
# data_path: 就是上面 get_data_path 返回的路径
# batch_size: 一次加载图像的数量, -1 表示加载全部
# zero_based: True, 四周扩展, False, 右边和底边扩展
# train_mode: 训练模式, 对应的是测试模式, 测试模式会返回 roi 矩形, 方便还原原始的尺寸
# shuffle_enable: 是否要打乱数据, 这个是上面 get_data_path 打乱有什么不一样呢, 这个打乱是每一个 epoch 会
#                 会打乱一次
def segment_reader(data_path, batch_size, zero_based = False, train_mode = True, shuffle_enable = True):
    if not train_mode:
        if osp.isdir(data_path):
            bkup = data_path
            test_files = os.listdir(data_path)
            
            data_path = []
            for f in test_files:
                data_path.append(osp.join(bkup, f))
        else:
            data_path = [data_path]
            
    data_nums = len(data_path)
    index_list = [x for x in range(data_nums)]
    
    stop_now = False
    
    while False == stop_now:
        if shuffle_enable:
            shuffle(index_list)
            
        train = []
        label = []
        
        for i in index_list:
            read_path = data_path[i]
            if train_mode:
                read_path = osp.join(data_path[i], "img.png")
                
            if not osp.exists(read_path):
                continue
                
            img_src = cv.imread(read_path)
                
            shape = img_src.shape
            
            max_rows = max(64, shape[0])
            max_rows = max_rows // 32 * 32 + 32 # +32 是为了扩展边缘, 不然边缘的像素可能分割不好
            max_cols = max(64, shape[1])
            max_cols = max_cols // 32 * 32 + 32

            b = max_rows - shape[0]
            r = max_cols - shape[1]

            half_padding_x = 0
            half_padding_y = 0
            
            if False == zero_based:
                half_padding_x = r >> 1
                half_padding_y = b >> 1

            # 扩展边界
            big_x = cv.copyMakeBorder(img_src,
                                      half_padding_y, b - half_padding_y,
                                      half_padding_x, r - half_padding_x,
                                      cv.BORDER_REPLICATE, (0, 0, 0))

            # 转换成 0~1 的范围
            big_x = np.array(big_x).astype(np.float32) / 255.0

            if train_mode:                
                read_path = osp.join(data_path[i], "label.png")
                
                # 如果标签图像像素值就是类别的话, 可以这样读
                '''
                img_label = cv.imread(read_path, cv.IMREAD_GRAYSCALE)
                big_y = cv.copyMakeBorder(img_label,
                                          half_padding_y, b - half_padding_y,
                                          half_padding_x, r - half_padding_x,
                                          cv.BORDER_REPLICATE, (0, 0, 0))
                
                big_y = np.array(big_y).astype(np.float32) # 注意这里不用除以 255
                
                '''
                # 这里如果标签图像是索引图像的话, 可以这样读                
                img_label = np.array(Image.open(read_path, 'r'))
                big_y = cv.copyMakeBorder(img_label,
                                          half_padding_y, b - half_padding_y,
                                          half_padding_x, r - half_padding_x,
                                          cv.BORDER_REPLICATE, (0, 0, 0))
                big_y = big_y.astype(np.float32) # 注意这里不用除以 255
                
               
                train.append(big_x)
                label.append(big_y)
            else:
                roi = (half_padding_x, half_padding_y, shape[1], shape[0])

                train.append(big_x)
                label.append(roi)
            

            if len(train) == batch_size:
                if train_mode:
                    yield (np.array(train), np.array(label))
                else:
                    yield (np.array(train), label)

                train = []
                label = []
                
        if train:
            if train_mode:
                yield (np.array(train), np.array(label))
            else:
                yield (np.array(train), label)
            
            train = []
            label = []
            
        if False == train_mode:
            stop_now = True

虽然可以用更简单的代码一次性把图像都载入内存, 但是用 Generator 的方式可以适应内存或显存不够用的情况, 也可以全部载入, 所以更方便. 以不变应万变.

在上面 segment_reader 函数中, 读标签图像用的是 img_label = cv.imread(read_path, cv.IMREAD_GRAYSCALE). 这是对应自己写代码的情况, 此时标签图像是单通道的灰度图像. 如果用 Labelme 的生成工具的话, 就要注意了. 此时标签是索引图像, 读的时候要用 img_label = np.array(Image.open(read_path, ‘r’)), 不然标签图像会变成 3 通道的图像, 就不能训练了

有了上面的生成器, 就可以按批量读数据了, 在训练之前我们显示一下读出来的图像看一下正确不

# 显示图像
show_reader = segment_reader(train_path, 1) # 读一张图, 当然你也可以多读几张
x, y = next(show_reader)

show_index = 0, # 不能超过 segment_reader 设置的 batch_size - 1
plt.figure("show_image", figsize = (8, 4))
plt.subplot(1, 2, 1)
plt.title("raccoon", color = "orange")
plt.imshow(x[show_index][..., : : -1]) # 因为 Opencv 内存的数据是 BGR, 所以要反成 RGB 才能正常显示

plt.subplot(1, 2, 2)
plt.title("ground_truth", color = "orange")
plt.imshow(np.squeeze(y[show_index]), cmap = "gray")
plt.show()

图像显示如下

如果 train_mode 设置为 False, 则返回的是原始图像和原始图像在扩展图像中的位置, 这样方便预测之后还原

# 显示图像
show_reader = segment_reader(train_path, 1, train_mode = False)
x, roi = next(show_reader)
print(roi)

打印结果如下

[(11, 15, 650, 417)]

三. 训练

假设你已经定义好了模型, 只差训练了. 下面讲如何把训练数据传到训练函数中, 如何定义模型之类的, 以后的文章会讲.

# 训练模型
epochs = 16
batch_size = 1	# 当训练图像的尺寸不一样时, 就只能是 1, 不然会把输入数据 shape 不对

# train_path 和 valid_path 由前面的 get_data_path 得到
# 这两个 reader 源源不断地提供训练所需要的数据, 这就是 yield 神奇的地方
train_reader = segment_reader(train_path, batch_size, shuffle_enable = True)
valid_reader = segment_reader(valid_path, batch_size, shuffle_enable = True)

history = model.fit(x = train_reader, 	# 训练数据
                    steps_per_epoch = len(train_path) // batch_size,
                    epochs = epochs,
                    verbose = 1,
                    validation_data = valid_reader,	# 验证数据
                    validation_steps = max(1, len(valid_path) // batch_size),
                    max_queue_size = 8,
                    workers = 1)

有了 get_data_path 和 segment_reader 之后, 训练就变成套路了

1. 定义好神经网络
2. 调用 get_data_path 得到训练数据的路径
3. 由 segment_reader 读入训练数据
4. 调用 fit 函数就开始训练了

四. 其他

有了 Generator, 妈妈不用再担心你的内存显存不够用. 就算够用, 你也可以用它一招对付不同的情况. Generator 是非常灵活的, 你可以在 Generator 中做数据的增强, 可以裁切出目标区域参与训练而排除其他区域等. 如果你理解了 Generator 的原理, 就可以实现更灵活与复杂的功能, 不一定按文章中的来, 但是套路是一样的.

如果不能理解 Generator, 一定要去把它弄明白

上一篇: 语义分割之 标签生成
下一篇: Keras 实现 FCN 语义分割并训练自己的数据之 FCN-32s