Spaces:

robin0307
/

MMOCR

Build error

App Files Files Community

MMOCR / app.py

robin0307

Update app.py

f7de557 over 1 year ago

raw history blame contribute delete

No virus

3.66 kB

	import os
	os.system('pip install gradio==3.9.1')
	import torch

	print(torch.__version__)
	# torch_ver, cuda_ver = torch.__version__.split('+')
	# os.system('pip list')
	# os.system(f'pip install opencv-contrib-python==4.5.5.62 --no-cache-dir')
	# os.system('pip list')
	# os.system(f'pip install pycocotools==2.0.0 mmdet mmcv-full==1.5.0 -f https://download.openmmlab.com/mmcv/dist/{cuda_ver}/torch1.10.0/index.html --no-cache-dir')
	# os.system('wget -nv -c https://download.openmmlab.com/mmocr/data/wildreceipt.tar; mkdir -p data; tar -xf wildreceipt.tar --directory data; rm -f wildreceipt.tar')

	import datetime
	import gradio as gr
	import pandas as pd
	from mmocr.utils.ocr import MMOCR
	import os

	def inference(img, det, recog):
	print(datetime.datetime.now(), 'start')
	# ocr = MMOCR(det='PS_CTW', recog='SAR', kie='SDMGR')
	# ocr = MMOCR(det=det, recog=recog, kie='SDMGR')
	if det == 'None':
	det = None
	if recog == 'None':
	recog = None
	ocr = MMOCR(det=det, recog=recog)
	print(datetime.datetime.now(), 'start read:', img.name)
	results = ocr.readtext(img.name, details=True, output='/tmp')
	result_file = '/tmp/out_{}.png'.format(os.path.splitext(os.path.basename(img.name))[0])
	print(datetime.datetime.now(), results)
	# return result_file, pd.DataFrame(results[0]['result']).iloc[: , 2:]
	return result_file, results

	description = 'Gradio demo for MMOCR. MMOCR is an open-source toolbox based on PyTorch and mmdetection for text detection, text recognition, and the corresponding downstream tasks including key information extraction. To use it, simply upload your image or click one of the examples to load them. Read more at the links below.'
	article = "<p style='text-align: center'><a href='https://mmocr.readthedocs.io/en/latest/'>MMOCR is an open-source toolbox based on PyTorch and mmdetection for text detection, text recognition, and the corresponding downstream tasks including key information extraction.</a> \| <a href='https://github.com/open-mmlab/mmocr'>Github Repo</a></p>"


	examples = []
	path = './images'

	files = os.listdir(path)
	files.sort()
	for f in files:
	file = os.path.join(path, f)
	if os.path.isfile(file):
	examples.append([file, 'PS_CTW', 'SAR'])

	det = gr.inputs.Dropdown(choices=[
	'DB_r18',
	'DB_r50',
	'DBPP_r50',
	'DRRG',
	'FCE_IC15',
	'FCE_CTW_DCNv2',
	'MaskRCNN_CTW',
	'MaskRCNN_IC15',
	'MaskRCNN_IC17',
	'PANet_CTW',
	'PANet_IC15',
	'PS_CTW',
	'PS_IC15',
	'TextSnake',
	'None'
	], type="value", default='PS_CTW', label='det')

	recog = gr.inputs.Dropdown(choices=[
	'CRNN',
	'SAR',
	'SAR_CN',
	'NRTR_1/16-1/8',
	'NRTR_1/8-1/4',
	'RobustScanner',
	'SATRN',
	'SATRN_sm',
	'ABINet',
	'ABINet_Vision',
	'SEG',
	'CRNN_TPS',
	'MASTER',
	'None'
	], type="value", default='SAR', label='recog')

	gr.Interface(inference,
	[gr.inputs.Image(type='file', label='Input'), det, recog ],
	# [gr.outputs.Image(type='pil', label='Output'), gr.outputs.Dataframe(headers=['text', 'text_score', 'label', 'label_score'])],
	[gr.outputs.Image(type='pil', label='Output'), gr.outputs.Textbox(type='str', label='Prediction')],
	title='MMOCR',
	description=description,
	article=article,
	examples=examples,
	css=".output_image, .input_image {height: 40rem !important; width: 100% !important;}",
	enable_queue=True
	).launch(debug=True)