TensorFlow之deeplab語義分割API介面除錯
在之前的文章中，對tensorflow目標檢測API進行了詳細的測試，成功應用其模型做簡單的檢測任務。本文對另一模組DeepLab的API進行測試，實現語義分割。
經過了好幾天的吐血折騰，終於將該模組調通，其中的bug真是數不勝數……
1 檔案結構
首先在research/deeplab/datasets下新建一個資料夾，這裡我建的是loulan，用來做漏纜的語義分割。然後在資料夾下新建dataset、init_models、tfrecord和train子資料夾
loulan . ├── dataset │ ├── changepixel.py │ ├── ImageSets │ ├── JPEGImages │ ├── SegmentationClass │ ├── SegmentationClassRaw │ ├── SegmentationClass(RGBA) │ └── trans2raw.py ├── export │ ├── frozen_inference_graph-2675.pb │ ├── frozen_inference_graph-2675.tar.gz ├── export_model.sh ├── init_models │ ├── deeplabv3_pascal_train_aug │ └── deeplabv3_pascal_train_aug_2018_01_04.tar.gz ├── tfrecord │ ├── train-00000-of-00004.tfrecord │ ├── train-00001-of-00004.tfrecord │ ├── train-00002-of-00004.tfrecord │ ├── train-00003-of-00004.tfrecord │ ├── trainval-00000-of-00004.tfrecord │ ├── trainval-00001-of-00004.tfrecord │ ├── trainval-00002-of-00004.tfrecord │ ├── trainval-00003-of-00004.tfrecord │ ├── val-00000-of-00004.tfrecord │ ├── val-00001-of-00004.tfrecord │ ├── val-00002-of-00004.tfrecord │ └── val-00003-of-00004.tfrecord └── train ├── checkpoint ├── graph.pbtxt ├── model.ckpt-2675.data-00000-of-00001 ├── model.ckpt-2675.index ├── model.ckpt-2675.meta
2 準備資料
這一步操作都在loulan/dataset資料夾中。
2.1 原始資料
將原始JPEG圖片打包放在JPEGImages資料夾下，標註檔案放在SegmentationClass資料夾下。這塊要注意原始圖片和標註圖片都是RGB格式的
神坑一：我找人標註的圖片格式是RGBA，而且是用A元素來區分的類別，RGB分量是無規律的，因此我要先把圖片的畫素更改過來，用RGB格式來區分圖片中的種類。具體轉化的程式碼，見我另一篇部落格《Python之修改圖片畫素值》。
2.2 轉化灰度圖
上一步將標註檔案也轉化為了RGB格式的，接下來要將標註檔案轉化為單通道的灰度圖，用不同的畫素值來表示不同的類別。程式碼如下。（由於我目標只有1類，因此只需要用2種畫素即可）
import tensorflow as tf from PIL import Image from tqdm import tqdm import numpy as np import os, shutil # palette (color map) describes the (R, G, B): Label pair palette = {(0,0,0) : 0 , #0表示背景 (255,255,255) : 1 #1表示類別 } def convert_from_color_segmentation(arr_3d): arr_2d = np.zeros((arr_3d.shape[0], arr_3d.shape[1]), dtype=np.uint8) for c, i in palette.items(): m = np.all(arr_3d == np.array(c).reshape(1, 1, 3), axis=2) arr_2d[m] = i return arr_2d label_dir = './SegmentationClass/' new_label_dir = './SegmentationClassRaw/' if not os.path.isdir(new_label_dir): print("creating folder: ",new_label_dir) os.mkdir(new_label_dir) else: print("Folder alread exists. Delete the folder and re-run the code!!!") label_files = os.listdir(label_dir) for l_f in tqdm(label_files): arr = np.array(Image.open(label_dir + l_f)) arr = arr[:,:,0:3] arr_2d = convert_from_color_segmentation(arr) Image.fromarray(arr_2d).save(new_label_dir + l_f)
2.3 資料分類
在loulan/dataset/ImageSets資料夾下新建三個文字檔案，分別是train.txt、trainval.txt、val.txt，然後將標註的資料的檔名按照訓練集：驗證集 = 4:1的比例分別保存於train.txt和val.txt，然後把所有的名字儲存在trainval.txt中。可參考我以下方法儲存
# -*- coding:utf8 -*-
import os
path = 'JPEGImages/'
filelist = os.listdir(path)
i = 0
#訓練集儲存方法
with open("train.txt","a") as f:
for item in filelist:
i += 1
#print('item name is ',item)
if i%5 != 0:
name = item.split('.',2)[0] + '.' + item.split('.',2)[1]
print('name is ',name)
f.write(name + '\n')其它兩個檔案的生成方法類似，稍加改動就行
2.4 生成tfrecord
在deeplab/datasets資料夾下，有生成tfrecord的程式，因為我們採用的是voc資料結構，因此需執行build_voc2012_data.py，具體引數如下
python ./build_voc2012_data.py --image_folder="./loulan/dataset/JPEGImages" --semantic_segmentation_folder="./loulan/dataset/SegmentationClassRaw" --list_folder="./loulan/dataset/ImageSets" --image_format='jpg' --output_dir="./loulan/tfrecord"3 網路訓練
3.1 預訓練模型下載
官方提供了不少預訓練模型，這裡以deeplabv3_pascal_train_aug_2018_01_04為例
在loulan/init_models資料夾中下載解壓模型，得到deeplabv3_pascal_train_aug資料夾
wget http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz
tar zxf deeplabv3_pascal_train_aug_2018_01_04.tar.gz3.2 資料集描述
在deeplab/datasets資料夾下，有一個segmentation_dataset.py檔案，用來描述我們資料集，開啟之後新增如下程式碼
首先在_DATASETS_INFORMATION中新增如下
_DATASETS_INFORMATION = {
'cityscapes': _CITYSCAPES_INFORMATION,
'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
'ade20k': _ADE20K_INFORMATION,
'loulan': _LOULAN_INFORMATION
}然後與_ADE20K_INFORMATION並列關係，新增如下程式碼
_PQR_SEG_INFORMATION = DatasetDescriptor(
splits_to_sizes={
'train': 824, # 訓練集的圖片數量
'trainval': 1030, #所有圖片數量
'val': 206, #驗證集的圖片數量
},
num_classes=4, # 資料類別，包含背景
ignore_label=255, # 忽略的類別
)其中，將自己圖片的數量依次對應修改，然後在num_classes中，我寫的是4，參考了《tensorflow下deeplab模型訓練資料集過程（續）》，他說：
“ 資料集的類別設定為4，是因為還有兩個預設的類別，分別是ignored_label（255）和-1. 從名字可以看出，ignored_label表示忽略該類別，即不考慮該類別，主要是在製作資料集的時候需要用到，而-1是為了保證過程中出現一些未知的類別。兩者都是為了保證訓練過程不報錯。
於是可能您會說，維護什麼不將訓練中的非道路類別歸為255或者-1.這個我的確做了，我一開始的確是將非道路視為255的，這樣實際的目標類別就只有一類，但是訓練出來的效果並不好。於是再將類別多加一類後，效果好了很多。”
具體還沒有測試過，這裡先這麼修改
還有一篇文章裡提到要修改deeplab/train.py檔案：
# Set to False if one does not want to re-use the trained classifier weights.
flags.DEFINE_boolean('initialize_last_layer', True,
'Initialize the last layer.')
把True改為False，重新訓練最後一層
3.3 開始訓練
在deeplab下，有一個train.py檔案，我們只需要把引數設定好即可，具體如下
python train.py --logtostderr --train_split="train" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18 --output_stride=16 --decoder_output_stride=4 --train_crop_size=512 --train_crop_size=512 --train_batch_size=4 --training_number_of_steps=30000 --fine_tune_batch_norm=false --tf_initial_checkpoint="./datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/model.ckpt" --train_logdir="./datasets/loulan/train/" --dataset_dir="./datasets/loulan/tfrecord" --num_clones=4訓練的時候有幾個引數要注意：--num_clones=4 表示在4塊GPU上進行訓練，要根據自己顯示卡數量進行配置
這裡要注意正確的設定train_batch_size的大小，我剛開始設定了16，執行時會報視訊記憶體不足的錯誤。然後修改為4，同時要設定fine_tune_batch_norm為false，如果顯示卡足夠強，batch_size能設定的大於12時，要將fine_tune_batch_norm設定為True。batch_size設定的大的時候，訓練速度會變慢
訓練的時候，可以用tensorboard --logdir='train' 來觀察loss的變化，如果loss趨於穩定，就可以停止訓練了
3.4 匯出模型
訓練後在loulan/train資料夾下生成一些結果檔案，如
- graph.pbtxt
- model.ckpt-1000.data-00000-of-00001
- model.ckpt-1000.info
- model.ckpt-1000.meta
其中meta檔案儲存了graph和metadata,ckpt檔案儲存了網路的weights，進行預測時有模型的權重就夠了，可以使用官方提供的指令碼來生成模型檔案，指令碼檔案是deeplab資料夾下export_model.py，我們在deeplab/datasets/loulan資料夾下新建export_model.sh指令碼，輸入以下內容：
python ../../export_model.py --logtostderr --checkpoint_path="train/model.ckpt-$1" --export_path="export/frozen_inference_graph-$1.pb" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18 --output_stride=16 --decoder_output_stride=4 --num_classes=21 --crop_size=513 --crop_size=513 --inference_scales=1.0然後執行 sh export_model.sh 2675即可，其中2675是train資料夾下最新的生成檔案，即有以下三個檔案
- model.ckpt-2675.data-00000-of-000001
- model.ckpt-2675.index
- model.ckpt-2675.meta
執行時要根據自己的檔案來修改數字，執行後在export資料夾下生成了frozen_inference_graph-2675.pb檔案，該模型可以在後邊用來做檢測。
4 模型測試
4.1 ipynb方法
在deeplab目錄下，有一個deeplab_demo.ipynb檔案，我們只需要修改裡邊的路徑就可以
4.1.1 修改類別
找到以下程式碼修改為自己的類別
LABEL_NAMES = np.asarray([
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'
])比如我只有兩個類別，則改為
LABEL_NAMES = np.asarray([
'background', 'lane'
])4.1.2 載入模型
首先要將上邊生成的pb檔案壓縮一下，我直接匯入測試的時候報錯，檢視原始碼發現是有一個解壓的過程，因此首先把檔案壓縮後再匯入，壓縮指令如下
tar zcvf frozen_inference_graph-2675.tar.gz frozen_inference_graph-2675.pb然後把下載模型的部分註釋掉，新增載入模型
#@title Select and download models {display-mode: "form"}
'''
MODEL_NAME = 'mobilenetv2_coco_voctrainaug' # @param
中間省略
MODEL = DeepLabModel(download_path)
'''
model_path ='./datasets/loulan/export/frozen_inference_graph-8070.tar.gz'
MODEL = DeepLabModel(model_path)
print('model loaded successfully!')4.1.3 新增測試圖片
將最後一段下載url圖片的程式註釋掉，新增自己圖片路徑執行模型即可，如下
image_path = './datasets/loulan/dataset/JPEGImages/19678.15(20181011135748557_0).jpg'
original_im = Image.open(image_path)
resized_im, seg_map = MODEL.run(original_im)
vis_segmentation(resized_im, seg_map)然後點選cell-run all即可執行出結果，效果如下
參考：https://lijiancheng0614.github.io/2018/03/13/2018_03_13_TensorFlow-DeepLab/