1. 程式人生 > >基於Tensorflow和TF-Slim影象分割示例

基於Tensorflow和TF-Slim影象分割示例

引言

在上一篇文章裡(如何用TensorFlow和TF-Slim實現影象分類與分割),我們介紹瞭如何擷取圖片的中央區域,然後用標準的分類模型對圖片的類別進行預測。隨後,我們又介紹瞭如何將網路模型改為全卷積模式,對整張圖片進行預測。我們通過這種方法可以得到原始圖片的一張降取樣預測圖 —— 降取樣是由於網路結構含有最大池化層。這種預測圖可以被視為一種全圖的快速預測。你也可以把它看作是一種影象分割的方法,但並不算嚴格意義上的影象分割,因為標準網路模型是以影象分類的目標而訓練得到的。為了實現真正的影象分割任務,我們需要在影象分割的資料集上重新訓練模型,訓練的方法參照Long等人發表的論文《Fully convolutional networks for semantic segmentation》。影象分割常用的兩種資料集分別是:

Microsoft COCOPASCAL VOC

在這篇文章裡,我們將對預測圖進行上取樣,從而使得輸出的圖片與輸入圖片尺寸保持一致。值得注意的一點是不要把我們的方法與所謂的deconvolution混淆,它們分別是兩種不同的操作方法。我們的方法準確的來說應該叫fractionally strided convolution,或者稱之為微步長卷積吧。下文中,我們會適當介紹一些理論知識以便於理解。

也許大家現在會有一個問題:為什麼我們需要用微步長卷積做上取樣操作呢?為什麼不能選用其它的上取樣方法呢?簡而言之:我們需要把上取樣操作定義為模型的一層結構。為什麼需要增加這一層呢?因為我們的訓練資料上已經標註了影象分割的資訊,我們需要根據這些資訊來訓練模型。

眾所周知,網路模型的每一層結構需要進行三種操作:前向傳播、反向傳播和引數更新。為了實現轉置卷積的上取樣,我們需要對這三種操作都給出定義,然後才能開始訓練。

在文章的最後,我們將會實現上取樣的過程,並且將結果與scikit-image庫的實現做比較。具體說來,我們會實現《Fully convolutional networks for semantic segmentation》一文中的FCN-32影象分割網路。

本文是在jupyter notebook上完成的,所以大家也可以直接下載ipynb檔案執行。

環境配置

為了成功執行文中的程式碼,請先輸入一下命令,指定GPU和輸入檔案路徑。可以參考上一篇文章的介紹。

from __future__ import division
import sys
import os
import numpy as np
from matplotlib import pyplot as plt

os.environ["CUDA_VISIBLE_DEVICES"] = '0'
sys.path.append("/home/dpakhom1/workspace/models/slim")

# A place where you have downloaded a network checkpoint -- look at the previous post
checkpoints_dir = '/home/dpakhom1/checkpoints'
sys.path.append("/home/dpakhom1/workspace/models/slim")

圖片上取樣

圖片上取樣是一種特殊的重取樣方法。這篇論文中給出了重取樣的定義:

重取樣的思路就是用更多的樣本(又稱為插值或是上取樣)或者更少的樣本(又稱為抽樣或是降取樣)來重建連續的訊號。

換句話說,我們可以用已有的資料點估計連續訊號,然後根據重建的訊號取樣得到新的樣本點。就我們的問題而言,我們已有降取樣的預測圖 —— 它們就是訊號重建的資料來源。如果我們能估計出原始的訊號,那麼就可以從中取樣出更多的樣本,即實現了上取樣。

我們強調“估計”是因為重建得到的原始連續訊號效果可能並不太好。只有在某些條件限制下,訊號才能被完全重建。根據夏農取樣定律,只有取樣頻率不小於模擬訊號頻譜中最高頻率的2倍時,訊號才能不失真地還原。

但是,訊號還原的公式是什麼?我們參考這篇文章中給出的公式:

s(x) = sum_n s(nT) * sinc((x-nT)/T),
其中當x!=0時,sinc(x) = sin(pix)/(pix),當x=0時 sinc(x)=1。

取樣率為fs = 1/T,s(n*T)是連續訊號s(x)的取樣點,sinc(x)是重取樣的核。

維基百科對上述公式做了完美的解釋。

因此,我們需要將已有的資料點代入重取樣核函式並求和(為了便於理解,忽略了一些細節),以此來還原連續的訊號。重取樣的核函式並不一定用sinc函式。比如,也可以用雙線性重取樣核。這裡給出了許多例子。上述解釋的另外一個關鍵點是可以將核函式與Dirac脈衝訊號序列做卷積運算,權重等於樣本的值,它們在數學上是等價的。這兩種等價的訊號重建方式非常重要,因為它們有助於我們理解轉置卷積的原理,以及每個轉置卷積都有一種等價的卷積。

用scikit-image庫內建的函式進行影象上取樣運算。我們需要以此為參照,來驗證我們雙線性上取樣的實現方式是正確的。這裡是在合併到程式碼庫之前的雙線性上取樣實現的驗證過程。本文中的部分程式碼來源於此。接下去,我們會進行三倍上取樣,也就是說輸出圖片的尺寸是輸入圖片的三倍。

%matplotlib inline

from numpy import ogrid, repeat, newaxis

from skimage import io

# Generate image that will be used for test upsampling
# Number of channels is 3 -- we also treat the number of
# samples like the number of classes, because later on
# that will be used to upsample predictions from the network
imsize = 3
x, y = ogrid[:imsize, :imsize]
img = repeat((x + y)[..., newaxis], 3, 2) / float(imsize + imsize)
io.imshow(img, interpolation='none')

import skimage.transform

def upsample_skimage(factor, input_img):

    # Pad with 0 values, similar to how Tensorflow does it.
    # Order=1 is bilinear upsampling
    return skimage.transform.rescale(input_img,
                                     factor,
                                     mode='constant',
                                     cval=0,
                                     order=1)


upsampled_img_skimage = upsample_skimage(factor=3, input_img=img)
io.imshow(upsampled_img_skimage, interpolation='none')

轉置卷積

在Long的這篇論文裡,他提到可以用微步長卷積(轉置卷積)實現上取樣。但是,首先要理解轉置卷積的工作原理。我們先以普通的卷積為例,理解哪些引數會影響卷積結果的圖片尺寸。如果我們想實現逆向操作該怎麼辦 —— 輸入小尺寸的圖片,輸出大尺寸的圖片,並且保留圖片的連通性模式。下面是示例:

卷積是一類線性運算,因此它可以表示為矩陣的乘法。為了實現上述效果,我們只需對卷積運算的矩陣做轉置。這種運算不再屬於卷積,但仍能用卷積來表現。大家若想要更詳細瞭解轉置卷積的知識,推薦閱讀這篇論文和這份指南。

所以,如果定義了雙線性上取樣核,並且對圖片做了微步長卷積運算,我們就可以得到上取樣的輸出,它往往是網路結構的其中一層,能通過反向傳播運算。對FCN-32網路,我們選用雙線性上取樣核為初始化,網路模型可以在反向傳播過程中學會更合適的核。

為了便於理解程式碼,可以參考這個頁面閱讀。上取樣的factor引數等於轉置卷積的步長。上取樣核的定義為:2 * factor - factor % 2。

我們在Tensorflow裡用轉置卷積運算來定義雙線性插值。這部分運算放在CPU上,因為我們在後續篇幅中還會用到這段程式碼來進行記憶體開銷較大的運算,因此不適合使用GPU。插值運算結束後,將我們的結果與sciki-image的運算結果做比較。

from __future__ import division
import numpy as np
import tensorflow as tf


def get_kernel_size(factor):
    """
    Find the kernel size given the desired factor of upsampling.
    """
    return 2 * factor - factor % 2


def upsample_filt(size):
    """
    Make a 2D bilinear kernel suitable for upsampling of the given (h, w) size.
    """
    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):
    """
    Create weights matrix for transposed convolution with bilinear filter
    initialization.
    """

    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


def upsample_tf(factor, input_img):

    number_of_classes = input_img.shape[2]

    new_height = input_img.shape[0] * factor
    new_width = input_img.shape[1] * factor

    expanded_img = np.expand_dims(input_img, axis=0)

    with tf.Graph().as_default():
        with tf.Session() as sess:
            with tf.device("/cpu:0"):

                upsample_filt_pl = tf.placeholder(tf.float32)
                logits_pl = tf.placeholder(tf.float32)

                upsample_filter_np = bilinear_upsample_weights(factor,
                                        number_of_classes)

                res = tf.nn.conv2d_transpose(logits_pl, upsample_filt_pl,
                        output_shape=[1, new_height, new_width, number_of_classes],
                        strides=[1, factor, factor, 1])

                final_result = sess.run(res,
                                feed_dict={upsample_filt_pl: upsample_filter_np,
                                           logits_pl: expanded_img})

    return final_result.squeeze()

upsampled_img_tf = upsample_tf(factor=3, input_img=img)
io.imshow(upsampled_img_tf)

# Test if the results of upsampling are the same
np.allclose(upsampled_img_skimage, upsampled_img_tf)

True
for factor in xrange(2, 10):

    upsampled_img_skimage = upsample_skimage(factor=factor, input_img=img)
    upsampled_img_tf = upsample_tf(factor=factor, input_img=img)

    are_equal = np.allclose(upsampled_img_skimage, upsampled_img_tf)

    print("Check for factor {}: {}".format(factor, are_equal))
Check for factor 2: True
Check for factor 3: True
Check for factor 4: True
Check for factor 5: True
Check for factor 6: True
Check for factor 7: True
Check for factor 8: True
Check for factor 9: True

上取樣預測

將我們的上取樣結果轉為真正的預測內容。我們選用上一篇文章中用過的VGG-16模型進行分類,然後將上取樣方法應用於低解析度的模型預測結果上。

在執行下面的程式碼之前,我要先修改VGG-16模型程式碼的其中一行,防止它繼續減小圖片的尺寸。具體來說,要把7x7卷積層的padding選項改為SAME。對於作者來說,他們希望得到輸入圖片的唯一預測結果,而我們的圖片分割任務並不需要到這一步,否則上取樣得到的圖片與輸入圖片的尺寸還是不一致。經過上述修改之後問題就解決了。

由於引數特別多,執行下面的程式碼會佔用很多記憶體,大概15GB,所以要留有足夠的空間。

%matplotlib inline

from matplotlib import pyplot as plt

import numpy as np
import os
import tensorflow as tf
import urllib2

from datasets import imagenet
from nets import vgg
from preprocessing import vgg_preprocessing

checkpoints_dir = '/home/dpakhom1/checkpoints'

slim = tf.contrib.slim

# Load the mean pixel values and the function
# that performs the subtraction
from preprocessing.vgg_preprocessing import (_mean_image_subtraction,
                                            _R_MEAN, _G_MEAN, _B_MEAN)

slim = tf.contrib.slim
# Function to nicely print segmentation results with
# colorbar showing class names
def discrete_matshow(data, labels_names=[], title=""):

    fig_size = [7, 6]
    plt.rcParams["figure.figsize"] = fig_size

    #get discrete colormap
    cmap = plt.get_cmap('Paired', np.max(data)-np.min(data)+1)

    # set limits .5 outside true range
    mat = plt.matshow(data,
                      cmap=cmap,
                      vmin = np.min(data)-.5,
                      vmax = np.max(data)+.5)
    #tell the colorbar to tick at integers
    cax = plt.colorbar(mat,
                       ticks=np.arange(np.min(data),np.max(data)+1))

    # The names to be printed aside the colorbar
    if labels_names:
        cax.ax.set_yticklabels(labels_names)

    if title:
        plt.suptitle(title, fontsize=15, fontweight='bold')


with tf.Graph().as_default():

    url = ("https://upload.wikimedia.org/wikipedia/commons/d/d9/"
           "First_Student_IC_school_bus_202076.jpg")

    image_string = urllib2.urlopen(url).read()
    image = tf.image.decode_jpeg(image_string, channels=3)

    # Convert image to float32 before subtracting the
    # mean pixel value
    image_float = tf.to_float(image, name='ToFloat')

    # Subtract the mean pixel value from each pixel
    processed_image = _mean_image_subtraction(image_float,
                                              [_R_MEAN, _G_MEAN, _B_MEAN])

    input_image = tf.expand_dims(processed_image, 0)

    with slim.arg_scope(vgg.vgg_arg_scope()):

        # spatial_squeeze option enables to use network in a fully
        # convolutional manner
        logits, _ = vgg.vgg_16(input_image,
                               num_classes=1000,
                               is_training=False,
                               spatial_squeeze=False)

    # For each pixel we get predictions for each class
    # out of 1000. We need to pick the one with the highest
    # probability. To be more precise, these are not probabilities,
    # because we didn't apply softmax. But if we pick a class
    # with the highest value it will be equivalent to picking
    # the highest value after applying softmax
    pred = tf.argmax(logits, dimension=3)

    init_fn = slim.assign_from_checkpoint_fn(
        os.path.join(checkpoints_dir, 'vgg_16.ckpt'),
        slim.get_model_variables('vgg_16'))

    with tf.Session() as sess:
        init_fn(sess)
        segmentation, np_image, np_logits = sess.run([pred, image, logits])

# Remove the first empty dimension
segmentation = np.squeeze(segmentation)

names = imagenet.create_readable_names_for_imagenet_labels()

# Let's get unique predicted classes (from 0 to 1000) and
# relable the original predictions so that classes are
# numerated starting from zero
unique_classes, relabeled_image = np.unique(segmentation,
                                            return_inverse=True)

segmentation_size = segmentation.shape

relabeled_image = relabeled_image.reshape(segmentation_size)

labels_names = []

for index, current_class_number in enumerate(unique_classes):

    labels_names.append(str(index) + ' ' + names[current_class_number+1])

# Show the downloaded image
plt.figure()
plt.imshow(np_image.astype(np.uint8))
plt.suptitle("Input Image", fontsize=14, fontweight='bold')
plt.axis('off')
plt.show()

discrete_matshow(data=relabeled_image, labels_names=labels_names, title="Segmentation")

接著,我們用雙線性上取樣核對預測結果進行上取樣操作。

upsampled_logits = upsample_tf(factor=32, input_img=np_logits.squeeze())
upsampled_predictions = upsampled_logits.squeeze().argmax(axis=2)

unique_classes, relabeled_image = np.unique(upsampled_predictions,
                                            return_inverse=True)

relabeled_image = relabeled_image.reshape(upsampled_predictions.shape)

labels_names = []

for index, current_class_number in enumerate(unique_classes):

    labels_names.append(str(index) + ' ' + names[current_class_number+1])

# Show the downloaded image
plt.figure()
plt.imshow(np_image.astype(np.uint8))
plt.suptitle("Input Image", fontsize=14, fontweight='bold')
plt.axis('off')
plt.show()

discrete_matshow(data=relabeled_image, labels_names=labels_names, title="Segmentation")

我們得到的結果混雜了很多噪音資訊,但是我們基本上把校車給分割出來了。更確切的說,這並不是影象分割的結果,而是神經網路預測的分類標籤所組成的區域。

總結

在本文中,我們介紹了轉置卷積,特別是雙線性差值的實現。我們對降取樣的預測結果進行上取樣還原,得到了整幅圖片的預測類別。

另外,我們還發現對很多很多類別做影象分割時,會遇到記憶體消耗過大的問題,因此只能使用CPU計算。

130+位講師,16大分論壇,中國科學院院士陳潤生、滴滴出行高階副總裁章文嵩、聯想集團高階副總裁兼CTO芮勇、上交所前總工程師白碩等專家將親臨2016中國大資料技術大會,票價折扣即將結束,預購從速

圖片描述