ui2code中的深度學習+傳統演算法應用
作者：閒魚技術-雲聽
背景
在之前的 ofollow,noindex" target="_blank">文章 中，我們已經提到過團隊在UI自動化這方面的嘗試，我們的目標是實現基於 單一圖片到程式碼 的轉換，在這個過程不可避免會遇到一個問題，就是為了從單一圖片中提取出足夠的有意義的結構資訊，我們必須要擁有從圖片中切割出想要區塊（文字、按鈕、商品圖片等）的能力，而傳統切割演算法遇到複雜背景圖片往往就捉襟見肘了（見下圖），這個時候，我們就需要有能力把複雜前後景的圖片劃分為各個層級圖層，再交給切割演算法去處理，拿到我們期望的結構資訊。
經過傳統切割演算法處理，會無法獲取圖片結構資訊，最終只會當成一張圖片處理。
在業界，圖片前後景分離一直是個很麻煩的命題，業界目前比較普遍採用的解決方案是計算機視覺演算法提取，或是引入人工智慧來解決，但直到現在，都沒有百分百完美的解決方案。那是否能引入AI來解決這個問題呢，我們來看一下，目前使用AI並拿到比較不錯結果的解法是fcn+crf，基本上能夠把目標物體的前景輪廓框出來，但缺點也很明顯：
- 準確率只有80%左右
- 邊緣切割無法達到畫素級別
- 打標成本非常大
- 難以訓練
- AI是個黑盒，結果不可控
在考慮到使用AI伴隨的問題之外，咱們也一起來思考下，難道AI真的是解決前後景分離的最佳解法嗎？
其實不是的，我們知道，一個頁面，或者說設計稿，一個有意義的前景，是具有比較明顯特徵的，比如說：
- 規則的形狀：線段、矩形、圓形、圓角、是否對稱等
- 形狀上是否有文字，或者說是類似於文字的資訊
- 是否閉合
讓我們一起來驗證下這個思路的可行性。
實踐結果
在嘗試了非常的多計算機視覺演算法之後，你會發現，沒有一種演算法是能夠解決掉這個問題的，基本上是可能一種演算法，在某種場景下是有效的，到了另外一個場景，就又失效了，而且就算是有效的場景，不同顏色複雜度下，所需要的最佳演算法引數又是不相同的。如果case by case來解決的話，可以預期未來的工程會變得越來越冗雜且不好維護。
那是不是可以這樣呢，找到儘可能多的前景區域，加一層過濾器過濾掉前景可能性低的，再加一層層級分配器，對搜尋到的全部前景進行前後層級劃分，最後對影象進行修復，填補空白後景。
咱們先來看看效果，以下查詢前景的過程：
為了避免有的前景被忽略（圖片大部分是有多層的，前景裡面還會巢狀前景），所以一個前景被檢測到之後不會去隱藏它，導致會出現一個前景被多次檢測到的情況，不過這塊加一層層級分配演算法就能解決了，最終得到出來的分離結果如下：
邏輯概要
文書處理
OCR獲取文字粗略位置
來看看例子，以下左圖是閒魚首頁，右圖是基於OCR給出的文字位置資訊對文字區域進行標記（圖中白色部分），可以看到，大致上位置是準確的 但比較粗糙 無法精確到每個文字本身 而且同一行的不同文字片段 OCR會當成一行去處理。
同時，也會有部分非文字的部分 也被當成文字，比如圖中的banner文案：
切割、CNN鑑別器
對以上結果標註的位置進行切割，切割出儘可能小的單個文字區域，交給CNN判斷，該文字是否是可編輯的文字，還是屬於圖片文案，後者將當作圖片進行處理，以下是CNN程式碼：
"""
ui基礎元素識別
"""
# TODO 載入模型
with ui_sess.as_default():
with g2.as_default():
tf.global_variables_initializer().run()
# Loads label file, strips off carriage return
ui_label_lines = [line.rstrip() for line in tf.gfile.GFile("AI_models/CNN/ui-elements-NN/tf_files/retrained_labels.txt")]
# Unpersists graph from file
with tf.gfile.FastGFile("AI_models/CNN/ui-elements-NN/tf_files/retrained_graph.pb", 'rb') as f:
ui_graph_def = tf.GraphDef()
ui_graph_def.ParseFromString(f.read())
tf.import_graph_def(ui_graph_def, name='')
# Feed the image_data as input to the graph and get first prediction
ui_softmax_tensor = ui_sess.graph.get_tensor_by_name('final_result:0')
# TODO 呼叫模型
with ui_sess.as_default():
with ui_sess.graph.as_default():
# UI原子級元素識別
def ui_classify(image_path):
# Read the image_data
image_data = tf.gfile.FastGFile(image_path, 'rb').read()
predictions = ui_sess.run(ui_softmax_tensor, {'DecodeJpeg/contents:0': image_data})
# Sort to show labels of first prediction in order of confidence
top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
for node_id in top_k:
human_string = ui_label_lines[node_id]
score = predictions[0][node_id]
print('%s (score = %s)' % (human_string, score))
return human_string, score
文字抽離
如果是純色背景，文字區域很好抽離，但如果是複雜背景就比較麻煩了。舉個例子：
Line"/>
基於以上，我們能拿到準確的文字資訊，我們逐一對各個文字資訊做處理，文字的特徵還是比較明顯的，比如說含有多個角點，在嘗試了多種演算法：Harris角點檢測、Canny邊緣檢測、SWT演算法，KNN演算法（把區域色塊分成兩部分）之後，發現KNN的效果是最好的。程式碼如下：
Z = gray_region.reshape((-1,1)) Z = np.float32(Z) criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) ret,label,center=cv2.kmeans(Z,K,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS) center = np.uint8(center) res = center[label.flatten()] res2 = res.reshape((gray_region.shape))
抽離後結果如下：
查詢前景
強化圖片邊緣，弱化非邊緣區域
使用卷積核對原圖進行卷積，該卷積核可以強化邊緣，影象平滑區域會被隱藏。
conv_kernel = [ [-1, -1, -1], [-1,8, -1], [-1, -1, -1] ]
卷積後，位與操作隱藏文字區域，結果如下：
降噪
對卷積後的圖，加一層降噪處理，首先把影象轉為灰度圖，接著二值化，小於10畫素值的噪點將被隱藏，最後使用cv2.connectedComponentsWithStats()演算法消除小的噪點連通區域。
基於文字位置，開始查詢輪廓
我們基於前面拿到的文字資訊，選中文字左上角座標，以這個點為種子點執行漫水填充演算法，之後我們會得到一個區域，我們用cv2.findContours()來獲取這個區域的外部輪廓，對輪廓進行鑑別，是否符合有效前景的特徵，之後對區域取反，重新執行cv2.findContours()獲取輪廓，並鑑別。
判斷內外部輪廓
如果文字在輪廓內部，那拿到的區域將不會包含該區域的border邊框，如果文字在輪廓外部，就能拿到包含邊框的一整個有效區域（邊框應該隸屬於前景），所以咱們要判斷文字和輪廓的位置關係（cv2.pointPolygonTest），如果在內部，會使輪廓往外擴散，知道拿到該輪廓的邊框資訊為止。
前景鑑別器
基於前面的步驟，我們會拿到非常多非常多的輪廓，其實絕大部分是無效輪廓以及重複檢測到的輪廓，咱們需要加一層鑑別器來對這些輪廓進行過濾，來判斷它是否是有效前景。
定義有效shape
我們會預先定義我們認為有意義的形狀shape，比如說矩形、正方形、圓形，只要檢測到的輪廓與這三個的相似度達到了設定的閥值要求，並且輪廓中還包含了文字資訊，我們就認為這是一個有意義的前景，見程式碼：
# TODO circle circle = cv2.imread(os.getcwd()+'/fgbgIsolation/utils/shapes/circle.png', 0) _, contours, _ = cv2.findContours(circle, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) self.circle = contours[0] # TODO square square = cv2.imread(os.getcwd()+'/fgbgIsolation/utils/shapes/square.png', 0) _, contours, _ = cv2.findContours(square, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) self.square = contours[0] # TODO rect rect = cv2.imread(os.getcwd()+'/fgbgIsolation/utils/shapes/rect.png', 0) _, contours, _ = cv2.findContours(rect, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) self.rect = contours[0]
匹配shape相似度
多次嘗試之後 發現score設定為3的效果是最好的。程式碼如下：
# TODO 檢測圖形相似度 def detect(self, cnt): shape = "unidentified" types = [self.square, self.rect, self.circle] names = ['square', 'rect', 'circle'] for i in range(len(types)): type = types[i] score = cv2.matchShapes(type, cnt, 1, 0.0)# score越小越相似 # TODO 一般小於3是有意義的 if score<3: shape = names[i] break return shape, score
單一匹配shape相似度的魯棒性還是不夠健壯，所以還引入了其他過濾邏輯，這裡不展開。
影象修復
可以預見的，我們傳入的圖片只有一張，但我們劃分圖層之後，底層的圖層肯定會出現“空白”區域，我們需要對這些區域進行修復。
計算重疊區域
需要修復的區域只在於重疊（重疊可以是多層的）的部分，其他部分我們不應該去修復。計算重疊區域的解決方案沿用了mask遮罩的思路，我們只需要計算當前層有效區域和當前層之上層有效區域的交集即可，使用cv2.bitwise_and
# mask是當前層的mask layers_merge是集合了所有前景的集合i代表當前層的層級數 # inpaint_mask 是要修復的區域遮罩 # TODO 尋找重疊關係 UPPER_level_mask = np.zeros(mask.shape, np.uint8)# 頂層的前景 UPPER_level_mask = np.where(layers_merge>i, 255, 0) UPPER_level_mask = UPPER_level_mask.astype(np.uint8) _, contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # 查詢當前層的每個前景外輪廓 overlaps_mask = np.zeros(mask.shape, np.uint8)# 當前層的所有前景的重疊區域 for cnt in contours: cnt_mask = np.zeros(mask.shape, np.uint8) cv2.drawContours(cnt_mask, [cnt], 0, (255, 255, 255), cv2.FILLED, cv2.LINE_AA) overlap_mask = cv2.bitwise_and(inpaint_mask, cnt_mask, mask=UPPER_level_mask) overlaps_mask = cv2.bitwise_or(overlaps_mask, overlap_mask) # TODO 將當前層重疊區域的mask賦值給修復mask inpaint_mask = overlaps_mask
修復
使用修復演算法cv2.INPAINT_TELEA，演算法思路是：先處理待修復區域邊緣上的畫素點，然後層層向內推進，直到修復完所有的畫素點。
# img是要修復的影象 inpaint_mask是上面提到的遮罩dst是修復好的影象 dst = cv2.inpaint(img, inpaint_mask, 3, cv2.INPAINT_TELEA)
延展
本文大概介紹了通過計算機視覺為主，深度學習為輔的圖片複雜前後景分離的解決方案，除了文中提到的部分，還有幾層輪廓捕獲的邏輯因為篇幅原因，未加展開，針對比較複雜的case，本方案已經能夠很好的實現圖層分離，但對於更加複雜的場景，比如邊緣顏色複雜度高，噪點多，邊緣輪廓不明顯等更復雜的case，分離的精確度還有很大的提升空間。
期待能夠聽到更多有趣的解決方案，歡迎交流。