本文對EAST算法中的重要模塊進行梳理,包括多gpu訓練、算法pipline、輸入數據和輸出數據的格式。
- 數據並行。使用多gpu進行訓練,將batch的訓練數據根據設定的gpu數量進行劃分,每個gpu訓練batch數據中的一部分,得到其結構風險,然後計算當前結構風險的梯度信息。當前batch數據的梯度全部計算完後,進行梯度更新。
for i, gpu_id in enumerate(gpus): #數據並行
with tf.device('/gpu:%d' % gpu_id):
with tf.name_scope('model_%d' % gpu_id) as scope:
iis = input_images_split[i]
isms = input_score_maps_split[i]
igms = input_geo_maps_split[i]
itms = input_training_masks_split[i]
total_loss, model_loss = tower_loss(iis, isms, igms, itms, reuse_variables)
batch_norm_updates_op = tf.group(*tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope))
reuse_variables = True
grads = opt.compute_gradients(total_loss)
tower_grads.append(grads) #彙總梯度
grads = average_gradients(tower_grads) #更新梯度
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
# save moving average
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# batch norm updates
with tf.control_dependencies([variables_averages_op, apply_gradient_op, batch_norm_updates_op]): #訓練tensor
train_op = tf.no_op(name='train_op')
- 模型pipline
- backbone爲PVANet;
- 特徵合併的規則採用了U-net的方法,從頂部特徵按照相應的規則向下進行合併,合併至h4層時,當前feature map的大小爲原圖的1/4。如上所示。特徵提取層與特徵合併層和起來類似FPN結構,不過FPN是在多層進行預測。
- 輸出層包含三大部分:score_map(類別置信度)、geo_map(每個像素點到其旋轉矩形框的四條邊的距離)、angle_map(旋轉矩形的角度)。 ps QUAD geometry代碼中未實現。
- 輸入數據
- 數據集icdar2015關於自然場景下標註信息的數據格式,x1, y1, x2, y2, x3, y3, x4, y4, transcription,其中transcription=###時表示不關心的區域。
- 輸入四邊形與其外接旋轉矩形的關係。如下圖,(a) 中黃色虛線表示標註區域,綠色實線表示壓縮後的正樣本區域;(b)將壓縮後的區域作爲score_map;(c)粉色框表示標註區域的最小外接矩形,黑色點表示壓縮多邊形內部的點,四種顏色的線分別表示當前點到外接矩形邊的距離,淺藍色的線表示旋轉矩形的角度;(d)共4個channel,分別表示圖上每個位置到外接矩形的某一條邊的距離值;(e)表示旋轉的角度。標註數據到網絡輸入數據的轉換過程詳見icdar.py文件,主函數generate_rbox。功能函數包括:壓縮區域的函數爲shrink_poly, 獲得標註區域的外接矩形框的函數爲rectangle_from_parallelogram(先獲得標註區域所在的平行四邊形,然後由平行四邊形獲得矩形區域),矩形框旋轉角度的函數爲sort_rectangle。
- 外接矩形的構建過程。由標註的四邊形得到其所在的平行四邊形,然後利用平行四邊形得到矩形。平行四邊形的構建過程如下所示,選擇edge和forward_edge兩條邊作爲所在平行四邊形中的兩條相鄰邊,然後計算平行四邊形的頂點位置。矩形框選擇與水平方向夾角小於90度的。可對比generate_rbox代碼中fitted_parallelograms變量的生成過程理解平行四邊形的生成過程。 ps:fit_line函數擬合的直線爲ax+by+c=0, 其中p1 is [x1, x2], and p2 is [y1, y2],所以該函數中垂直於x軸的條件是p1[0]==p1[1];凸四邊形的外接平行四邊形不止一種情況,作者選擇的是將標註框的top和right邊作爲平行四邊形的兩條相鄰邊的情況;平行四邊形的外接矩形框有兩個,選擇與水平線夾角小於90度的矩形框;直線夾角的計算公式見這裏。
- 損失函數由兩部分組成:分類損失和迴歸損失,其中分類損失論文中給出的是帶權重的二分類交叉熵,代碼中給出的
dice_coefficient。迴歸部分採用iou損失函數。
- 輸出數據。損失函數決定了輸出信息爲每個位置到其所在旋轉矩形每條邊的距離。所以由網絡輸出到預測值之間需要j經過兩個處理:構建矩形區域和矩形區域的後處理(Locality-Aware NMS)。
- 矩形局域的構建。eval.py文件中detect函數爲預測時的入口,矩形局域的構建函數在icdar.py文件的restore_rectangle_rbox函數中實現。由於矩形與橫軸的夾角不同時,旋轉矩陣的構建也不相同,所以先根據矩陣與橫軸夾角分爲兩種情況,即angle>0和angel<0;接着,構建每種角度下當前輸出像素點所在矩形的座標信息(0, -h, w, -h, w, 0, 0, 0, left, -bottom),其中前8個值分別表示當前矩陣順時針方向的4個點(和標註順序相同),(left, -bottom)表示當前像素點在矩形中的座標。然後計算矩陣旋轉後的值,並通過如下代碼得到矩形的四個頂點座標。
p3_in_origin = origin_0 - p_rotate[:, 4, :] #獲得矩形中每個像素點在圖像中的起始座標
new_p0 = p_rotate[:, 0, :] + p3_in_origin #根據座標的偏移量(p_rotate)得到矩形的四個頂點座標
new_p1 = p_rotate[:, 1, :] + p3_in_origin
new_p2 = p_rotate[:, 2, :] + p3_in_origin
new_p3 = p_rotate[:, 3, :] + p3_in_origin
- Locality-Aware NMS。生成的矩形區域數量巨大,且水平方向相鄰矩形區域高度相關。所以EAST先從行進行矩形框的合併,然後再進行標準的NMS。
參考文獻: