Spaces:
Sleeping
Sleeping
# Copyright (c) OpenMMLab. All rights reserved. | |
import argparse | |
import os.path as osp | |
from functools import partial | |
import mmengine | |
import numpy as np | |
from mmocr.utils import bezier2polygon, sort_points | |
# The default dictionary used by CurvedSynthText | |
dict95 = [ | |
' ', '!', '"', '#', '$', '%', '&', '\'', '(', ')', '*', '+', ',', '-', '.', | |
'/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=', | |
'>', '?', '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', | |
'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', | |
'\\', ']', '^', '_', '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', | |
'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', | |
'z', '{', '|', '}', '~' | |
] | |
UNK = len(dict95) | |
EOS = UNK + 1 | |
def digit2text(rec): | |
res = [] | |
for d in rec: | |
assert d <= EOS | |
if d == EOS: | |
break | |
if d == UNK: | |
print('Warning: Has a UNK character') | |
res.append('口') # Or any special character not in the target dict | |
res.append(dict95[d]) | |
return ''.join(res) | |
def modify_annotation(ann, num_sample, start_img_id=0, start_ann_id=0): | |
ann['text'] = digit2text(ann.pop('rec')) | |
# Get hide egmentation points | |
polygon_pts = bezier2polygon(ann['bezier_pts'], num_sample=num_sample) | |
ann['segmentation'] = np.asarray(sort_points(polygon_pts)).reshape( | |
1, -1).tolist() | |
ann['image_id'] += start_img_id | |
ann['id'] += start_ann_id | |
return ann | |
def modify_image_info(image_info, path_prefix, start_img_id=0): | |
image_info['file_name'] = osp.join(path_prefix, image_info['file_name']) | |
image_info['id'] += start_img_id | |
return image_info | |
def parse_args(): | |
parser = argparse.ArgumentParser( | |
description='Convert CurvedSynText150k to COCO format') | |
parser.add_argument('root_path', help='CurvedSynText150k root path') | |
parser.add_argument('-o', '--out-dir', help='Output path') | |
parser.add_argument( | |
'-n', | |
'--num-sample', | |
type=int, | |
default=4, | |
help='Number of sample points at each Bezier curve.') | |
parser.add_argument( | |
'--nproc', default=1, type=int, help='Number of processes') | |
args = parser.parse_args() | |
return args | |
def convert_annotations(data, | |
path_prefix, | |
num_sample, | |
nproc, | |
start_img_id=0, | |
start_ann_id=0): | |
modify_image_info_with_params = partial( | |
modify_image_info, path_prefix=path_prefix, start_img_id=start_img_id) | |
modify_annotation_with_params = partial( | |
modify_annotation, | |
num_sample=num_sample, | |
start_img_id=start_img_id, | |
start_ann_id=start_ann_id) | |
if nproc > 1: | |
data['annotations'] = mmengine.track_parallel_progress( | |
modify_annotation_with_params, data['annotations'], nproc=nproc) | |
data['images'] = mmengine.track_parallel_progress( | |
modify_image_info_with_params, data['images'], nproc=nproc) | |
else: | |
data['annotations'] = mmengine.track_progress( | |
modify_annotation_with_params, data['annotations']) | |
data['images'] = mmengine.track_progress( | |
modify_image_info_with_params, | |
data['images'], | |
) | |
data['categories'] = [{'id': 1, 'name': 'text'}] | |
return data | |
def main(): | |
args = parse_args() | |
root_path = args.root_path | |
out_dir = args.out_dir if args.out_dir else root_path | |
mmengine.mkdir_or_exist(out_dir) | |
anns = mmengine.load(osp.join(root_path, 'train1.json')) | |
data1 = convert_annotations(anns, 'syntext_word_eng', args.num_sample, | |
args.nproc) | |
# Get the maximum image id from data1 | |
start_img_id = max(data1['images'], key=lambda x: x['id'])['id'] + 1 | |
start_ann_id = max(data1['annotations'], key=lambda x: x['id'])['id'] + 1 | |
anns = mmengine.load(osp.join(root_path, 'train2.json')) | |
data2 = convert_annotations( | |
anns, | |
'emcs_imgs', | |
args.num_sample, | |
args.nproc, | |
start_img_id=start_img_id, | |
start_ann_id=start_ann_id) | |
data1['images'] += data2['images'] | |
data1['annotations'] += data2['annotations'] | |
mmengine.dump(data1, osp.join(out_dir, 'instances_training.json')) | |
if __name__ == '__main__': | |
main() | |