| 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": 1}, |
| "en": {"num_workers": 1}, |
| "fr": {"num_workers": 1}, |
| "german": {"num_workers": 1}, |
| "korean": {"num_workers": 1}, |
| "japan": {"num_workers": 1}, |
| } |
|
|
| CONCURRENCY_LIMIT = 2 |
|
|
|
|
| class PaddleOCRModelManager: |
| def __init__(self, num_workers, model_factory): |
| 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=True) |
| 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) |
|
|
| 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)) |
|
|
|
|
| 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_managers[lang] = PaddleOCRModelManager( |
| config["num_workers"], |
| functools.partial(create_model, lang=lang) |
| ) |
|
|
|
|
| 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) |
|
|
| if not result or not result[0]: |
| return Image.open(img) |
|
|
| result = result[0] |
|
|
| image = Image.open(img).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. |
| |
| Supported languages: |
| Chinese, English, French, German, Korean, Japanese. |
| |
| Upload an image and select the language. |
| """ |
|
|
| examples = [ |
| ["en_example.jpg", "en"], |
| ["cn_example.jpg", "ch"], |
| ["jp_example.jpg", "japan"], |
| ] |
|
|
| css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}" |
|
|
|
|
| demo = gr.Interface( |
| fn=inference, |
| inputs=[ |
| gr.Image(type="filepath", label="Input"), |
| gr.Dropdown( |
| choices=list(LANG_CONFIG.keys()), |
| value="en", |
| label="Language" |
| ), |
| ], |
| outputs=gr.Image(type="pil", label="Output"), |
| title=title, |
| description=description, |
| examples=examples, |
| cache_examples=False, |
| css=css, |
| concurrency_limit=CONCURRENCY_LIMIT, |
| ) |
|
|
| demo.launch() |
