如何用TensorFlow和TF-Slim實現影象分類與分割

2017年04月08日 11:29:56

閱讀數：1282

本文將介紹如何用近日釋出的TF-Slim工具包和預訓練的模型來完成影象分類和影象分割。

引言

筆者將和大家分享一個結合了TensorFlow和最近釋出的slim庫的小應用，來實現影象分類、影象標註以及影象分割的任務，圍繞著slim展開，包括其理論知識和應用場景。

之前自己嘗試過許多其它的庫，比如Caffe、Matconvnet、Theano和Torch等。它們各有優劣，而我想要一個可靠靈活的、自帶預訓練模型的Python庫。最近，新推出了一款名叫slim的庫，slim自帶了許多預訓練的模型，比如ResNet、VGG、Inception-ResNet-v2（

下文中會用到Tensorflow和卷積神經網路的知識。Tensorflow的網站上有兩者的完美教程，不瞭解的讀者可以前去閱讀。

我是用jupyter notebook完成的寫作。因此，每一段程式碼之後都列印了執行結果。讀者們也可以下載完整的notebook。本文中有一部分內容借鑑了此文

安裝

在執行程式碼之前，首先需要安裝Tensorflow。我用的是0.11版本。你可以從github的tensorflow/models程式碼庫克隆程式碼。

git clone https://github.com/tensorflow/models

我還會用到scikit-image和numpy等依賴，把它們都先裝上。在這裡我推薦先下載並安裝Anaconda，然後通過conda install命令安裝其它的python庫。

首先，我們指定tensorflow使用第一塊GPU。否則tensorflow預設會佔用所有可用的記憶體資源。其次，新增克隆下來的程式碼庫路徑，這樣python執行的時候就能找到需要的程式碼。

1. import sys

2. import os

3. os.environ["CUDA_VISIBLE_DEVICES"] = '0'

4. sys.path.append("/home/dpakhom1/workspace/models/slim")

接著，下載VGG-16模型，我們將用它來對影象做分類和分割。也可以選用其它佔用記憶體少的網路模型（比如，AlexNet）。關於模型的詳細內容，請參考此頁面。

1. from datasets import dataset_utils

2. import tensorflow as tf

3. url = "http://download.tensorflow.org/models/vgg_16_2016_08_28.tar.gz"

4. # 指定儲存路徑

5. checkpoints_dir = '/home/dpakhom1/checkpoints'

6. if not tf.gfile.Exists(checkpoints_dir):

7. tf.gfile.MakeDirs(checkpoints_dir)

8. dataset_utils.download_and_uncompress_tarball(url, checkpoints_dir)

9. >> Downloading vgg_16_2016_08_28.tar.gz 100.0%

10. Successfully downloaded vgg_16_2016_08_28.tar.gz 513324920 bytes.

影象分類

我們剛剛下載的模型可以將影象分成1000類。類別的覆蓋度非常廣。在本文中，我們就用這個預訓練的模型來給圖片分類、標註和分割，對映到這1000個類別。

下面是一個影象分類的例子。影象首先要做預處理，經過縮放和裁剪，輸入的影象尺寸與訓練集的圖片尺寸相同。

1. %matplotlib inline

2. from matplotlib import pyplot as plt

3. import numpy as np

4. import os

5. import tensorflow as tf

6. import urllib2

7. from datasets import imagenet

8. from nets import vgg

9. from preprocessing import vgg_preprocessing

10. checkpoints_dir = '/home/dpakhom1/checkpoints'

11. slim = tf.contrib.slim

12. # 網路模型的輸入影象有預設的尺寸

13. # 因此，我們需要先調整輸入圖片的尺寸

14. image_size = vgg.vgg_16.default_image_size

15. with tf.Graph().as_default():

16. url = ("https://upload.wikimedia.org/wikipedia/commons/d/d9/"

17. "First_Student_IC_school_bus_202076.jpg")

18. # 連線網址，下載圖片

19. image_string = urllib2.urlopen(url).read()

20. # 將圖片解碼成jpeg格式

21. image = tf.image.decode_jpeg(image_string, channels=3)

22. # 對圖片做縮放操作，保持長寬比例不變，裁剪得到圖片中央的區域

23. # 裁剪後的圖片大小等於網路模型的預設尺寸

24. processed_image = vgg_preprocessing.preprocess_image(image,

25. image_size,

26. image_size,

27. is_training=False)

28. # 可以批量匯入影象

29. # 第一個維度指定每批圖片的張數

30. # 我們每次只匯入一張圖片

31. processed_images = tf.expand_dims(processed_image, 0)

32. # 建立模型，使用預設的arg scope引數

33. # arg_scope是slim library的一個常用引數

34. # 可以設定它指定網路層的引數，比如stride, padding 等等。

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

36. logits, _ = vgg.vgg_16(processed_images,

37. num_classes=1000,

38. is_training=False)

39. # 我們在輸出層使用softmax函式，使輸出項是概率值

40. probabilities = tf.nn.softmax(logits)

41. # 建立一個函式，從checkpoint讀入網路權值

42. init_fn = slim.assign_from_checkpoint_fn(

43. os.path.join(checkpoints_dir, 'vgg_16.ckpt'),

44. slim.get_model_variables('vgg_16'))

45. with tf.Session() as sess:

46. # 載入權值

47. init_fn(sess)

48. # 圖片經過縮放和裁剪，最終以numpy矩陣的格式傳入網路模型

49. np_image, network_input, probabilities = sess.run([image,

50. processed_image,

51. probabilities])

52. probabilities = probabilities[0, 0:]

53. sorted_inds = [i[0] for i in sorted(enumerate(-probabilities),

54. key=lambda x:x[1])]

55. # 顯示下載的圖片

56. plt.figure()

57. plt.imshow(np_image.astype(np.uint8))

58. plt.suptitle("Downloaded image", fontsize=14, fontweight='bold')

59. plt.axis('off')

60. plt.show()

61. # 顯示最終傳入網路模型的圖片

62. # 影象的畫素值做了[-1, 1]的歸一化

63. # to show the image.

64. plt.imshow( network_input / (network_input.max() - network_input.min()) )

65. plt.suptitle("Resized, Cropped and Mean-Centered input to network",

66. fontsize=14, fontweight='bold')

67. plt.axis('off')

68. plt.show()

69. names = imagenet.create_readable_names_for_imagenet_labels()

70. for i in range(5):

71. index = sorted_inds[i]

72. # 列印top5的預測類別和相應的概率值。

73. print('Probability %0.2f => [%s]' % (probabilities[index], names[index+1]))

74. res = slim.get_model_variables()

1. Probability 1.00 => [school bus]

2. Probability 0.00 => [minibus]

3. Probability 0.00 => [passenger car, coach, carriage]

4. Probability 0.00 => [trolleybus, trolley coach, trackless trolley]

5. Probability 0.00 => [cab, hack, taxi, taxicab]

圖片標註和分割

從上面的例子中可以看到，網路模型只處理了原始影象中的一部分割槽域。這種方式只適用於單一預測結果的場景。

某些場景下，我們希望從圖片中獲得更多的資訊。舉個例子，我們想知道圖片中出現的所有物體。網路模型就告訴我們圖片中有一輛校車，還有幾輛小汽車和幾幢建築物。這些資訊可以協助我們搭建一個圖片搜尋引擎。以上就是一個圖片標註的簡單應用。

但是，如果我們也想得到物體的空間位置該怎麼辦。網路能告訴我們它在圖片的中央看到一輛校車，在右上角看到幾幢建築物？這樣，我們就可以建立一個更具體的搜尋查詢詞：“我想要找到中間有一輛校車，左上角有幾隻花盆的所有符合要求的圖片”。

某些情況下，我們需要對影象的每個畫素進行分類，也被稱作是影象的分割。想象一下，假如有一個巨大的圖片資料集，需要給人臉打上馬賽克，這樣我們就不必得到所有人的許可之後才能釋出這些照片。例如，谷歌街景都對行人的臉做了模糊化處理。當然，我們只需要對圖片中的人臉進行模糊處理，而不是所有的內容。圖片分割可以幫助我們實現類似的需求。我們可以分割得到屬於人臉的那部分畫素，並只對它們進行模糊處理。

下面將介紹一個簡單的圖片分割例子。我們可以使用現有的卷積神經網路，通過完全卷積的方式進行分割。若想要輸出的分割結果與輸入影象尺寸保持一致，可以增加一個去卷積層。

1. from preprocessing import vgg_preprocessing

2. # 載入畫素均值及相關函式

3. from preprocessing.vgg_preprocessing import (_mean_image_subtraction,

4. _R_MEAN, _G_MEAN, _B_MEAN)

5. # 展現分割結果的函式，以不同的顏色區分各個類別

6. def discrete_matshow(data, labels_names=[], title=""):

7. #獲取離散化的色彩表

8. cmap = plt.get_cmap('Paired', np.max(data)-np.min(data)+1)

9. mat = plt.matshow(data,

10. cmap=cmap,

11. vmin = np.min(data)-.5,

12. vmax = np.max(data)+.5)

13. #在色彩表的整數刻度做記號

14. cax = plt.colorbar(mat,

15. ticks=np.arange(np.min(data),np.max(data)+1))

16. # 新增類別的名稱

17. if labels_names:

18. cax.ax.set_yticklabels(labels_names)

19. if title:

20. plt.suptitle(title, fontsize=14, fontweight='bold')

21. with tf.Graph().as_default():

22. url = ("https://upload.wikimedia.org/wikipedia/commons/d/d9/"

23. "First_Student_IC_school_bus_202076.jpg")

24. image_string = urllib2.urlopen(url).read()

25. image = tf.image.decode_jpeg(image_string, channels=3)

26. # 減去均值之前，將畫素值轉為32位浮點

27. image_float = tf.to_float(image, name='ToFloat')

28. # 每個畫素減去畫素的均值

29. processed_image = _mean_image_subtraction(image_float,

30. [_R_MEAN, _G_MEAN, _B_MEAN])

31. input_image = tf.expand_dims(processed_image, 0)

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

33. # spatial_squeeze選項指定是否啟用全卷積模式

34. logits, _ = vgg.vgg_16(input_image,

35. num_classes=1000,

36. is_training=False,

37. spatial_squeeze=False)

38. # 得到每個畫素點在所有1000個類別下的概率值，挑選出每個畫素概率最大的類別

39. # 嚴格說來，這並不是概率值，因為我們沒有呼叫softmax函式

40. # 但效果等同於softmax輸出值最大的類別

41. pred = tf.argmax(logits, dimension=3)

42. init_fn = slim.assign_from_checkpoint_fn(

43. os.path.join(checkpoints_dir, 'vgg_16.ckpt'),

44. slim.get_model_variables('vgg_16'))

45. with tf.Session() as sess:

46. init_fn(sess)

47. segmentation, np_image = sess.run([pred, image])

48. # 去除空的維度

49. segmentation = np.squeeze(segmentation)

50. unique_classes, relabeled_image = np.unique(segmentation,

51. return_inverse=True)

52. segmentation_size = segmentation.shape

53. relabeled_image = relabeled_image.reshape(segmentation_size)

54. labels_names = []

55. for index, current_class_number in enumerate(unique_classes):

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

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

我們得到的結果顯示網路模型確實可以從圖片中找到校車，以及左上角顯示不太清晰的交通標誌。而且，模型可以找到左上角建築物的窗戶，甚至猜測說這是一個圖書館（我們無法判斷是否屬實）。它做出了一些不那麼正確的預測。這些通常是由於網路在預測的時候只能看到當前畫素周圍的一部分影象。網路模型表現出來的這種特性被稱為感受視野。在本文中，我們使用的網路模型的感受視野是404畫素。所以，當網路只能觀察到校車的一部分圖片時，與計程車和皮卡車混淆了。

正如我們在上面所看到的，我們得到了圖片的一個簡單分割結果。它不算很精確，因為最初訓練網路是用來進實現分類任務，而不是影象分割。如果想得到更好的結果，我們還是需要重新訓練一個模型。不管怎麼說，我們得到的結果是可以用作影象標註的。

使用卷積神經網路進行影象分割，可以被看作是對輸入影象的不同部分進行分類。我們將網路聚焦於某個畫素點，進行預測判斷，並輸出該畫素的類別標籤。這樣，我們給分類和分割的結果增加了空間資訊。

小結

本文介紹了用slim庫實現影象的分類和分割，並且簡要闡述了技術原理。

自帶預訓練模型的slim庫是一款強大而靈活的工具，可以配合tensorflow使用。由於最近剛剛釋出，文件不是很完善，有時候甚至要閱讀程式碼來幫助理解。Google正在加快進度完善後續的工作