卷積神經網絡應用:基於Tensorflow的CNN/CRF圖像分割技術

卷積神經網絡應用:基於Tensorflow的CNN/CRF圖像分割技術

【方向】 2016-12-26 14:27:03 瀏覽7815 評論0

Image TensorFlow 卷積神經網絡 條件隨機場 後處理 圖像分割 上採樣

摘要: 本篇文章驗證了卷積神經網絡應用於圖像分割領域時存在的一個問題——粗糙的分割結果。根據像素間交叉熵損失的定義,我們在簡化的場景下進行了模型的訓練,並使用後向傳播來更新權重。我們使用條件隨機場(CRFs)來解決分割結果粗糙的問題,並取得了很好的效果。本文中的代碼註釋詳細、功能完善,也便於讀者閱讀。

這是一篇翻譯文章。介紹了一種基於最近發佈的TF-Slim庫與預訓練模型來進行圖像分割的方法。本篇文章的內容包括基於條件隨機場的模型訓練與後處理過程。

引言

在之前的文章中,我們實現了上採樣操作,並通過將其與scikit-image庫中的對應實現作比較,以確保上採樣過程的正確性。更具體地說,我們實現了論文《Fully convolutional networks for semantic segmentation》中描述的FCN-32分割網絡。

在本篇文章中,我們將進行一個簡單的訓練:我們將從PASCAL VOC數據集中選取一張樣本圖像以及它的標註信息,基於樣本圖像與標註訓練來我們的網絡,並在同樣的圖像上測試我們訓練好的網絡。之所以這樣做,是因爲這樣我們就可以使用CPU來訓練模型——只需要10次迭代就能完成訓練。本篇文章的另一個目的在於,證明我們這個網絡(FCN-32s)的分割結果是很粗糙的——即使在用於訓練的圖像上進行分割,結果也同樣如此。在本篇文章中,我們使用條件隨機場的後處理階段來解決這個問題,其綜合考慮圖像中的RGB純色特徵以及模型輸出的預測概率,進而改善圖像的分割效果。總的來說,我們得到了改善後的分割結果。我們刻意在文章中使用了非常簡單的訓練環境。類似的分割方法在Chen等人發表的論文《Semantic Image Segmentation with Deep Convolutional Nets and Fully Connected CRFs》中也有所提及。請注意,本篇文章中的訓練環境僅爲了說明FCN-32s模型的侷限,如果要在真實場景中進行訓練,我們建議讀者參考這篇論文《Fully convolutional networks for semantic segmentation》。

本篇文章使用Jupyter Notebook創建。在每個代碼塊後,你可以看到模型的評估結果。你也可以從這裏獲得本文對應的筆記文檔。本篇文章的部分內容也借鑑了TF-Slim庫的演示教程

準備階段

爲了運行下面的代碼,你需要安裝Tensorflow。我使用的版本是r0.12。你可能需要使用Tensorflow的這個分支

在這篇教程中,我也使用了scikit-imag庫、numpy以及其他的一些依賴。你可以通過下載Anaconda的Python軟件包來安裝這些軟件。

此外,你還需要按照之前幾篇文章介紹的步驟來配置你的訓練環境——包括如何下載VGG-16模型以及其他所有的必需操作。

上採樣輔助函數與圖像加載

在這一部分中,我們定義前一篇文章使用的輔助函數。如果你還記得的話,我們使用了上採樣技術對我們從網絡中獲得的下采樣預測結果進行採樣。我們從VGG-16模型使用的最大值池化層中獲得模型的下采樣預測結果。

我們也爲圖像與各個實際分割區域的加載操作編寫了程序。這些代碼加入了足夠的註釋,所以不必擔心無法理解它。

import numpy as np

def get_kernel_size(factor):
    """
    給定所需的上採樣因子,確定卷積核的大小
    """
    return 2 * factor - factor % 2


def upsample_filt(size):
    """
    創建一個給定(h, w) 大小的適用於上採樣過程的二維雙線性卷積核
    """
    factor = (size + 1) // 2
    if size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:size, :size]
    return (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)


def bilinear_upsample_weights(factor, number_of_classes):
    """
    使用雙線性卷積核,爲轉置卷積創建權重矩陣
    初始化
    """

    filter_size = get_kernel_size(factor)

    weights = np.zeros((filter_size,
                        filter_size,
                        number_of_classes,
                        number_of_classes), dtype=np.float32)

    upsample_kernel = upsample_filt(filter_size)

    for i in xrange(number_of_classes):

        weights[:, :, i, i] = upsample_kernel

    return weights
%matplotlib inline

from __future__ import division

import os
import sys
import tensorflow as tf
import skimage.io as io
import numpy as np

os.environ["CUDA_VISIBLE_DEVICES"] = '1'
sys.path.append("/home/dpakhom1/workspace/my_models/slim/")
checkpoints_dir = '/home/dpakhom1/checkpoints'

image_filename = 'cat.jpg'
annotation_filename = 'cat_annotation.png'

image_filename_placeholder = tf.placeholder(tf.string)
annotation_filename_placeholder = tf.placeholder(tf.string)
is_training_placeholder = tf.placeholder(tf.bool)

feed_dict_to_use = {image_filename_placeholder: image_filename,
                    annotation_filename_placeholder: annotation_filename,
                    is_training_placeholder: True}

image_tensor = tf.read_file(image_filename_placeholder)
annotation_tensor = tf.read_file(annotation_filename_placeholder)

image_tensor = tf.image.decode_jpeg(image_tensor, channels=3)
annotation_tensor = tf.image.decode_png(annotation_tensor, channels=1)

# 對於每個類別,將其設置爲1而不是一個數字——我們在後續計算
# 交叉熵時會用到。有的時候,實際分割區域的掩碼可以有很多值,
# 不僅僅爲0和1
class_labels_tensor = tf.equal(annotation_tensor, 1)
background_labels_tensor = tf.not_equal(annotation_tensor, 1)

#將布爾值轉換爲浮點數——這樣才能正確地計算交叉熵損失
bit_mask_class = tf.to_float(class_labels_tensor)
bit_mask_background = tf.to_float(background_labels_tensor)

combined_mask = tf.concat(concat_dim=2, values=[bit_mask_class,
                                                bit_mask_background])

# 調整輸入數據的大小,使其與tf.softmax_cross_entropy_with_logits中
# [batch_size, num_classes]的要求保持一致
flat_labels = tf.reshape(tensor=combined_mask, shape=(-1, 2))

損失函數定義與基於Adam優化算法的訓練

在這一部分中,我們把準備好的所有東西整合到一起:爲我們的網絡添加上採樣層,定義可微分求導的損失函數,並進行模型的訓練。

參照論文《Fully convolutional networks for semantic segmentation》,我們將模型的損失定義爲像素間的交叉熵。我們之所以這樣定義,是因爲在上採樣後我們可以得到與輸入圖像尺寸相同的預測結果,進而我們可以將獲得的分割結果與實際的各個分割區域進行比較:

4a8f9e64c859fedd6c2346716e3a257f5ac383af

其中,N表示像素的數量,K表示類別的數量,變量Tnk表示第n個像素對應的實際分割區域,實際區域以1至K的數字表示,變量Ynk爲我們模型的預測結果(使用softmax函數來規範輸出)。

對於這種情況,由於Adam優化器僅調整較少的參數就能使模型取得很好的效果,我們將其應用到了模型的訓練當中。

在這個特殊的情況下,我們使用一張圖像來訓練模型並評估模型的效果——與真實場景相比,這種訓練方式非常簡單。我們這樣做是爲了展示這種方法的不足——只爲了說明這種方法糟糕的定位能力。如果在這種簡單場景下這種方法的效果尚且如此,那麼在未訓練的圖像上,這種方法只會產生類似的更糟糕的結果。

import numpy as np
import tensorflow as tf
import sys
import os
from matplotlib import pyplot as plt

fig_size = [15, 4]
plt.rcParams["figure.figsize"] = fig_size

import urllib2

slim = tf.contrib.slim

from nets import vgg
from preprocessing import vgg_preprocessing

# 加載像素均值以及爲每個像素進行減法運算的函數
from preprocessing.vgg_preprocessing import (_mean_image_subtraction,
                                            _R_MEAN, _G_MEAN, _B_MEAN)

upsample_factor = 32
number_of_classes = 2
log_folder = '/home/dpakhom1/tf_projects/segmentation/log_folder'

vgg_checkpoint_path = os.path.join(checkpoints_dir, 'vgg_16.ckpt')

# 在與像素均值做差前,將圖像轉換至float32類型
image_float = tf.to_float(image_tensor, name='ToFloat')

# 將每個像素的具體數值與像素均值做差
mean_centered_image = _mean_image_subtraction(image_float,
                                          [_R_MEAN, _G_MEAN, _B_MEAN])

processed_images = tf.expand_dims(mean_centered_image, 0)

upsample_filter_np = bilinear_upsample_weights(upsample_factor,
                                               number_of_classes)

upsample_filter_tensor = tf.constant(upsample_filter_np)

# 定義將要使用的模型——指定在最後一層僅使用兩個類別
with slim.arg_scope(vgg.vgg_arg_scope()):

    logits, end_points = vgg.vgg_16(processed_images,
                           num_classes=2,
                           is_training=is_training_placeholder,
                           spatial_squeeze=False,
                           fc_conv_padding='SAME')

downsampled_logits_shape = tf.shape(logits)

# 計算上採樣數據的輸出大小
upsampled_logits_shape = tf.pack([
                                  downsampled_logits_shape[0],
                                  downsampled_logits_shape[1] * upsample_factor,
                                  downsampled_logits_shape[2] * upsample_factor,
                                  downsampled_logits_shape[3]
                                 ])

# 進行上採樣處理
upsampled_logits = tf.nn.conv2d_transpose(logits, upsample_filter_tensor,
                                 output_shape=upsampled_logits_shape,
                                 strides=[1, upsample_factor, upsample_factor, 1])

# 展開預測結果,以便於我們計算每個像素的交叉熵,並獲得交叉熵的總和
flat_logits = tf.reshape(tensor=upsampled_logits, shape=(-1, number_of_classes))

cross_entropies = tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits,
                                                          labels=flat_labels)

cross_entropy_sum = tf.reduce_sum(cross_entropies)

# 獲得每個像素的最終預測結果——請注意,在這種情況下我們並不需要
# 使用softmax,因爲我們只需要得到最終的決策。如果我們還需要各
# 個類別的概率,那麼我們必須應用softmax
pred = tf.argmax(upsampled_logits, dimension=3)

probabilities = tf.nn.softmax(upsampled_logits)

# 在這裏我們定義了一個優化器,並添加了所有將要創建至命名
# 空間'adam_vars'下的變量。這樣做有利於我們後續輕鬆地訪問
# 它們。這些變量供adam優化器使用,並且與vgg模型中的變量無關

# 我們還獲得了每個變量的梯度數據
# 這樣,我們可以在tensorboard中可視化這些變量
# optimizer.compute_gradients與optimizer.apply_gradients
# 等價於執行:
# train_step = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(cross_entropy_sum)
with tf.variable_scope("adam_vars"):
    optimizer = tf.train.AdamOptimizer(learning_rate=0.0001)
    gradients = optimizer.compute_gradients(loss=cross_entropy_sum)

    for grad_var_pair in gradients:

        current_variable = grad_var_pair[1]
        current_gradient = grad_var_pair[0]

        # 替換原始變量名中的一些字符
        # tensorboard不支持':'符號
        gradient_name_to_save = current_variable.name.replace(":", "_")

        # 得到每一層的梯度直方圖,並隨後在tensorboard中可視化這些數據
        tensorboard
        tf.summary.histogram(gradient_name_to_save, current_gradient) 

    train_step = optimizer.apply_gradients(grads_and_vars=gradients)

# 在這裏,我們定義了一個函數,調用時會從VGG模型檢查點中讀取權重數據,並加載至變量中。
# 我們從負責類別預測的最後一層中剔除了權重。我們這樣做是因爲我們將有不同數量的
# 類別進行預測,我們不能在初始化時使用原先的類別。
vgg_except_fc8_weights = slim.get_variables_to_restore(exclude=['vgg_16/fc8', 'adam_vars'])

# 這裏我們得到了網絡中最後一層的權重變量
# 正如我們看到的,VGG最初訓練的類別數量與我們實際的類別數量
# 並不相同——在我們的情況下,總共只有兩類
vgg_fc8_weights = slim.get_variables_to_restore(include=['vgg_16/fc8'])

adam_optimizer_variables = slim.get_variables_to_restore(include=['adam_vars'])

# 爲模型損失添加一個summary OP——以便我們可以在tensorboard中看到它
tf.summary.scalar('cross_entropy_loss', cross_entropy_sum)

# 將所有summary OP合併至一個OP總
# 在運行程序時生成字符串
merged_summary_op = tf.summary.merge_all()

# 創建一個summary writer——用於將所有日誌寫入到一個特定文件中
# 這個文件後續可以由tensorboard讀取
summary_string_writer = tf.summary.FileWriter(log_folder)

# 如果日誌文件夾尚未存在,則創建一個新的文件夾
if not os.path.exists(log_folder):
    os.makedirs(log_folder)

# 創建一個OP,對VGG模型中各權重變量進行初始化操作
read_vgg_weights_except_fc8_func = slim.assign_from_checkpoint_fn(
                                   vgg_checkpoint_path,
                                   vgg_except_fc8_weights)

# 針對新的fc8層權重數據的初始化器——僅包括兩類
vgg_fc8_weights_initializer = tf.variables_initializer(vgg_fc8_weights)

# adam變量的初始化器
optimization_variables_initializer = tf.variables_initializer(adam_optimizer_variables)

with tf.Session() as sess:
    # 運行初始化器
    read_vgg_weights_except_fc8_func(sess)
    sess.run(vgg_fc8_weights_initializer)
    sess.run(optimization_variables_initializer)

    train_image, train_annotation = sess.run([image_tensor, annotation_tensor],
                                              feed_dict=feed_dict_to_use)

    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
    ax1.imshow(train_image)
    ax1.set_title('Input image')
    probability_graph = ax2.imshow(np.dstack((train_annotation,)*3)*100)
    ax2.set_title('Input Ground-Truth Annotation')
    plt.show()

    # 執行10次迭代
    for i in range(10):

        loss, summary_string = sess.run([cross_entropy_sum, merged_summary_op],
                                        feed_dict=feed_dict_to_use)

        sess.run(train_step, feed_dict=feed_dict_to_use)

        pred_np, probabilities_np = sess.run([pred, probabilities],
                                              feed_dict=feed_dict_to_use)

        summary_string_writer.add_summary(summary_string, i)

        cmap = plt.get_cmap('bwr')

        f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
        ax1.imshow(np.uint8(pred_np.squeeze() != 1), vmax=1.5, vmin=-0.4, cmap=cmap)
        ax1.set_title('Argmax. Iteration # ' + str(i))
        probability_graph = ax2.imshow(probabilities_np.squeeze()[:, :, 0])
        ax2.set_title('Probability of the Class. Iteration # ' + str(i))

        plt.colorbar(probability_graph)
        plt.show()

        print("Current Loss: " +  str(loss))

    feed_dict_to_use[is_training_placeholder] = False

    final_predictions, final_probabilities, final_loss = sess.run([pred,
                                                                   probabilities,
                                                                   cross_entropy_sum],
                                                         feed_dict=feed_dict_to_use)


    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)

    ax1.imshow(np.uint8(final_predictions.squeeze() != 1),
               vmax=1.5,
               vmin=-0.4,
               cmap=cmap)

    ax1.set_title('Final Argmax')


    probability_graph = ax2.imshow(final_probabilities.squeeze()[:, :, 0])
    ax2.set_title('Final Probability of the Class')
    plt.colorbar(probability_graph)

    plt.show()

    print("Final Loss: " +  str(final_loss))

summary_string_writer.close()
79727851db94aa0184782ea940b4723f29c735a3
bf3053a03868c2e7d63ea7a6f1c14c39ec7f1e2a

Current Loss: 201433.0

462d07f37ed83c991e48aeef025135bd9f848f26

Current Loss: 245565.0

2f867f2909ebb4aca17cf3937c41d6d15890cc4d

Current Loss: 135906.0

a2a0dc74a4c572f6bc8efcb0260ebdd1e639ba8a

Current Loss: 183353.0

fe55c967b1354af4bc828e9fea0f5eb6b1b839aa

Current Loss: 48563.9

3a7237f9170269f1c96438da61e1ec5b7a15bfc8

Current Loss: 37925.8

d792b37dc429a2bcf9e2d03c67ff407542b50888

Current Loss: 33199.1

49e91bcb453388860a66f5a9a572f25e4acfa274

Current Loss: 26540.3

1f639e4ba2a5936d84319584c34325d232b43ba8

Current Loss: 23658.0

3dd2c899e849908d7061c5204b7ea1bc79c33e5f

Current Loss: 29404.9

d352cff7f9f25d1732a3dca4f8c989f7526de8ab

Final Loss: 18177.5

正如你所看到的,結果非常粗糙——而且,這還是我們使用模型訓練同一張圖像來運行網絡預測的結果。這在圖像分割中一個非常常見的問題——分割結果通常很粗糙。有幾種不同的方法可以解決此類問題——其中之一便是使用跳躍連接。主要思路是根據融合網絡中不同層的預測結果來確定模型的最終預測結果。由於在網絡較靠前的層次中下采樣因子較小,所以,可以根據這些層來確定預測結果,進而取得更好的定位效果。Long等人發表的論文《Fully convolutional networks for semantic segmentation》介紹了這種方法。基於這種方法,研究人員設計了FCN-16s與FCN-8s架構。

另一種方法則基於帶孔卷積與全連接條件隨機場。Chen等人發表的論文《Semantic Image Segmentation with Deep Convolutional Nets and Fully Connected CRFs》介紹了這種方法。在本篇文章中,我們將僅使用條件隨機場後處理階段來展示它對模型效果的提高。

還有一點需要注意的是,當前模型訓練時在全連接層(我們映射到卷積層的全連接層)中應用了Dropout技術,這種方法在Srivastava等人發表的論文《Dropout: a simple way to prevent neural networks from overfitting》中有所提及。Dropout是一種用於模型訓練的正則化技術。它有一個非常優秀的理論描述,而且實現起來也非常簡單:我們只需要在每個訓練步驟中隨機地選擇一定數量的神經元,僅根據這些神經元來進行推斷並後向傳播。但是,從理論角度來看,Dropout可以看作是通過權重共享來訓練一個稀疏網絡的集合,每個網絡僅進行很少次數的訓練。在測試階段,我們對所有這些網絡的預測結果求均值。在論文中作者表明,Dropout在線性迴歸情況下預期能夠取得與嶺迴歸相同的效果。在我們的具體情況中,Dropout僅用於全連接層(我們映射到卷積層的全連接層)。這也解釋了爲什麼最終模型的損失幾乎比最後一次迭代的損失小了兩倍——因爲在最後的推斷中,我們使用了損失的均值。

上面提供的代碼用於處理單張圖像,但你可以很容易地在整個數據集上運行這些程序。唯一需要調整的是,在每個迭代步驟中提供不同的圖像。這種訓練方式與論文《Fully convolutional networks for semantic segmentation》的做法完全一致,其中,論文作者採用了數量爲1的批處理進行訓練。

總的來說,我們可以看到我們的分割結果仍然很粗糙,需要執行一些額外的處理步驟。在下一節,我們將應用條件隨機場的後處理步驟使模型的分割粒度更細。

條件隨機場後處理

條件隨機場是圖模型的一種特定類型。在我們的應用場景中,條件隨機場有助於根據網絡的預測以及圖像原生的RGB特徵估計模型預測結果的後驗分佈。它通過最小化用戶定義的能量函數來實現這一點。在我們的應用場景中,其效果與雙邊濾波器非常接近,雙邊濾波器綜合考慮了圖像中像素的空間鄰近性以及在RGB特徵空間(強度空間)中的相似性。

在一個非常簡單的層面上,這種方法使用RGB特徵來使圖像分割結果更加精準——例如,邊界通常表示爲極大的強度變化——這是一個關鍵因素,位於該邊界兩側的物體屬於不同的類別。這種方法也同樣對小的分割區域進行懲罰——例如,20像素或50像素大小的小區域通常不可能是一個正確的分割區域。物體通常由大的空間相鄰的區域表示。

下面你可以看到這個後處理步驟是如何影響我們的分割結果的。我們使用了論文《fully connected crfs with gaussian edge potentials》有效性推斷章節中描述的全連接條件隨機場對分割結果進行調整。

在這部分中,我使用了一個版本稍微有些老的全連接條件隨機場的實現,你可以在這裏找到它。

import sys

path = "/home/dpakhom1/dense_crf_python/"
sys.path.append(path)

import pydensecrf.densecrf as dcrf

from pydensecrf.utils import compute_unary, create_pairwise_bilateral, \
    create_pairwise_gaussian, softmax_to_unary

import skimage.io as io

image = train_image

softmax = final_probabilities.squeeze()

softmax = processed_probabilities.transpose((2, 0, 1))

# 輸入數據應爲概率值的負對數
# 你可以在softmax_to_unary函數的定義中找到更多信息
unary = softmax_to_unary(processed_probabilities)

# 輸入數據應爲C-連續的——我們使用了Cython封裝器
unary = np.ascontiguousarray(unary)

d = dcrf.DenseCRF(image.shape[0] * image.shape[1], 2)

d.setUnaryEnergy(unary)

# 潛在地對空間上相鄰的小塊分割區域進行懲罰——促使產生更多空間連續的分割區域
feats = create_pairwise_gaussian(sdims=(10, 10), shape=image.shape[:2])

d.addPairwiseEnergy(feats, compat=3,
                    kernel=dcrf.DIAG_KERNEL,
                    normalization=dcrf.NORMALIZE_SYMMETRIC)

# 這將創建與顏色相關的圖像特徵——因爲我們從卷積神經網絡中得到的分割結果非常粗糙,
# 我們可以使用局部的顏色特徵來改善分割結果
feats = create_pairwise_bilateral(sdims=(50, 50), schan=(20, 20, 20),
                                   img=image, chdim=2)

d.addPairwiseEnergy(feats, compat=10,
                     kernel=dcrf.DIAG_KERNEL,
                     normalization=dcrf.NORMALIZE_SYMMETRIC)
Q = d.inference(5)

res = np.argmax(Q, axis=0).reshape((image.shape[0], image.shape[1]))

cmap = plt.get_cmap('bwr')

f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
ax1.imshow(res, vmax=1.5, vmin=-0.4, cmap=cmap)
ax1.set_title('Segmentation with CRF post-processing')
probability_graph = ax2.imshow(np.dstack((train_annotation,)*3)*100)
ax2.set_title('Ground-Truth Annotation')
plt.show()
ba7b8385f296631a0afb1419f81b51ad0a1e8819

結論與討論

在本篇教程中,我們看到了卷積神經網絡應用於圖像分割領域時的一個缺陷——粗糙的分割結果。我們可以看到,這種情況是由於在VGG-16網絡架構中使用最大值池化層所導致的。

根據像素間交叉熵損失的定義,我們在簡化的場景下進行了模型的訓練,並使用後向傳播來更新權重。

我們使用條件隨機場(CRFs)來解決分割結果粗糙的問題,並取得了很好的效果。

以下爲譯文

**本文由北郵@愛可可-愛生活 老師推薦,阿里云云棲社區組織翻譯。

原標題《Image Segmentation with Tensorflow using CNNs and Conditional Random Fields》,作者:Daniil,譯者:6816816151**

文章爲簡譯,更爲詳細的內容,請查看原文


發佈了15 篇原創文章 · 獲贊 41 · 訪問量 12萬+
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章