mirror of https://github.com/open-mmlab/mmocr.git
109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import math
|
|
import os
|
|
import os.path as osp
|
|
from functools import partial
|
|
|
|
import mmcv
|
|
|
|
from mmocr.utils.fileio import list_to_file
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='Generate training and validation set of TextOCR '
|
|
'by cropping box image.')
|
|
parser.add_argument('root_path', help='Root dir path of TextOCR')
|
|
parser.add_argument(
|
|
'n_proc', default=1, type=int, help='Number of processes to run')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def process_img(args, src_image_root, dst_image_root):
|
|
# Dirty hack for multi-processing
|
|
img_idx, img_info, anns = args
|
|
src_img = mmcv.imread(osp.join(src_image_root, img_info['file_name']))
|
|
labels = []
|
|
for ann_idx, ann in enumerate(anns):
|
|
text_label = ann['utf8_string']
|
|
|
|
# Ignore illegible or non-English words
|
|
if text_label == '.':
|
|
continue
|
|
|
|
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)
|
|
dst_img = src_img[y:y + h, x:x + w]
|
|
dst_img_name = f'img_{img_idx}_{ann_idx}.jpg'
|
|
dst_img_path = osp.join(dst_image_root, dst_img_name)
|
|
mmcv.imwrite(dst_img, dst_img_path)
|
|
labels.append(f'{osp.basename(dst_image_root)}/{dst_img_name}'
|
|
f' {text_label}')
|
|
return labels
|
|
|
|
|
|
def convert_textocr(root_path,
|
|
dst_image_path,
|
|
dst_label_filename,
|
|
annotation_filename,
|
|
img_start_idx=0,
|
|
nproc=1):
|
|
|
|
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.')
|
|
src_image_root = root_path
|
|
|
|
# outputs
|
|
dst_label_file = osp.join(root_path, dst_label_filename)
|
|
dst_image_root = osp.join(root_path, dst_image_path)
|
|
os.makedirs(dst_image_root, exist_ok=True)
|
|
|
|
annotation = mmcv.load(annotation_path)
|
|
|
|
process_img_with_path = partial(
|
|
process_img,
|
|
src_image_root=src_image_root,
|
|
dst_image_root=dst_image_root)
|
|
tasks = []
|
|
for img_idx, img_info in enumerate(annotation['imgs'].values()):
|
|
ann_ids = annotation['imgToAnns'][img_info['id']]
|
|
anns = [annotation['anns'][ann_id] for ann_id in ann_ids]
|
|
tasks.append((img_idx + img_start_idx, img_info, anns))
|
|
labels_list = mmcv.track_parallel_progress(
|
|
process_img_with_path, tasks, keep_order=True, nproc=nproc)
|
|
final_labels = []
|
|
for label_list in labels_list:
|
|
final_labels += label_list
|
|
list_to_file(dst_label_file, final_labels)
|
|
return len(annotation['imgs'])
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
root_path = args.root_path
|
|
print('Processing training set...')
|
|
num_train_imgs = convert_textocr(
|
|
root_path=root_path,
|
|
dst_image_path='image',
|
|
dst_label_filename='train_label.txt',
|
|
annotation_filename='TextOCR_0.1_train.json',
|
|
nproc=args.n_proc)
|
|
print('Processing validation set...')
|
|
convert_textocr(
|
|
root_path=root_path,
|
|
dst_image_path='image',
|
|
dst_label_filename='val_label.txt',
|
|
annotation_filename='TextOCR_0.1_val.json',
|
|
img_start_idx=num_train_imgs,
|
|
nproc=args.n_proc)
|
|
print('Finish')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|