源碼解析-圖像分割-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.圖像分割的輸入數據和標籤數據都是圖像，操做起來比較簡單，不像目標檢測那樣還要對標籤框進行復雜的轉換才能計算損失值。能夠根據數據的存放方式，本身寫一個圖像讀取，擴充（裁剪、反轉、通道變換等）的類，建議本身嘗試寫一個熟悉一下圖像處理的基本知識。