首先十分感谢原博客博主的耐心解答！！帮我解决了问题。
当你想使用自己的数据集训练一个语义分割神经网络时, 可能会遇到下面的问题：
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)
放下效果图：：
最后再次感谢源博客博主！！！