天天看點

深度網絡圖像分割通俗指南

深度網絡圖像分割通俗指南

給定一張包含家具的照片,你可以讓程式自動将家具與背景分開嗎?

在這篇文章中,我将介紹如何使用目前最先進的深度學習來嘗試解決這個問題。我不是機器學習的專家,是以我希望這個文章對于其他希望使用這個強大新工具的非專家們有一定的幫助作用。

這個問題稱為分割。也就是說,從這張圖開始:

深度網絡圖像分割通俗指南

到這張圖:

深度網絡圖像分割通俗指南

我們可以将這個遮罩應用到源圖像上,獲得沒有背景的椅子。我們将使用一些工具來簡化這個工作:

第一步是安裝資料集。由于它托管在Brine上,是以可以用一個簡單的指令來實作:<code>brine install rohan/chairs-with-masks</code>。

下一步是加載資料集。可以通過Brine的<code>load_dataset</code>函數來執行此操作:<code>chairs = brine.load_dataset('rohan/chairs-with-masks')</code>。該資料集包含了97個樣本,每個樣本是圖像及其遮罩。遮罩是一個隻有兩種顔色的圖像,藍色代表背景,紅色代表前景。

我們可以看到輸入的形狀是<code>(None, 256, 256, 3)</code>,輸出的形狀是<code>(None, 256, 256, 1)</code>。元組的第一個元素是批量的大小,是以我們現在可以忽略它。這告訴我們,網絡期望的輸入是一批256x256的三通道圖像,并将輸出一批256x256個單通道遮罩。我們的遮罩也需要比對這個形狀。

下一步是準備樣本,使之與網絡一起使用。我們将為訓練資料定義一個處理函數,在樣本傳給網絡之前需應用于每個樣本。

這裡做了很多事情,我會一步一步進行解釋。樣本作為一個元組被傳遞進去,是以首先要進行解包。接下來的兩行使用numpy切片來確定圖像隻有3個通道,如果有第四個alpha通道,則忽略。然後,使用cv2(OpenCV的python綁定)将遮罩轉換為灰階圖,這樣,我們現在就有了一個單通道遮罩,這是網絡所期望的。這裡沒有為兩種顔色使用兩個随機的灰階數字,而是強制使用<code>fix_mask</code>函數将掩碼設定為0和255以表示背景和前景。然後,我們将圖像和掩碼的大小調整為256x256,以比對網絡指定的大小。

由于資料集較小,是以我們将使用資料擴充。資料擴充是指在訓練期間在保留原始資訊的基礎上随機修改圖像,人為地生成更多的資料。例如,旋轉5度的椅子仍然是椅子,是以網絡應該能夠正确地識别出來。在代碼中,我們使用petrosgk的Carvana示例中的3個函數來随機改變圖像的色相、飽和度和值,并随機旋轉和翻轉圖像。如果旋轉或翻轉了圖像,則必須對遮罩執行相同的操作,以使遮罩與原始圖像保持一緻。

最後,我們通過将所有像素值除以255來對資料進行歸一化操作,這樣,所有的值都在0和1之間了。如果在此時列印<code>image.shape</code>,結果是256x256x3,這正是網絡所需要的。盡管<code>mask.shape</code>是256x256,但網絡是256x256x1,是以我們使用<code>np.expand_dims</code>來讓遮罩比對這個形狀。最後,傳回新的圖像和遮罩。

在開始訓練網絡之前,還需要用一些樣本進行驗證。驗證集不能用于訓練,隻能用于檢查模型的性能。我們可以使用Brine的<code>create_folds</code>來建立:

我們将樣本集中的20個圖像作為驗證集,剩下的77個圖像作為訓練集。

我們還定義了用于樣本驗證的處理函數,這與用于訓練的處理函數非常相似,除了在驗證時不使用資料擴充之外。

最後,Brine傳回一個可用于訓練和校驗集的圖檔生成器:

這些生成器将傳回樣品的批處理,以及之前定義的處理函數。它們可以被直接傳給Keras的<code>fit_generator</code>方法來訓練模型。

現在,準備開始訓練模型了:

這些回調會改變Keras的訓練過程。 <code>EarlyStopping</code>将在驗證損失的改進停止後停止訓練,<code>ReduceLROnPlateau</code>将降低學習率,<code>ModelCheckpoint</code>将對在驗證集上表現最佳的模型版本進行儲存。

最後,可以開始訓練了:

Keras開始訓練模型了,它将在資料集上運作多次(由于資料增加和混洗,每次運作的結果會略有不同),并輸出訓練和驗證集的損失和DICE分數。在某些時候,運作會停止,可能是因為<code>EarlyStopping</code>回調,或是因為達到了100次疊代。

好了,就這些!我們訓練了一個U-Net網絡,并嘗試在圖像中把椅子分割出來。讓我們來看看它的實際表現如何吧。

要産生一些預測,首先使用<code>model.load_weights('weights/best_weights.hdf5')</code>來加載最佳模型的權重,然後在驗證集中的某些圖像上使用模型的<code>predict</code>方法:

以下是一些樣本的結果:

深度網絡圖像分割通俗指南

在原始圖像上覆寫了網絡預測出來的遮罩。顔色越黃,則說明網絡對該像素的預測具有更高的置信度

雖然這對于生産系統來說還不夠好,但網絡已經學到了有關前景與背景、椅子與非椅子的一些内容。我個人覺得這個網絡在沒有進行預先訓練,并且隻有77張圖檔的情況下表現得已經非常不錯了。如果有更大的資料集,網絡的準确性可能會大大提高,或許就能用于生産系統了。

文章原标題《A Non-Expert’s Guide to Image Segmentation Using Deep Neural Nets》,作者:Rohan Relan,譯者:夏天,審校:主題曲。

本文由北郵@愛可可-愛生活老師推薦,阿裡雲雲栖社群組織翻譯。

繼續閱讀