mirror of
https://github.com/open-mmlab/mmocr.git
synced 2025-06-03 21:54:47 +08:00
33cbc9b92f
* [Enhancement] Support batch visualization & dumping in Inferencer * fix empty det output * Update mmocr/apis/inferencers/base_mmocr_inferencer.py Co-authored-by: liukuikun <24622904+Harold-lkk@users.noreply.github.com> * [Docs] Inferencer docs * fix * Support weight_list * add req * improve md * inferencers.md * update * add tab * refine * polish * add cn docs * js * js * js * fix ch docs * translate * translate * finish * fix * fix * fix * update * standard inferencer * update docs * update docs * update docs * update docs * update docs * update docs * en * update * update * update * update * fix * apply sugg --------- Co-authored-by: liukuikun <24622904+Harold-lkk@users.noreply.github.com>
101 lines
3.2 KiB
Python
Executable File
101 lines
3.2 KiB
Python
Executable File
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from argparse import ArgumentParser
|
|
|
|
from mmocr.apis.inferencers import MMOCRInferencer
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser()
|
|
parser.add_argument(
|
|
'inputs', type=str, help='Input image file or folder path.')
|
|
parser.add_argument(
|
|
'--out-dir',
|
|
type=str,
|
|
default='results/',
|
|
help='Output directory of results.')
|
|
parser.add_argument(
|
|
'--det',
|
|
type=str,
|
|
default=None,
|
|
help='Pretrained text detection algorithm. It\'s the path to the '
|
|
'config file or the model name defined in metafile.')
|
|
parser.add_argument(
|
|
'--det-weights',
|
|
type=str,
|
|
default=None,
|
|
help='Path to the custom checkpoint file of the selected det model. '
|
|
'If it is not specified and "det" is a model name of metafile, the '
|
|
'weights will be loaded from metafile.')
|
|
parser.add_argument(
|
|
'--rec',
|
|
type=str,
|
|
default=None,
|
|
help='Pretrained text recognition algorithm. It\'s the path to the '
|
|
'config file or the model name defined in metafile.')
|
|
parser.add_argument(
|
|
'--rec-weights',
|
|
type=str,
|
|
default=None,
|
|
help='Path to the custom checkpoint file of the selected recog model. '
|
|
'If it is not specified and "rec" is a model name of metafile, the '
|
|
'weights will be loaded from metafile.')
|
|
parser.add_argument(
|
|
'--kie',
|
|
type=str,
|
|
default=None,
|
|
help='Pretrained key information extraction algorithm. It\'s the path'
|
|
'to the config file or the model name defined in metafile.')
|
|
parser.add_argument(
|
|
'--kie-weights',
|
|
type=str,
|
|
default=None,
|
|
help='Path to the custom checkpoint file of the selected kie model. '
|
|
'If it is not specified and "kie" is a model name of metafile, the '
|
|
'weights will be loaded from metafile.')
|
|
parser.add_argument(
|
|
'--device',
|
|
type=str,
|
|
default=None,
|
|
help='Device used for inference. '
|
|
'If not specified, the available device will be automatically used.')
|
|
parser.add_argument(
|
|
'--batch-size', type=int, default=1, help='Inference batch size.')
|
|
parser.add_argument(
|
|
'--show',
|
|
action='store_true',
|
|
help='Display the image in a popup window.')
|
|
parser.add_argument(
|
|
'--print-result',
|
|
action='store_true',
|
|
help='Whether to print the results.')
|
|
parser.add_argument(
|
|
'--save_pred',
|
|
action='store_true',
|
|
help='Save the inference results to out_dir.')
|
|
parser.add_argument(
|
|
'--save_vis',
|
|
action='store_true',
|
|
help='Save the visualization results to out_dir.')
|
|
|
|
call_args = vars(parser.parse_args())
|
|
|
|
init_kws = [
|
|
'det', 'det_weights', 'rec', 'rec_weights', 'kie', 'kie_weights',
|
|
'device'
|
|
]
|
|
init_args = {}
|
|
for init_kw in init_kws:
|
|
init_args[init_kw] = call_args.pop(init_kw)
|
|
|
|
return init_args, call_args
|
|
|
|
|
|
def main():
|
|
init_args, call_args = parse_args()
|
|
ocr = MMOCRInferencer(**init_args)
|
|
ocr(**call_args)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|