Spaces:
Runtime error
Runtime error
File size: 2,505 Bytes
7f4886d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import time
from pathlib import Path
import cv2
import uvicorn
import numpy as np
import pandas as pd
from PIL import Image
from rapidocr_onnxruntime import RapidOCR
from fastapi import FastAPI, File, UploadFile, Query
from fastapi.responses import RedirectResponse
app = FastAPI()
det_models = [
"ch_PP-OCRv4_det_infer.onnx",
"ch_PP-OCRv3_det_infer.onnx",
"ch_PP-OCRv2_det_infer.onnx",
"ch_ppocr_server_v2.0_det_infer.onnx",
]
rec_models = [
"ch_PP-OCRv4_rec_infer.onnx",
"ch_PP-OCRv3_rec_infer.onnx",
"ch_PP-OCRv2_rec_infer.onnx",
"ch_PP-OCRv4_det_server_infer.onnx",
"ch_ppocr_server_v2.0_rec_infer.onnx",
"en_PP-OCRv3_rec_infer.onnx",
"en_number_mobile_v2.0_rec_infer.onnx",
"korean_mobile_v2.0_rec_infer.onnx",
"japan_rec_crnn_v2.onnx",
]
def get_text(
image,
text_det=None,
text_rec=None,
box_thresh=0.5,
unclip_ratio=1.6,
text_score=0.5,
):
det_model_path = str(Path("models") / "text_det" / text_det)
rec_model_path = str(Path("models") / "text_rec" / text_rec)
if (
"v2" in rec_model_path
or "korean" in rec_model_path
or "japan" in rec_model_path
):
rec_image_shape = [3, 32, 320]
else:
rec_image_shape = [3, 48, 320]
rapid_ocr = RapidOCR(
det_model_path=det_model_path,
rec_model_path=rec_model_path,
rec_img_shape=rec_image_shape,
)
ocr_result, infer_elapse = rapid_ocr(
image, box_thresh=box_thresh, unclip_ratio=unclip_ratio, text_score=text_score
)
if not ocr_result or not infer_elapse:
return None
det_cost, cls_cost, rec_cost = infer_elapse
dt_boxes, rec_res, scores = list(zip(*ocr_result))
out_df = pd.DataFrame(
[[rec, score] for rec, score in zip(rec_res, scores)],
columns=("Rec", "Score"),
)
return out_df
@app.get("/", include_in_schema=False)
def docs_redirect():
return RedirectResponse(url='/docs')
@app.post("/ocr")
def create_upload_file(
file: UploadFile,
text_det: str = Query("ch_PP-OCRv4_det_infer.onnx", enum=det_models),
text_rec: str = Query("en_number_mobile_v2.0_rec_infer.onnx", enum=rec_models),
box_thresh: float = 0.5,
unclip_ratio: float = 1.6,
text_score: float = 0.5,
):
resp = get_text(
file.file.read(), text_det, text_rec, box_thresh, unclip_ratio, text_score
)
return resp.to_dict("list")
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |