|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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() |
|
|
|
