File size: 3,105 Bytes
c9019cd |
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 |
# Copyright (c) 2022, National Diet Library, Japan
#
# This software is released under the CC BY 4.0.
# https://creativecommons.org/licenses/by/4.0/
import click
import json
import os
import sys
from cli.core import OcrInferencer
from cli.core import utils
@click.group()
@click.option('--debug', is_flag=True)
@click.pass_context
def cmd(ctx, debug):
ctx.obj['DEBUG'] = debug
@cmd.command()
@click.pass_context
def help(ctx):
if ctx.obj['DEBUG']:
click.echo('DEBUG MODE!')
click.echo('help!')
@cmd.command()
@click.pass_context
@click.argument('input_root')
@click.argument('output_root')
@click.option('-s', '--input_structure', type=click.Choice(['s', 'i', 't', 'w', 'f'], case_sensitive=True), default='s', help='Input directory structure type. s(single), i(intermediate_output), t(tosho_data), w(workstation), and f(image_file).')
@click.option('-p', '--proc_range', type=str, default='0..3', help='Inference process range to run. Default is "0..3".')
@click.option('-c', '--config_file', type=str, default='config.yml', help='Configuration yml file for inference. Default is "config.yml".')
@click.option('-i', '--save_image', type=bool, default=False, is_flag=True, help='Output result image file with text file.')
@click.option('-x', '--save_xml', type=bool, default=False, is_flag=True, help='Output result XML file with text file.')
@click.option('-d', '--dump', type=bool, default=False, is_flag=True, help='Dump all intermediate process output.')
def infer(ctx, input_root, output_root, config_file, proc_range, save_image, save_xml, input_structure, dump):
"""
\b
INPUT_ROOT \t: Input data directory for inference.
OUTPUT_ROOT \t: Output directory for inference result.
"""
click.echo('start inference !')
click.echo('input_root : {0}'.format(input_root))
click.echo('output_root : {0}'.format(output_root))
click.echo('config_file : {0}'.format(config_file))
cfg = {
'input_root': input_root,
'output_root': output_root,
'config_file': config_file,
'proc_range': proc_range,
'save_image': save_image,
'save_xml': save_xml,
'dump': dump,
'input_structure': input_structure
}
# check if input_root exists
if not os.path.exists(input_root):
print('INPUT_ROOT not found :{0}'.format(input_root), file=sys.stderr)
exit(0)
# parse command line option
infer_cfg = utils.parse_cfg(cfg)
if infer_cfg is None:
print('[ERROR] Config parse error :{0}'.format(input_root), file=sys.stderr)
exit(1)
# prepare output root derectory
infer_cfg['output_root'] = utils.mkdir_with_duplication_check(infer_cfg['output_root'])
# save inference option
with open(os.path.join(infer_cfg['output_root'], 'opt.json'), 'w') as fp:
json.dump(infer_cfg, fp, ensure_ascii=False, indent=4,
sort_keys=True, separators=(',', ': '))
# do inference
inferencer = OcrInferencer(infer_cfg)
inferencer.run()
def main():
cmd(obj={})
if __name__ == '__main__':
main()
|