学习前言
- 数据集和标签详细信息
- LOSS函数计算
- 2.训练代码
- 1.数据集的制作
- 2.模型的训练
- 3.模型的预测
- 结果显示
前言
在第一部分完成了主干网络VGGnet还有segnet模型，这节主要简单讲一下模型的训练以及预测过程。如有遗忘参考小生不财-语义分割1：基于VGGNet模型的segnet讲解
数据集和标签详细信息
用于训练语义分割模型的文件分为两部分。
第一部分是原始图像，如下所示：
标签的第二部分，如下所示：
当你们看到这个标签的时候你们会说，我靠，你给我看的什么辣鸡，全黑的算什么标签，其实并不是这样的，这个标签看起来全黑，但是实际上在斑马线的部分其RGB三个通道的值都是1。
其实给你换一张图，就能看得更清楚。
这是它的标签。
为什么这里的标签看起来就清楚的多呢，因为在voc中，其一共需要分21类，所以火车的RGB的值可能都大于10了，当然看得见。所以，在训练集中，如果像本文一样分两类，那么背景的RGB就是000，斑马线的RGB就是111，如果分多类，那么还会存在222，333，444这样的。这说明其属于不同的类。
2.代码训练
1、数据集的制作
这里我自己制作了191例数据集，这里介绍如何把数据集制作成训练集和测试集，具体代码和注释如下：
def train_data(NC):
# 读取文件名，一部分是原图名字，另一部分是标签图片
# txt是按行存的，读取时也是按照行读，中间使用的是';'分割的
with open('dataset2/train.txt') as f:
data = f.readlines()
train = []
label = []
for ele in data:
# 第一部分，读取训练集图片
# 读取训练图片
img = Image.open("dataset2/jpg/" + ele.split(';')[0])
# 调整图片大小
img = img.resize((WIDTH, HEIGHT))
img = np.array(img)
# 对图像进行归一化处理
img = img / 255
# 将图像存入train
train.append(img)
# 第二部分，读取label 这里必须和网络输出的格式一致，也就是说图片的长宽都要变为原来的1/2
img = Image.open('dataset2/png/' + ele.split(';')[1].replace('\n', ''))
# 调整图片大小
img = img.resize((int(WIDTH / 2), int(HEIGHT / 2)))
# 将图像转化为array ，然后进行归一化
img = np.array(img)
img = img / 255
# 生成一个和segnet网络输出结构一致的图像
seg_lable = np.zeros([int(WIDTH / 2), int(HEIGHT / 2), NC])
# 因为斑马线三个通道都是1，背景为0
for i in range(NC):
# 因为标签图片，RGB第三个通道的值为1
seg_lable[:, :, i] = (img[:, :, 0] == i).astype(int)
# 进行一次reshape,将通道数进行调整，调整到和decoder格式一样
seg_lable = np.reshape(seg_lable, (-1, NC))
# 存放表情包
label.append(seg_lable)
return train, label2、完整训练代码
from net.segnet import mobilenet_segnet
from tensorflow.keras.optimizers import Adam
from PIL import Image
import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import binary_crossentropy
from tensorflow.keras import backend as k
import os
# 定义分类种类
NCLASSES = 2
HEIGHT = 416
WIDTH = 416
# 制作数据集
def train_data(NC):
# 读取文件名，一部分是原图名字，另一部分是标签图片
# txt是按行存的，读取时也是按照行读，中间使用的是';'分割的
with open('dataset2/train.txt') as f:
data = f.readlines()
train = []
label = []
for ele in data:
# 第一部分，读取训练集图片
# 读取训练图片
img = Image.open("dataset2/jpg/" + ele.split(';')[0])
# 调整图片大小
img = img.resize((WIDTH, HEIGHT))
img = np.array(img)
# 对图像进行归一化处理
img = img / 255
# 将图像存入train
train.append(img)
# 第二部分，读取label 这里必须和网络输出的格式一致，也就是说图片的长宽都要变为原来的1/2
img = Image.open('dataset2/png/' + ele.split(';')[1].replace('\n', ''))
# 调整图片大小
img = img.resize((int(WIDTH / 2), int(HEIGHT / 2)))
# 将图像转化为array ，然后进行归一化
img = np.array(img)
img = img / 255
# 生成一个和segnet网络输出结构一致的图像
seg_lable = np.zeros([int(WIDTH / 2), int(HEIGHT / 2), NC])
# 因为斑马线三个通道都是1，背景为0
for i in range(NC):
# 因为标签图片，RGB第三个通道的值为1
seg_lable[:, :, i] = (img[:, :, 0] == i).astype(int)
# 进行一次reshape,将通道数进行调整，调整到和decoder格式一样
seg_lable = np.reshape(seg_lable, (-1, NC))
# 存放表情包
label.append(seg_lable)
return train, label
# 自定义损失函数
def loss(y_true, ypred):
crossloss = binary_crossentropy(y_true, ypred)
# 计算损失平均值
loss = 4 * k.sum(crossloss) / HEIGHT / WIDTH
return loss
if __name__ == "__main__":
# 加载数据
train, labels = train_data(NCLASSES)
train = np.array(train)
labels = np.array(labels)
# 对数据集进行打乱处理
np.random.seed(11)
np.random.shuffle(train)
np.random.seed(11)
np.random.shuffle(labels)
print(train.shape)
print(labels.shape)
# 划分数据集
total_num = train.shape[0]
train_num = int(total_num * 0.9)
x_train = train[:train_num, :, :, :]
y_train = labels[:train_num, :, :]
x_test = train[train_num:, :, :, :]
y_test = labels[train_num:, :, :]
# 获取模型
model = mobilenet_segnet(n_classes=NCLASSES, input_width=WIDTH, input_height=HEIGHT)
# 设置超参数
model.compile(optimizer=Adam(), loss=loss, metrics=['accuracy'])
# 保持模型路径
model_save_path = './checks/crosswalk.h5'
# 如果模型存在，则加载模型，在之前的基础上继续训练
if os.path.exists(model_save_path + '.index'):
print('--------------加载模型--------------')
model.load_weights(model_save_path)
# 设置模型模型要求
cp_callback = tf.keras.callbacks.ModelCheckpoint(
# 模型保存陆军
filepath=model_save_path,
# 只保存最好的一次
save_best_only=True,
# 保存模型
save_weights_only=True
)
model.fit(x_train, y_train, batch_size=4, epochs=50, validation_data=(x_test, y_test), validation_freq=1,
callbacks=[cp_callback])
model.summary()
3、模型预测代码实现
from net.segnet import mobilenet_segnet
from PIL import Image
import numpy as np
import copy
import os
# 给不同类别的目标，渲染颜色
class_colors = [[0, 0, 0],
[0, 255, 0]]
NCLASSES = 2
HEIGHT = 416
WIDTH = 416
# 加载模型
model = mobilenet_segnet(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH)
# 读取权重
model.load_weights("./check/classess.ckpt")
# 读取要预测的图片
imgs = os.listdir("./img/")
for jpg in imgs:
# 对图片格式进行处理
img = Image.open("./img/" + jpg)
old_img = copy.deepcopy(img)
orininal_h = np.array(img).shape[0]
orininal_w = np.array(img).shape[1]
img = img.resize((WIDTH, HEIGHT))
img = np.array(img)
img = img / 255
img = img.reshape(-1, HEIGHT, WIDTH, 3)
# 进行预测
pr = model.predict(img)[0]
# 调整通道，并返回通道上，每个轴上面的最大值
pr = pr.reshape((int(HEIGHT / 2), int(WIDTH / 2), NCLASSES)).argmax(axis=-1)
# 生成三通道黑色图片
seg_img = np.zeros((int(HEIGHT / 2), int(WIDTH / 2), 3))
colors = class_colors
for c in range(NCLASSES):
seg_img[:, :, 0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8')
seg_img[:, :, 1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8')
seg_img[:, :, 2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8')
seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h))
# 合成一副图像
image = Image.blend(old_img, seg_img, 0.3)
image.save("./img_out/" + jpg)
3、效果展示
原图
检测后的图片
总结
本文，在进行训练时，由于样本数和迭代次数较少，检测效果不是太好，希望大家学习之后可以改进，利用其他一些经典的卷积神经网络进行测试，比如ResNet、LeNet等去实现。也可以解码部分利用另外解码器进行解码。后续将更新一些新的方法。
训练代码和数据集地址：基于VGGnet模型的segnet斑马线划分-数据集文档类资源-CSDN下载
文章出处登录后可见！
已经登录？立即刷新
