|
|
|
|
|
import argparse |
|
|
import math |
|
|
import os.path as osp |
|
|
|
|
|
import mmcv |
|
|
|
|
|
from mmocr.utils import convert_annotations |
|
|
|
|
|
|
|
|
def parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description='Generate training and validation set of TextOCR ') |
|
|
parser.add_argument('root_path', help='Root dir path of TextOCR') |
|
|
args = parser.parse_args() |
|
|
return args |
|
|
|
|
|
|
|
|
def collect_textocr_info(root_path, annotation_filename, print_every=1000): |
|
|
|
|
|
annotation_path = osp.join(root_path, annotation_filename) |
|
|
if not osp.exists(annotation_path): |
|
|
raise Exception( |
|
|
f'{annotation_path} not exists, please check and try again.') |
|
|
|
|
|
annotation = mmcv.load(annotation_path) |
|
|
|
|
|
|
|
|
img_infos = [] |
|
|
for i, img_info in enumerate(annotation['imgs'].values()): |
|
|
if i > 0 and i % print_every == 0: |
|
|
print(f'{i}/{len(annotation["imgs"].values())}') |
|
|
|
|
|
img_info['segm_file'] = annotation_path |
|
|
ann_ids = annotation['imgToAnns'][img_info['id']] |
|
|
anno_info = [] |
|
|
for ann_id in ann_ids: |
|
|
ann = annotation['anns'][ann_id] |
|
|
|
|
|
|
|
|
text_label = ann['utf8_string'] |
|
|
iscrowd = 1 if text_label == '.' else 0 |
|
|
|
|
|
x, y, w, h = ann['bbox'] |
|
|
x, y = max(0, math.floor(x)), max(0, math.floor(y)) |
|
|
w, h = math.ceil(w), math.ceil(h) |
|
|
bbox = [x, y, w, h] |
|
|
segmentation = [max(0, int(x)) for x in ann['points']] |
|
|
anno = dict( |
|
|
iscrowd=iscrowd, |
|
|
category_id=1, |
|
|
bbox=bbox, |
|
|
area=ann['area'], |
|
|
segmentation=[segmentation]) |
|
|
anno_info.append(anno) |
|
|
img_info.update(anno_info=anno_info) |
|
|
img_infos.append(img_info) |
|
|
return img_infos |
|
|
|
|
|
|
|
|
def main(): |
|
|
args = parse_args() |
|
|
root_path = args.root_path |
|
|
print('Processing training set...') |
|
|
training_infos = collect_textocr_info(root_path, 'TextOCR_0.1_train.json') |
|
|
convert_annotations(training_infos, |
|
|
osp.join(root_path, 'instances_training.json')) |
|
|
print('Processing validation set...') |
|
|
val_infos = collect_textocr_info(root_path, 'TextOCR_0.1_val.json') |
|
|
convert_annotations(val_infos, osp.join(root_path, 'instances_val.json')) |
|
|
print('Finish') |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
