File size: 3,905 Bytes
1da285f | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import multiprocessing
import glob
import time
import json
from tqdm import tqdm
import argparse
import os
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))
def main(args):
# 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_img_root = args.input_img_root
output_root = args.output_root
with open(args.ann, 'r') as f:
data = json.load(f)
os.makedirs(output_root, exist_ok=True)
# 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
input_imgs = [pjoin(input_img_root, img['file_name']) for img in data['images']]
# 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-bar': True}
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-bar': False}
is_ip = True
is_clf = False
is_ocr = True
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
start_time = time.time()
for input_img in tqdm(input_imgs, total=len(input_imgs)):
name = input_img.split('/')[-1].split('.')[-2]
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, method='paddle')
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 detect_merge.merge as merge
# compo_path = pjoin(output_root, 'ip', str(index) + '.json')
# ocr_path = pjoin(output_root, 'ocr', str(index) + '.json')
compo_path = pjoin(output_root, 'ip', name + '.json')
ocr_path = pjoin(output_root, 'ocr', name + '.json')
merge.merge(input_img, compo_path, ocr_path, output_root, is_remove_bar=key_params['remove-bar'], show=False)
num += 1
end_time = time.time()
print(f'Total inference time: {end_time - start_time} s')
os.makedirs(pjoin(output_root, 'time'), exist_ok=True)
with open(pjoin(output_root, 'time', 'inference time.txt'), 'w') as f:
f.write(f'First stage total inference time: {end_time - start_time} s')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input_img_root', type=str, required=True)
parser.add_argument('--output_root', type=str, required=True)
parser.add_argument('--ann', type=str, required=True)
args = parser.parse_args()
main(args)
|