源码解析-图像分割-FCN网络

FCN网络简介

FCN网络出自文章Fully Convolutional Networks for Semantic Segmentation。传统的CNN网络多以全链接层结尾，用于图像级别的分类问题。而语义分割(Semantic Segmentation)问题须要对图像中的每一个像素点进行分类。全链接层虽然可以获取全局语义信息，可是却会破坏像素点的局部位置信息。FCN网络率先使用卷积层替代全链接层，实现端到端的像素级预测，以下所示。

FCN网络源码

关于FCN网络的结构等知识请参考其余博客，本文主要介绍Tensorflow框架下的FCN代码，代码源自(shekkizh/FCN.tensorflow)，结构比较清晰，很是适合入门级的分析。代码主要分为inference（构建网络结构，返回预测结果）、train（计算梯度，更新参数）和main（数据操做，loss计算等）三个部分。

inference部分

def vgg_net(weights, image):
    """
    使用vgg_net做为基网络，使用预训练模型中的参数初始化网络，返回网络各层的结果
    :param weights:     预训练模型参数，做为权重和偏置初始值
    :param image:       网络的输入
    :return:            网络各层的计算结果
    """
    # vgg_net19网络的结构，舍弃了后面的全链接层
    layers = (
        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',

        'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',

        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3',
        'relu3_3', 'conv3_4', 'relu3_4', 'pool3',

        'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3',
        'relu4_3', 'conv4_4', 'relu4_4', 'pool4',

        'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3',
        'relu5_3', 'conv5_4', 'relu5_4'
    )

    net = {}
    current = image
    for i, name in enumerate(layers):
        kind = name[:4]     # 根据
        if kind == 'conv':
            kernels, bias = weights[i][0][0][0][0]      # 从预训练模型中提取出权重和偏置参数

            # matconvnet: weights are [width, height, in_channels, out_channels]
            # tensorflow: weights are [height, width, in_channels, out_channels]
            # 预训练模型mat文件中的存储格式位[W,H,N,C]，tensorflow中对卷积核的格式应为[N,C,H,W]
            kernels = utils.get_variable(np.transpose(kernels, (1, 0, 2, 3)), name=name + "_w")
            bias = utils.get_variable(bias.reshape(-1), name=name + "_b")

            '''使用weights的值建立一个变量
            若是不想使用预训练模型中的值，可使用截断正态分布来初始化权重参数，使用常量初始化偏置参数
            def get_variable(weights, name):
                init = tf.constant_initializer(weights, dtype=tf.float32)
                var = tf.get_variable(name=name, initializer=init,  shape=weights.shape)
                return var
            '''

            current = utils.conv2d_basic(current, kernels, bias)    # stride=1的卷积
        elif kind == 'relu':
            current = tf.nn.relu(current, name=name)
            if FLAGS.debug:
                # 将current加入到tensorboard中，用于可视化参数在训练过程当中的分布及变化状况
                utils.add_activation_summary(current)
        elif kind == 'pool':
            current = utils.avg_pool_2x2(current)   # 2*2平均池化
        net[name] = current     # 保存该层的输出结果

    return net


def inference(image, keep_prob):
    """
    计算FCN网络的输出
    Semantic segmentation network definition
    :param image: 输入的彩色图像，范围0-255，shape  =[N, IMAGE_SIZE, IMAGE_SIZE, 3]
    :param keep_prob: dropout比例
    :return:
    """

    print("setting up vgg initialized conv layers ...")
    # 从FLAGS.model_dir下读取imagenet-vgg-verydeep-19.mat文件，若是不存在则先根据MODEL_URL连接下载文件
    # 这个函数看着复杂，可是其实就是一个读取mat文件（scipy.io.loadmat函数）,返回结果的过程。
    model_data = utils.get_model_data(FLAGS.model_dir, MODEL_URL)

    mean = model_data['normalization'][0][0][0]
    mean_pixel = np.mean(mean, axis=(0, 1))     # 从预训练模型文件中提取图像文件的均值

    # 这个操做与预训练模型mat文件的存储格式相关，想了解它的格式可写个脚本读取，一步一步解析，但可能不值得
    # 参考：https://zhuanlan.zhihu.com/p/28897952
    weights = np.squeeze(model_data['layers'])

    processed_image = utils.process_image(image, mean_pixel)    # image - mean_pixel，减去均值

    with tf.variable_scope("inference"):
        image_net = vgg_net(weights, processed_image)   # 计算VGG的各层输出结果
        conv_final_layer = image_net["conv5_3"]         # 提取conv5_3层的结果，图像分辨率为输入图像的1/16

        pool5 = utils.max_pool_2x2(conv_final_layer)    # max pool，分辨率再降

        # 下面这种代码块包括建立变量，而后计算。推荐使用TensorFlow-Slim代码会简洁不少
        W6 = utils.weight_variable([7, 7, 512, 4096], name="W6")    # conv7*7-4096, relu, dropout
        b6 = utils.bias_variable([4096], name="b6")
        conv6 = utils.conv2d_basic(pool5, W6, b6)
        relu6 = tf.nn.relu(conv6, name="relu6")
        if FLAGS.debug:
            utils.add_activation_summary(relu6)
        relu_dropout6 = tf.nn.dropout(relu6, keep_prob=keep_prob)

        W7 = utils.weight_variable([1, 1, 4096, 4096], name="W7")   # conv1*1-4096, relu, dropout
        b7 = utils.bias_variable([4096], name="b7")
        conv7 = utils.conv2d_basic(relu_dropout6, W7, b7)
        relu7 = tf.nn.relu(conv7, name="relu7")
        if FLAGS.debug:
            utils.add_activation_summary(relu7)
        relu_dropout7 = tf.nn.dropout(relu7, keep_prob=keep_prob)

        W8 = utils.weight_variable([1, 1, 4096, NUM_OF_CLASSESS], name="W8")    # conv1*1-4096, relu, dropout
        b8 = utils.bias_variable([NUM_OF_CLASSESS], name="b8")
        conv8 = utils.conv2d_basic(relu_dropout7, W8, b8)
        # annotation_pred1 = tf.argmax(conv8, dimension=3, name="prediction1")

        '''FCN网络中conv8已是预测的结果，可是此时图像的分辨率为原始的1/32。为了得到更加准确的边界信息，
        直接进行线性插值上采样获得的结果并很差，所以使用转置卷积并与中间层的结果融合，提升输出的分辨率。
        （也可以使用插值上采样加3*3卷积替代转置卷积）'''

        # now to upscale to actual image size
        # 与VGG中间层pool4的结果融合，分辨率从1/32变为1/16
        deconv_shape1 = image_net["pool4"].get_shape()                  # pool4的尺寸
        W_t1 = utils.weight_variable([4, 4, deconv_shape1[3].value, NUM_OF_CLASSESS], name="W_t1")
        b_t1 = utils.bias_variable([deconv_shape1[3].value], name="b_t1")
        # k_size=4, stride=2，则图像分辨率扩大2倍
        conv_t1 = utils.conv2d_transpose_strided(conv8, W_t1, b_t1, output_shape=tf.shape(image_net["pool4"]))
        fuse_1 = tf.add(conv_t1, image_net["pool4"], name="fuse_1")     # 与pool4的结果相加

        # 与VGG中间层pool3的结果融合，分辨率从1/16变为1/8
        deconv_shape2 = image_net["pool3"].get_shape()
        W_t2 = utils.weight_variable([4, 4, deconv_shape2[3].value, deconv_shape1[3].value], name="W_t2")
        b_t2 = utils.bias_variable([deconv_shape2[3].value], name="b_t2")
        # k_size=4, stride=2，则图像分辨率扩大2倍
        conv_t2 = utils.conv2d_transpose_strided(fuse_1, W_t2, b_t2, output_shape=tf.shape(image_net["pool3"]))
        fuse_2 = tf.add(conv_t2, image_net["pool3"], name="fuse_2")     # 与pool3的结果相加

        # 还原到原始图像大小，分辨率从1/8变为1/1
        shape = tf.shape(image)
        deconv_shape3 = tf.stack([shape[0], shape[1], shape[2], NUM_OF_CLASSESS])
        W_t3 = utils.weight_variable([16, 16, NUM_OF_CLASSESS, deconv_shape2[3].value], name="W_t3")
        b_t3 = utils.bias_variable([NUM_OF_CLASSESS], name="b_t3")
        # k_size=16, stride=8，则图像分辨率扩大8倍（SAME模式下，分辨率变化仅与stride有关）
        conv_t3 = utils.conv2d_transpose_strided(fuse_2, W_t3, b_t3, output_shape=deconv_shape3, stride=8)

        # conv_t3.shape=[N, H, W, NUM_OF_CLASSESS]
        # 取NUM_OF_CLASSESS中值最大的做为预测的分类结果，annotation_pred.shape=[N,H,W]
        annotation_pred = tf.argmax(conv_t3, dimension=3, name="prediction")

    return tf.expand_dims(annotation_pred, dim=3), conv_t3

1.FLAGS.debug==True中的内容为TensorBoard相关的，主要是为了可视化（变量分布、变量变化状况等），前期可先忽略这部分代码。
2.转置卷积容易出现棋盘格效应，避免的方法为k_size为stride的整数倍关系，也可用上采样加卷积替代转置卷积。
3.与中间层的融合，能够用tf.add()直接相加，也可使用tf.concat()拼接在一块儿。

train部分

def train(loss_val, var_list):
    """
    计算并更新梯度。Optimizer.minimize()函数中直接包含compute_gradients()和apply_gradients()两步
    :param loss_val: loss值
    :param var_list: 需更新的变量，通常都是用tf.trainable_variables()获取全部需训练的变量
    :return: 梯度更新操做
    """
    optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)             # 建立Adam优化器
    grads = optimizer.compute_gradients(loss_val, var_list=var_list)    # 计算变量的梯度
    # 此处可增长tf.clip_by_value()操做，对梯度进行修建，避免梯度消失或者爆炸之类的状况
    if FLAGS.debug:
        # print(len(var_list))
        for grad, var in grads:
            utils.add_gradient_summary(grad, var)       # 将变量加到tensorboard中可视化
    # 更新梯度。
    return optimizer.apply_gradients(grads)

1.这部分能改动的很少，可选择其余的优化方法（Momentum、Adagrad等），也可以使用学习率衰减（指数衰减、多项式衰减等）。

main部分

def main(argv=None):
    # 设置dropout，输入图像，真实值的占位，实际训练或预测时须要传入sess.run()中
    keep_probability = tf.placeholder(tf.float32, name="keep_probabilty")
    image = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 3], name="input_image")
    annotation = tf.placeholder(tf.int32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 1], name="annotation")

    # pred_annotation.shape=[N, H, W, 1]，网络的预测结果，像素对应的得分值最大的类别
    # logits.shape=[N, H, W, NUM_OF_CLASSESS]，网络每一个像素对应每类标签的得分值
    pred_annotation, logits = inference(image, keep_probability)    # 计算FCN网络的预测结果

    tf.summary.image("input_image", image, max_outputs=2)
    tf.summary.image("ground_truth", tf.cast(annotation, tf.uint8), max_outputs=2)
    tf.summary.image("pred_annotation", tf.cast(pred_annotation, tf.uint8), max_outputs=2)  # 仍是可视化

    # 计算logits与annotation的交叉熵
    '''sparse_softmax_cross_entropy_with_logits()函数要求labels.shape=[d_0, d_1, ..., d_{r-1}],
    logits.shape=[d_0, d_1, ..., d_{r-1}, num_classes]。该函数比softmax_cross_entropy_with_logits()多一个
    标签稀疏化的步骤，计算的时候会先将labels进行稀疏化(one-hot编码，例如labels=[3]的10分类问题，
    会将其转化为[0,0,0,1,0,0,0,0,0,0])，而后再计算softmax值，再计算交叉熵，返回每一个神经元的结果，形状与labels一致'''
    loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
                                                                          labels=tf.squeeze(annotation, squeeze_dims=[3]),
                                                                          name="entropy")))
    loss_summary = tf.summary.scalar("entropy", loss)

    trainable_var = tf.trainable_variables()    # 获取全部的可训练变量
    if FLAGS.debug:
        for var in trainable_var:
            utils.add_to_regularization_and_summary(var)
    train_op = train(loss, trainable_var)       # 计算更新梯度

    print("Setting up summary op...")
    summary_op = tf.summary.merge_all()         # 汇总全部的summary

    print("Setting up image reader...")
    train_records, valid_records = scene_parsing.read_dataset(FLAGS.data_dir)   # 读取数据集，此处只是获取图像和标签的路径名
    print(len(train_records))
    print(len(valid_records))

    # 此处是的数据集操做了解一下就行。大体为如下步骤：获取文件名，而后读取图像，调整图像大小，按照序列每次取batch_size个图像
    # 做为网络的输入。一个epoch结束后，打乱序列后从新选取数据做为输入。
    # 此处读取所有训练集和验证集的图像数据，加载到内存中。可是若是内存有限或者数据集很是大，能够在实际输入时再读取图像数据。
    print("Setting up dataset reader")
    image_options = {'resize': True, 'resize_size': IMAGE_SIZE}     # 图像读取时需进行缩放
    if FLAGS.mode == 'train':
        train_dataset_reader = dataset.BatchDatset(train_records, image_options)
    validation_dataset_reader = dataset.BatchDatset(valid_records, image_options)

    sess = tf.Session()

    print("Setting up Saver...")
    saver = tf.train.Saver()

    # create two summary writers to show training loss and validation loss in the same graph
    # need to create two folders 'train' and 'validation' inside FLAGS.logs_dir
    train_writer = tf.summary.FileWriter(FLAGS.logs_dir + '/train', sess.graph)     # 用于将summary写入到文件中
    validation_writer = tf.summary.FileWriter(FLAGS.logs_dir + '/validation')

    sess.run(tf.global_variables_initializer())     # 变量初始化
    ckpt = tf.train.get_checkpoint_state(FLAGS.logs_dir)
    if ckpt and ckpt.model_checkpoint_path:         # 如何存在模型文件，则加载文件还原当前sess中的变量
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("Model restored...")

    # 训练模式
    if FLAGS.mode == "train":
        for itr in xrange(MAX_ITERATION):
            # 从训练集中读取FLAGS.batch_size大小的图像和标签数据，函数内部会自行调整读取的图像序列
            train_images, train_annotations = train_dataset_reader.next_batch(FLAGS.batch_size)
            feed_dict = {image: train_images, annotation: train_annotations, keep_probability: 0.85}

            sess.run(train_op, feed_dict=feed_dict)         # 训练

            if itr % 10 == 0:   # 每10次计算当前的loss，并保存到文件中
                train_loss, summary_str = sess.run([loss, loss_summary], feed_dict=feed_dict)
                print("Step: %d, Train_loss:%g" % (itr, train_loss))
                train_writer.add_summary(summary_str, itr)  # 将summary_str写入到文件中

            if itr % 500 == 0:  # 每500次使用验证集中的数据计算当前网络的loss，并保存到文件中
                valid_images, valid_annotations = validation_dataset_reader.next_batch(FLAGS.batch_size)
                valid_loss, summary_sva = sess.run([loss, loss_summary],
                                                   feed_dict={image: valid_images, annotation: valid_annotations,
                                                              keep_probability: 1.0})
                print("%s ---> Validation_loss: %g" % (datetime.datetime.now(), valid_loss))

                # add validation loss to TensorBoard
                validation_writer.add_summary(summary_sva, itr)
                saver.save(sess, FLAGS.logs_dir + "model.ckpt", itr)
    # 测试模式，
    elif FLAGS.mode == "visualize":
        # 从测试集中随机获取FLAGS.batch_size个图像数据
        valid_images, valid_annotations = validation_dataset_reader.get_random_batch(FLAGS.batch_size)
        pred = sess.run(pred_annotation, feed_dict={image: valid_images, annotation: valid_annotations,
                                                    keep_probability: 1.0})     # 计算预测结果
        valid_annotations = np.squeeze(valid_annotations, axis=3)               # 真实值
        pred = np.squeeze(pred, axis=3)

        # 保存原始图像、真实分割图像和预测的分割图像
        for itr in range(FLAGS.batch_size):
            utils.save_image(valid_images[itr].astype(np.uint8), FLAGS.logs_dir, name="inp_" + str(5+itr))
            utils.save_image(valid_annotations[itr].astype(np.uint8), FLAGS.logs_dir, name="gt_" + str(5+itr))
            utils.save_image(pred[itr].astype(np.uint8), FLAGS.logs_dir, name="pred_" + str(5+itr))
            print("Saved image: %d" % itr)

1.图像数据操做，建议使用opencv库进行处理，代码中的scipy.misc好像最近版本的兼容性有点问题。
2.图像分割的输入数据和标签数据都是图像，操做起来比较简单，不像目标检测那样还要对标签框进行复杂的转换才能计算损失值。能够根据数据的存放方式，本身写一个图像读取，扩充（裁剪、反转、通道变换等）的类，建议本身尝试写一个熟悉一下图像处理的基本知识。