小生不财-语义分割2：训练自己的segnet模型（斑马线划分）_训练segnet

---

前言

在第一部分完成了主干网络VGGnet还有segnet模型，这节主要简单讲一下模型的训练以及预测过程。如有遗忘参考小生不财-语义分割1：基于VGGNet模型的segnet讲解

数据集以及标签详解

语义分割模型训练的文件分为两部分。

第一部分是原图，像这样：

 第二部分标签，像这样：

 当你们看到这个标签的时候你们会说，我靠，你给我看的什么辣鸡，全黑的算什么标签，其实并不是这样的，这个标签看起来全黑，但是实际上在斑马线的部分其RGB三个通道的值都是1。

其实给你们换一个图你们就可以更明显的看到了。

这是它的标签。

为什么这里的标签看起来就清楚的多呢，因为在voc中，其一共需要分21类，所以火车的RGB的值可能都大于10了，当然看得见。所以，在训练集中，如果像本文一样分两类，那么背景的RGB就是000，斑马线的RGB就是111，如果分多类，那么还会存在222，333，444这样的。这说明其属于不同的类。

二、代码训练

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, label

2、完整训练代码

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)

三、效果展示

原始图片

检测后图片

 

---

总结

        本文，在进行训练时，由于样本数和迭代次数较少，检测效果不是太好，希望大家学习之后可以改进，利用其他一些经典的卷积神经网络进行测试，比如ResNet、LeNet等去实现。也可以解码部分利用另外解码器进行解码。后续将更新一些新的方法。

训练代码和数据集地址：基于VGGnet模型的segnet斑马线划分-数据集文档类资源-CSDN下载