Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import argparse | |
| import os | |
| import os.path as osp | |
| import xml.etree.ElementTree as ET | |
| import cv2 | |
| from mmocr.utils.fileio import list_to_file | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='Generate testset of svt by cropping box image.') | |
| parser.add_argument( | |
| 'root_path', | |
| help='Root dir path of svt, where test.xml in,' | |
| 'for example, "data/mixture/svt/svt1/"') | |
| parser.add_argument( | |
| '--resize', | |
| action='store_true', | |
| help='Whether resize cropped image to certain size.') | |
| parser.add_argument('--height', default=32, help='Resize height.') | |
| parser.add_argument('--width', default=100, help='Resize width.') | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = parse_args() | |
| root_path = args.root_path | |
| # inputs | |
| src_label_file = osp.join(root_path, 'test.xml') | |
| if not osp.exists(src_label_file): | |
| raise Exception( | |
| f'{src_label_file} not exists, please check and try again.') | |
| src_image_root = root_path | |
| # outputs | |
| dst_label_file = osp.join(root_path, 'test_label.txt') | |
| dst_image_root = osp.join(root_path, 'image') | |
| os.makedirs(dst_image_root, exist_ok=True) | |
| tree = ET.parse(src_label_file) | |
| root = tree.getroot() | |
| index = 1 | |
| lines = [] | |
| total_img_num = len(root) | |
| i = 1 | |
| for image_node in root.findall('image'): | |
| image_name = image_node.find('imageName').text | |
| print(f'[{i}/{total_img_num}] Process image: {image_name}') | |
| i += 1 | |
| lexicon = image_node.find('lex').text.lower() | |
| lexicon_list = lexicon.split(',') | |
| lex_size = len(lexicon_list) | |
| src_img = cv2.imread(osp.join(src_image_root, image_name)) | |
| for rectangle in image_node.find('taggedRectangles'): | |
| x = int(rectangle.get('x')) | |
| y = int(rectangle.get('y')) | |
| w = int(rectangle.get('width')) | |
| h = int(rectangle.get('height')) | |
| rb, re = max(0, y), max(0, y + h) | |
| cb, ce = max(0, x), max(0, x + w) | |
| dst_img = src_img[rb:re, cb:ce] | |
| text_label = rectangle.find('tag').text.lower() | |
| if args.resize: | |
| dst_img = cv2.resize(dst_img, (args.width, args.height)) | |
| dst_img_name = f'img_{index:04}' + '.jpg' | |
| index += 1 | |
| dst_img_path = osp.join(dst_image_root, dst_img_name) | |
| cv2.imwrite(dst_img_path, dst_img) | |
| lines.append(f'{osp.basename(dst_image_root)}/{dst_img_name} ' | |
| f'{text_label} {lex_size} {lexicon}') | |
| list_to_file(dst_label_file, lines) | |
| print(f'Finish to generate svt testset, ' | |
| f'with label file {dst_label_file}') | |
| if __name__ == '__main__': | |
| main() | |