|
import gradio as gr |
|
import numpy as np |
|
import cv2 |
|
import onnxruntime as ort |
|
import os |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
def download_model_files(): |
|
model_repo = "SWHL/RapidOCR" |
|
det_model_filename = "PP-OCRv4/en_PP-OCRv3_det_infer.onnx" |
|
rec_model_filename = "PP-OCRv4/ch_PP-OCRv4_rec_server_infer.onnx" |
|
cls_model_filename = "PP-OCRv3/ch_ppocr_mobile_v2.0_cls_train.onnx" |
|
|
|
det_model_path = hf_hub_download(repo_id=model_repo, filename=det_model_filename) |
|
rec_model_path = hf_hub_download(repo_id=model_repo, filename=rec_model_filename) |
|
cls_model_path = hf_hub_download(repo_id=model_repo, filename=cls_model_filename) |
|
|
|
return det_model_path, rec_model_path, cls_model_path |
|
|
|
|
|
det_model_path, rec_model_path, cls_model_path = download_model_files() |
|
|
|
|
|
det_session = ort.InferenceSession(det_model_path) |
|
rec_session = ort.InferenceSession(rec_model_path) |
|
cls_session = ort.InferenceSession(cls_model_path) |
|
|
|
def preprocess_image(image): |
|
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) |
|
|
|
return gray |
|
|
|
import cv2 |
|
from rapidocr_onnxruntime import RapidOCR |
|
|
|
|
|
ocr_engine = RapidOCR() |
|
|
|
def ocr_predict(image): |
|
""" |
|
Perform OCR on the input image and return the extracted text. |
|
""" |
|
|
|
image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
|
|
|
|
|
result, _ = ocr_engine(image_bgr) |
|
|
|
|
|
extracted_text = "\n".join([item[1] for item in result]) |
|
|
|
return extracted_text |
|
|
|
|
|
iface = gr.Interface( |
|
fn=ocr_predict, |
|
inputs=gr.Image(type="numpy", label="Upload Image"), |
|
outputs=gr.Textbox(label="Detected Text"), |
|
title="RapidOCR Gradio Demo", |
|
description="Upload an image to extract text using RapidOCR." |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |