|
|
import atexit
|
|
|
import functools
|
|
|
from queue import Queue
|
|
|
from threading import Event, Thread
|
|
|
|
|
|
from paddleocr import PaddleOCR, draw_ocr
|
|
|
from PIL import Image
|
|
|
import gradio as gr
|
|
|
|
|
|
|
|
|
LANG_CONFIG = {
|
|
|
"ch": {"num_workers": 2},
|
|
|
"en": {"num_workers": 2},
|
|
|
"fr": {"num_workers": 1},
|
|
|
"german": {"num_workers": 1},
|
|
|
"korean": {"num_workers": 1},
|
|
|
"japan": {"num_workers": 1},
|
|
|
}
|
|
|
CONCURRENCY_LIMIT = 8
|
|
|
|
|
|
|
|
|
class PaddleOCRModelManager(object):
|
|
|
def __init__(self,
|
|
|
num_workers,
|
|
|
model_factory):
|
|
|
super().__init__()
|
|
|
self._model_factory = model_factory
|
|
|
self._queue = Queue()
|
|
|
self._workers = []
|
|
|
self._model_initialized_event = Event()
|
|
|
for _ in range(num_workers):
|
|
|
worker = Thread(target=self._worker, daemon=False)
|
|
|
worker.start()
|
|
|
self._model_initialized_event.wait()
|
|
|
self._model_initialized_event.clear()
|
|
|
self._workers.append(worker)
|
|
|
|
|
|
def infer(self, *args, **kwargs):
|
|
|
|
|
|
result_queue = Queue(maxsize=1)
|
|
|
self._queue.put((args, kwargs, result_queue))
|
|
|
success, payload = result_queue.get()
|
|
|
if success:
|
|
|
return payload
|
|
|
else:
|
|
|
raise payload
|
|
|
|
|
|
def close(self):
|
|
|
for _ in self._workers:
|
|
|
self._queue.put(None)
|
|
|
for worker in self._workers:
|
|
|
worker.join()
|
|
|
|
|
|
def _worker(self):
|
|
|
model = self._model_factory()
|
|
|
self._model_initialized_event.set()
|
|
|
while True:
|
|
|
item = self._queue.get()
|
|
|
if item is None:
|
|
|
break
|
|
|
args, kwargs, result_queue = item
|
|
|
try:
|
|
|
result = model.ocr(*args, **kwargs)
|
|
|
result_queue.put((True, result))
|
|
|
except Exception as e:
|
|
|
result_queue.put((False, e))
|
|
|
finally:
|
|
|
self._queue.task_done()
|
|
|
|
|
|
|
|
|
def create_model(lang):
|
|
|
return PaddleOCR(lang=lang, use_angle_cls=True, use_gpu=False)
|
|
|
|
|
|
|
|
|
model_managers = {}
|
|
|
for lang, config in LANG_CONFIG.items():
|
|
|
model_manager = PaddleOCRModelManager(config["num_workers"], functools.partial(create_model, lang=lang))
|
|
|
model_managers[lang] = model_manager
|
|
|
|
|
|
|
|
|
def close_model_managers():
|
|
|
for manager in model_managers.values():
|
|
|
manager.close()
|
|
|
|
|
|
|
|
|
|
|
|
atexit.register(close_model_managers)
|
|
|
|
|
|
|
|
|
def inference(img, lang):
|
|
|
ocr = model_managers[lang]
|
|
|
result = ocr.infer(img, cls=True)[0]
|
|
|
img_path = img
|
|
|
image = Image.open(img_path).convert("RGB")
|
|
|
boxes = [line[0] for line in result]
|
|
|
txts = [line[1][0] for line in result]
|
|
|
scores = [line[1][1] for line in result]
|
|
|
im_show = draw_ocr(image, boxes, txts, scores,
|
|
|
font_path="./simfang.ttf")
|
|
|
return im_show
|
|
|
|
|
|
|
|
|
title = 'PaddleOCR'
|
|
|
description = '''
|
|
|
- Gradio demo for PaddleOCR. PaddleOCR demo supports Chinese, English, French, German, Korean and Japanese.
|
|
|
- To use it, simply upload your image and choose a language from the dropdown menu, or click one of the examples to load them. Read more at the links below.
|
|
|
- [Docs](https://paddlepaddle.github.io/PaddleOCR/), [Github Repository](https://github.com/PaddlePaddle/PaddleOCR).
|
|
|
'''
|
|
|
|
|
|
examples = [
|
|
|
['en_example.jpg','en'],
|
|
|
['cn_example.jpg','ch'],
|
|
|
['jp_example.jpg','japan'],
|
|
|
]
|
|
|
|
|
|
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
|
|
|
gr.Interface(
|
|
|
inference,
|
|
|
[
|
|
|
gr.Image(type='filepath', label='Input'),
|
|
|
gr.Dropdown(choices=list(LANG_CONFIG.keys()), value='en', label='language')
|
|
|
],
|
|
|
gr.Image(type='pil', label='Output'),
|
|
|
title=title,
|
|
|
description=description,
|
|
|
examples=examples,
|
|
|
cache_examples=False,
|
|
|
css=css,
|
|
|
concurrency_limit=CONCURRENCY_LIMIT,
|
|
|
).launch(debug=False)
|
|
|
|
