def det_recog_pp(self, result):
final_results = []
args = self.args
for arr, output, export, det_recog_result in zip(
args.arrays, args.output, args.export, result):
if output or args.imshow:
if self.kie_model:
res_img = det_recog_show_result(arr, det_recog_result)
else:
res_img = det_recog_show_result(arr,
det_recog_result,
out_file=output)
if args.imshow and not self.kie_model:
mmcv.imshow(res_img, 'inference results')
if not args.details:
simple_res = {}
simple_res['filename'] = det_recog_result['filename']
simple_res['text'] = [
x['text'] for x in det_recog_result['result']
]
final_result = simple_res
else:
final_result = det_recog_result
if export:
mmcv.dump(final_result, export, indent=4)
if args.print_result:
print(final_result, end='\n\n')
final_results.append(final_result)
return final_results
def main():
parser = ArgumentParser()
parser.add_argument('img', type=str, help='Input Image file.')
parser.add_argument('out_file',
type=str,
help='Output file name of the visualized image.')
parser.add_argument('--det-config',
type=str,
default='./configs/textdet/psenet/'
'psenet_r50_fpnf_600e_icdar2015.py',
help='Text detection config file.')
parser.add_argument('--det-ckpt',
type=str,
default='https://download.openmmlab.com/'
'mmocr/textdet/psenet/'
'psenet_r50_fpnf_600e_icdar2015_pretrain-eefd8fe6.pth',
help='Text detection checkpint file (local or url).')
parser.add_argument('--recog-config',
type=str,
default='./configs/textrecog/sar/'
'sar_r31_parallel_decoder_academic.py',
help='Text recognition config file.')
parser.add_argument('--recog-ckpt',
type=str,
default='https://download.openmmlab.com/'
'mmocr/textrecog/sar/'
'sar_r31_parallel_decoder_academic-dba3a4a3.pth',
help='Text recognition checkpint file (local or url).')
parser.add_argument('--batch-mode',
action='store_true',
help='Whether use batch mode for text recognition.')
parser.add_argument('--batch-size',
type=int,
default=4,
help='Batch size for text recognition inference '
'if batch_mode is True above.')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference.')
parser.add_argument('--imshow',
action='store_true',
help='Whether show image with OpenCV.')
parser.add_argument('--ocr-in-lines',
action='store_true',
help='Whether group ocr results in lines.')
args = parser.parse_args()
if args.device == 'cpu':
args.device = None
# build detect model
detect_model = init_detector(args.det_config,
args.det_ckpt,
device=args.device)
if hasattr(detect_model, 'module'):
detect_model = detect_model.module
if detect_model.cfg.data.test['type'] == 'ConcatDataset':
detect_model.cfg.data.test.pipeline = \
detect_model.cfg.data.test['datasets'][0].pipeline
# build recog model
recog_model = init_detector(args.recog_config,
args.recog_ckpt,
device=args.device)
if hasattr(recog_model, 'module'):
recog_model = recog_model.module
if recog_model.cfg.data.test['type'] == 'ConcatDataset':
recog_model.cfg.data.test.pipeline = \
recog_model.cfg.data.test['datasets'][0].pipeline
det_recog_result = det_and_recog_inference(args, detect_model, recog_model)
print(f'result: {det_recog_result}')
mmcv.dump(det_recog_result,
args.out_file + '.json',
ensure_ascii=False,
indent=4)
if args.ocr_in_lines:
res = det_recog_result['result']
res = stitch_boxes_into_lines(res, 10, 0.5)
det_recog_result['result'] = res
mmcv.dump(det_recog_result,
args.out_file + '.line.json',
ensure_ascii=False,
indent=4)
img = det_recog_show_result(args.img, det_recog_result)
mmcv.imwrite(img, args.out_file)
if args.imshow:
mmcv.imshow(img, 'predicted results')