Initial commit

UIED/run_batch.py

@@ -0,0 +1,74 @@
+import multiprocessing
+import glob
+import time
+import json
+from tqdm import tqdm
+from os.path import join as pjoin, exists
+import cv2
+
+import detect_compo.ip_region_proposal as ip
+
+
+def resize_height_by_longest_edge(img_path, resize_length=800):
+    org = cv2.imread(img_path)
+    height, width = org.shape[:2]
+    if height > width:
+        return resize_length
+    else:
+        return int(resize_length * (height / width))
+
+
+if __name__ == '__main__':
+    # initialization
+    input_img_root = "E:/Mulong/Datasets/rico/combined"
+    output_root = "E:/Mulong/Result/rico/rico_uied/rico_new_uied_v3"
+    data = json.load(open('E:/Mulong/Datasets/rico/instances_test.json', 'r'))
+
+    input_imgs = [pjoin(input_img_root, img['file_name'].split('/')[-1]) for img in data['images']]
+    input_imgs = sorted(input_imgs, key=lambda x: int(x.split('/')[-1][:-4]))  # sorted by index
+
+    key_params = {'min-grad': 10, 'ffl-block': 5, 'min-ele-area': 50, 'merge-contained-ele': True,
+                  'max-word-inline-gap': 10, 'max-line-ingraph-gap': 4, 'remove-top-bar': True}
+
+    is_ip = False
+    is_clf = False
+    is_ocr = False
+    is_merge = True
+
+    # Load deep learning models in advance
+    compo_classifier = None
+    if is_ip and is_clf:
+        compo_classifier = {}
+        from cnn.CNN import CNN
+        # compo_classifier['Image'] = CNN('Image')
+        compo_classifier['Elements'] = CNN('Elements')
+        # compo_classifier['Noise'] = CNN('Noise')
+    ocr_model = None
+    if is_ocr:
+        import detect_text.text_detection as text
+
+    # set the range of target inputs' indices
+    num = 0
+    start_index = 30800  # 61728
+    end_index = 100000
+    for input_img in input_imgs:
+        resized_height = resize_height_by_longest_edge(input_img)
+        index = input_img.split('/')[-1][:-4]
+        if int(index) < start_index:
+            continue
+        if int(index) > end_index:
+            break
+
+        if is_ocr:
+            text.text_detection(input_img, output_root, show=False)
+
+        if is_ip:
+            ip.compo_detection(input_img, output_root, key_params,  classifier=compo_classifier, resize_by_height=resized_height, show=False)
+
+        if is_merge:
+            import merge
+            compo_path = pjoin(output_root, 'ip', str(index) + '.json')
+            ocr_path = pjoin(output_root, 'ocr', str(index) + '.json')
+            merge.merge(input_img, compo_path, ocr_path, output_root, is_remove_top=key_params['remove-top-bar'], show=True)
+
+        num += 1