ui2code中的深度學習+傳統算法應用

背景

在以前的文章中，咱們已經提到過團隊在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

文字抽離

若是是純色背景，文字區域很好抽離，但若是是複雜背景就比較麻煩了。舉個例子：

基於以上，咱們能拿到準確的文本信息，咱們逐一對各個文本信息作處理，文本的特徵仍是比較明顯的，好比說含有多個角點，在嘗試了多種算法：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，分離的精確度還有很大的提高空間。

原文連接