Spaces:
Running
Running
File size: 2,506 Bytes
a383d0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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
|