Spaces:
Runtime error
Runtime error
| # Author: Parag Mali | |
| # This script scales the detected boxes back to given width and height of the page | |
| # read the image | |
| import sys | |
| sys.path.extend(['/home/psm2208/code', '/home/psm2208/code']) | |
| import cv2 | |
| import os | |
| import numpy as np | |
| from multiprocessing import Pool | |
| from gtdb import feature_extractor | |
| import argparse | |
| def parse_args(): | |
| # Default parameters | |
| parser = argparse.ArgumentParser( | |
| description='Stitching method') | |
| parser.add_argument('--data_file', default='test', | |
| type=str, help='choose one') | |
| parser.add_argument('--output_dir', default='.', | |
| help='Output directory path') | |
| parser.add_argument('--math_dir', required=True, | |
| type=str, help='detections dir') | |
| parser.add_argument('--math_ext', default='.csv', | |
| help='Extention of detection files') | |
| parser.add_argument('--home_data', default='/home/psm2208/data/GTDB/', type = str, | |
| help='database dir') | |
| parser.add_argument('--home_eval', default='/home/psm2208/code/eval/', type = str, | |
| help='Eval dir') | |
| parser.add_argument('--home_images', default='/home/psm2208/data/GTDB/images/', type = str, | |
| help='Images dir') | |
| parser.add_argument('--home_anno', default='/home/psm2208/data/GTDB/annotations/', type = str, | |
| help='Annotations dir') | |
| parser.add_argument('--home_char', default='/home/psm2208/data/GTDB/char_annotations/', type = str, | |
| help='Char anno dir') | |
| parser.add_argument('--num_workers', default=4, type=int, help='Number of workers') | |
| parser.add_argument('--type', default='math', type=str, help='Math or text') | |
| return parser.parse_args() | |
| def read_math(args, pdf_name): | |
| math_file = os.path.join(args.math_dir, pdf_name + args.math_ext) | |
| data = np.array([]) | |
| if os.path.exists(math_file): | |
| data = np.genfromtxt(math_file, delimiter=',') | |
| # if there is only one entry convert it to correct form required | |
| if len(data.shape) == 1: | |
| data = data.reshape(1, -1) | |
| return data.astype(int) | |
| def scale(params): | |
| try: | |
| args, math_regions, pdf_name, page_num = params | |
| print('Processing ', pdf_name, ' > ', page_num) | |
| image = cv2.imread(os.path.join(args.home_images, pdf_name, str(int(page_num + 1)) + ".png")) | |
| height = image.shape[0] | |
| width = image.shape[1] | |
| new_math = [] | |
| for math in math_regions: | |
| box = [0,0,0,0] | |
| box[0] = math[0] * width / 512 # 512 is ip image size | |
| box[1] = math[1] * height / 512 | |
| box[2] = math[2] * width / 512 | |
| box[3] = math[3] * height / 512 | |
| if feature_extractor.width(box) > 0 and feature_extractor.height(box) > 0: | |
| new_math.append(box) | |
| return new_math | |
| except Exception as e: | |
| print('Error while processing ', pdf_name, ' > ', page_num, sys.exc_info()) | |
| return [] | |
| def scale_boxes(args): | |
| pdf_list = [] | |
| pdf_names_file = open(args.data_file, 'r') | |
| for pdf_name in pdf_names_file: | |
| pdf_name = pdf_name.strip() | |
| if pdf_name != '': | |
| pdf_list.append(pdf_name) | |
| regions = {} | |
| for pdf_name in pdf_list: | |
| regions[pdf_name] = read_math(args, pdf_name) | |
| voting_ip_list = [] | |
| for pdf_name in pdf_list: | |
| pages = np.unique(regions[pdf_name][:, 0]) | |
| #args, math_regions, pdf_name, page_num | |
| for page_num in pages: | |
| current_math = regions[pdf_name][regions[pdf_name][:,0] == page_num] | |
| current_math = np.delete(current_math, 0, 1) | |
| voting_ip_list.append([args, current_math, pdf_name, page_num]) | |
| pool = Pool(processes=args.num_workers) | |
| out = pool.map(scale, voting_ip_list) | |
| for ip, final_math in zip(voting_ip_list, out): | |
| pdf_name = ip[2] | |
| page_num = ip[3] | |
| col = np.array([int(page_num)] * len(final_math)) | |
| final_math = np.concatenate((col[:, np.newaxis], final_math), axis=1) | |
| math_file_path = os.path.join(args.output_dir, pdf_name + '.csv') | |
| if not os.path.exists(os.path.dirname(math_file_path)): | |
| os.makedirs(os.path.dirname(math_file_path)) | |
| math_file = open(math_file_path, 'a') | |
| np.savetxt(math_file, final_math, fmt='%.2f', delimiter=',') | |
| math_file.close() | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| scale_boxes(args) | |