add id for each element; save the children id instead of entire eleme…

…nts; check containment at the end

detect_merge/Element.py

@@ -3,15 +3,16 @@
 
 
 class Element:
-    def __init__(self, corner, category, text_content=None):
+    def __init__(self, id, corner, category, text_content=None):
+        self.id = id
         self.category = category
         self.col_min, self.row_min, self.col_max, self.row_max = corner
         self.width = self.col_max - self.col_min
         self.height = self.row_max - self.row_min
         self.area = self.width * self.height
 
         self.text_content = text_content
-        self.is_child = False
+        self.parent_id = None
         self.children = []  # list of elements
 
     def init_bound(self):
@@ -23,15 +24,17 @@ def put_bbox(self):
         return self.col_min, self.row_min, self.col_max, self.row_max
 
     def wrap_info(self):
-        info = {'class': self.category, 'height': self.height, 'width': self.width,
+        info = {'id':self.id, 'class': self.category, 'height': self.height, 'width': self.width,
                 'position': {'column_min': self.col_min, 'row_min': self.row_min, 'column_max': self.col_max,
                              'row_max': self.row_max}}
         if self.text_content is not None:
             info['text_content'] = self.text_content
         if len(self.children) > 0:
             info['children'] = []
             for child in self.children:
-                info['children'].append(child.wrap_info())
+                info['children'].append(child.id)
+        if self.parent_id is not None:
+            info['parent'] = self.parent_id
         return info
 
     def resize(self, resize_ratio):
@@ -41,12 +44,12 @@ def resize(self, resize_ratio):
         self.row_max = int(self.row_max * resize_ratio)
         self.init_bound()
 
-    def element_merge(self, element_b, new_element=False, new_category=None):
+    def element_merge(self, element_b, new_element=False, new_category=None, new_id=None):
         col_min_a, row_min_a, col_max_a, row_max_a = self.put_bbox()
         col_min_b, row_min_b, col_max_b, row_max_b = element_b.put_bbox()
         new_corner = (min(col_min_a, col_min_b), min(row_min_a, row_min_b), max(col_max_a, col_max_b), max(row_max_a, row_max_b))
         if new_element:
-            return Element(new_corner, new_category)
+            return Element(new_id, new_corner, new_category)
         else:
             self.col_min, self.row_min, self.col_max, self.row_max = new_corner
             self.init_bound()
@@ -81,19 +84,19 @@ def element_relation(self, element_b, bias=0):
         if ioa == 0:
             return 0
         # a in b
-        if ioa > 0.6:
+        if ioa >= 1:
             return -1
         # b in a
-        if iob > 0.6:
+        if iob >= 1:
             return 1
         return 2
 
     def visualize_element(self, img, color=(0, 255, 0), line=1, show=False):
         loc = self.put_bbox()
         cv2.rectangle(img, loc[:2], loc[2:], color, line)
-        for child in self.children:
-            child.visualize_element(img, color=(255, 0, 255), line=line)
+        # for child in self.children:
+        #     child.visualize_element(img, color=(255, 0, 255), line=line)
         if show:
             cv2.imshow('element', img)
             cv2.waitKey(0)
-            cv2.destroyWindow('element')
+            cv2.destroyWindow('element')

detect_merge/merge.py

@@ -10,7 +10,7 @@
 
 
 def show_elements(org_img, eles, show=False, win_name='element', wait_key=0, shown_resize=None, line=2):
-    color_map = {'Text':(0, 0, 255), 'Compo':(0, 255, 0), 'Text Content':(255, 0, 255)}
+    color_map = {'Text':(0, 0, 255), 'Compo':(0, 255, 0), 'Block':(0, 255, 255), 'Text Content':(255, 0, 255)}
     img = org_img.copy()
     for ele in eles:
         color = color_map[ele.category]
@@ -21,21 +21,24 @@ def show_elements(org_img, eles, show=False, win_name='element', wait_key=0, sho
     if show:
         cv2.imshow(win_name, img_resize)
         cv2.waitKey(wait_key)
+        if wait_key == 0:
+            cv2.destroyWindow(win_name)
     return img_resize
 
 
-def save_elements(output_dir, elements, img_shape):
+def save_elements(output_file, elements, img_shape):
     components = {'compos': [], 'img_shape': img_shape}
-    clip_dir = pjoin(output_dir, 'clips')
-
     for i, ele in enumerate(elements):
         c = ele.wrap_info()
-        c['id'] = i
-        c['clip_path'] = pjoin(clip_dir, c['class'], str(i) + '.jpg')
+        # c['id'] = i
         components['compos'].append(c)
+    json.dump(components, open(output_file, 'w'), indent=4)
+    return components
+
 
-    json.dump(components, open(pjoin(output_dir, 'compo.json'), 'w'), indent=4)
-    return components['compos']
+def reassign_ids(elements):
+    for i, element in enumerate(elements):
+        element.id = i
 
 
 def refine_texts(texts, img_shape):
@@ -69,17 +72,30 @@ def refine_elements(compos, texts, intersection_bias=2, containment_ratio=0.8):
                 if iob >= containment_ratio:
                     contained_texts.append(text)
         if is_valid and text_area / compo.area < containment_ratio:
-            for t in contained_texts:
-                t.is_child = True
-            compo.children += contained_texts
+            # for t in contained_texts:
+            #     t.parent_id = compo.id
+            # compo.children += contained_texts
             elements.append(compo)
 
-    for text in texts:
-        if not text.is_child:
-            elements.append(text)
+    elements += texts
+    # for text in texts:
+    #     if not text.is_child:
+    #         elements.append(text)
     return elements
 
 
+def check_containment(elements):
+    for i in range(len(elements) - 1):
+        for j in range(i + 1, len(elements)):
+            relation = elements[i].element_relation(elements[j], bias=2)
+            if relation == -1:
+                elements[j].children.append(elements[i])
+                elements[i].parent_id = elements[j].id
+            if relation == 1:
+                elements[i].children.append(elements[j])
+                elements[j].parent_id = elements[i].id
+
+
 def remove_top_bar(elements, img_height):
     new_elements = []
     for ele in elements:
@@ -138,14 +154,17 @@ def merge(img_path, compo_path, text_path, output_root=None, is_remove_top=True,
     text_json = json.load(open(text_path, 'r'))
 
     # load text and non-text compo
+    ele_id = 0
     compos = []
     for compo in compo_json['compos']:
-        element = Element((compo['column_min'], compo['row_min'], compo['column_max'], compo['row_max']), compo['class'])
+        element = Element(ele_id, (compo['column_min'], compo['row_min'], compo['column_max'], compo['row_max']), compo['class'])
         compos.append(element)
+        ele_id += 1
     texts = []
     for text in text_json['texts']:
-        element = Element((text['column_min'], text['row_min'], text['column_max'], text['row_max']), 'Text', text_content=text['content'])
+        element = Element(ele_id, (text['column_min'], text['row_min'], text['column_max'], text['row_max']), 'Text', text_content=text['content'])
         texts.append(element)
+        ele_id += 1
     if compo_json['img_shape'] != text_json['img_shape']:
         resize_ratio = compo_json['img_shape'][0] / text_json['img_shape'][0]
         for text in texts:
@@ -160,15 +179,13 @@ def merge(img_path, compo_path, text_path, output_root=None, is_remove_top=True,
     texts = refine_texts(texts, compo_json['img_shape'])
     elements = refine_elements(compos, texts)
     if is_remove_top: elements = remove_top_bar(elements, img_height=compo_json['img_shape'][0])
+    reassign_ids(elements)
+    check_containment(elements)
     board = show_elements(img_resize, elements, show=show, win_name='elements after merging', wait_key=wait_key)
 
     # save all merged elements, clips and blank background
-    if output_root is not None:
-        compos_json = save_elements(output_root, elements, img_resize.shape)
-        compos_clip_and_fill(pjoin(output_root, 'clips'), img_resize, compos_json)
-        cv2.imwrite(pjoin(output_root, 'result.jpg'), board)
-        print('Merge Complete and Save to', pjoin(output_root, 'result.jpg'), time.ctime(), '\n')
-    else:
-        print('Merge Complete', time.ctime(), '\n')
-
-# merge('../data/input/2.jpg', '../data/output/ip/2.json', '../data/output/ocr/2.json', '../data/output', show=True)
+    save_elements(pjoin(output_root, 'compo.json'), elements, img_resize.shape)
+    # compos_clip_and_fill(pjoin(output_root, 'clips'), img_resize, compos_json)
+    cv2.imwrite(pjoin(output_root, 'result.jpg'), board)
+    print('[Merge Completed] Input: %s Output: %s' % (img_path, pjoin(output_root, 'compo.jpg')), time.ctime(), '\n')
+