语义分割之-加载训练数据

偏执的太偏执、 2022-12-12 04:55 227阅读 0赞

首先十分感谢原博客博主的耐心解答！！帮我解决了问题。  
当你想使用自己的数据集训练一个语义分割神经网络时, 可能会遇到下面的问题：  
1：很多教程在开始时都是使用 mnist 这样的已经处理好的数据, 一个加载函数就实现了数据的加载. 但是要使用自己的数据集时, 很多新手就蒙了, 不知道如何操作  
2：训练的数据比较大, 一次性载入内存又不够  
3：训练的图像尺寸不一样  
对于以上的问题, 就跟大家一起捋一捋  
假设标签都是由 语义分割之 标签生成（https://blog.csdn.net/yx123919804/article/details/107285881） 中讲的那样生成的, 且图像都放到了 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 文件同名的全部文件夹 ##

在 语义分割之 标签生成 （https://blog.csdn.net/yx123919804/article/details/107285881）中已经生成了所需要的标签图像, 所以只需要将其列出来加载就可以了  
这里我们只列出与 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 os.path.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):
        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 = 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]) # x, y, width, height
    
                    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)\]

## 三. 训练 ##

假设你已经定义好了模型, 只差训练了. 下面讲如何把训练数据传到训练函数中, 如何定义模型之类的, 以后的文章会讲.（https://blog.csdn.net/yx123919804/article/details/104654323）  
https://blog.csdn.net/yx123919804/article/details/104811087

# 训练模型
    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, 一定要去把它弄明白

按照上面的代码会出现不能预测图片：所以需要把代码改一下“

# 读图像和标签
    # 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

如果发现输入图像和自己不一样 把代码中的+32这个去掉 然后我的项目出现维度不正常的现象 我把这句话改正了如下：我的图像大小是448的：

if len(train) == batch_size:
            if train_mode:
                yield (np.array(train).reshape(1, 448, 448, 3), np.array(label).reshape(1, 448, 448, 1))
            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)

放下效果图：：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 1]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 2]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 3]  
最后再次感谢源博客博主！！！

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/75b89e85d28a47ffb872e29563a305eb.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 1]: /images/20221123/ec577ebbfaad424fa3d16ec2bb19c5ea.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 2]: /images/20221123/90aab4158fca45e485bcd9acb3acc081.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pxeDk1MTEwMg_size_16_color_FFFFFF_t_70_pic_center 3]: /images/20221123/3d7c01818cf341589a10681bf75f1b6c.png