Spaces:
Running
Running
File size: 3,375 Bytes
9bf4bd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import json
from typing import List, Tuple
from mmocr.datasets import RecogLMDBDataset
from mmocr.utils import StringStripper, dump_ocr_data, recog_anno_to_imginfo
def parse_legacy_data(in_path: str,
format: str) -> Tuple[List[str], List[str]]:
"""Load legacy data and return a list of file paths and labels.
Args:
in_path (str): Path to annotation file.
format (str): Annotation format. Choices are 'txt', 'json' and 'lmdb'.
For 'lmdb' format, the lmdb file should only contains labels. For
lmdb file with labels and images, the conversion is unnecessary.
Returns:
tuple(list[str], list[str]): File paths and labels.
"""
file_paths = []
labels = []
strip_cls = StringStripper()
if format == 'lmdb':
dataset = RecogLMDBDataset(
in_path,
parser_cfg=dict(type='LineJsonParser', keys=['filename', 'text']))
for data_info in dataset:
file_path = data_info['img_path']
label = data_info['instances'][0]['text']
file_path = strip_cls(file_path)
label = strip_cls(label)
# MJ's file_path starts with './'
if file_path.startswith('./'):
file_path = file_path[2:]
file_paths.append(file_path)
labels.append(label)
return file_paths, labels
else:
with open(in_path) as f:
if format == 'txt':
for line in f:
line = strip_cls(line)
file_path, label = line.split()[:2]
# MJ's file_path starts with './'
if file_path.startswith('./'):
file_path = file_path[2:]
file_paths.append(file_path)
labels.append(label)
elif format == 'jsonl':
for line in f:
datum = json.loads(line)
file_path = datum['filename']
# MJ's file_path starts with './'
if file_path.startswith('./'):
file_path = file_path[2:]
file_paths.append(file_path)
labels.append(datum['text'])
return file_paths, labels
def parse_args():
"""Parse input arguments."""
parser = argparse.ArgumentParser(
description='Convert annotations for'
'text recognition tasks in MMOCR 0.x into the latest openmmlab format.'
)
parser.add_argument(
'in_path', help='The path to legacy recognition data file')
parser.add_argument(
'out_path', help='The output json path in openmmlab format')
parser.add_argument(
'--format',
choices=['txt', 'jsonl', 'lmdb'],
type=str,
default='txt',
help='Legacy data format')
args = parser.parse_args()
if args.out_path.split('.')[-1] != 'json':
raise ValueError('The output path must be a json file.')
return args
def main():
args = parse_args()
file_paths, labels = parse_legacy_data(args.in_path, args.format)
img_infos = recog_anno_to_imginfo(file_paths, labels)
dump_ocr_data(img_infos, args.out_path, 'textrecog')
print('finish')
if __name__ == '__main__':
main()
|