Initial (no) commit (v0.0.1)

app.py

@@ -0,0 +1,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)

models/text_det/ch_PP-OCRv2_det_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43d01120eab9f00e8d7541ad0f906dd9c7f28394233e47279073fab73f82da87
+size 2341867

models/text_det/ch_PP-OCRv3_det_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3439588c030faea393a54515f51e983d8e155b19a2e8aba7891934c1cf0de526
+size 2432880

models/text_det/ch_PP-OCRv4_det_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2a7720d45a54257208b1e13e36a8479894cb74155a5efe29462512d42f49da9
+size 4745517

models/text_det/ch_PP-OCRv4_det_server_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfa39a3f298f6d3fc71789834d15da36d11a6c59b489fc16ea4733728012f786
+size 113352104

models/text_det/ch_ppocr_server_v2.0_det_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a7e74057c1f4e17f38bf3e0a9f500e64e83159e4fd58775f4953f13438370ca
+size 48866380

models/text_rec/ch_PP-OCRv2_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86eec59d02644a8890d54f4aa5fe7402cf2848450fc7ac613d4677cd6025b3da
+size 8402009

models/text_rec/ch_PP-OCRv3_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:897a3ededb38fee0dae2c1ccee38241f37df202c9509e3abca02e9217c5ee615
+size 10690752

models/text_rec/ch_PP-OCRv4_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48fc40f24f6d2a207a2b1091d3437eb3cc3eb6b676dc3ef9c37384005483683b
+size 10857958

models/text_rec/ch_ppocr_server_v2.0_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d0714a2b44f5d7ac7a17e44d0e1961393447590c895f4197f37e855e3464f53
+size 111590126

models/text_rec/en_PP-OCRv3_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef7abd8bd3629ae57ea2c28b425c1bd258a871b93fd2fe7c433946ade9b5d9ea
+size 8967018

models/text_rec/en_number_mobile_v2.0_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e679ba625c544444be78292a50d9e1af9caa1569239a88bb8b864cb688b11c01
+size 1882607

models/text_rec/japan_rec_crnn_v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b0495059f5738166e606d864b04ff00093f67a807efb02cddf472839cae970c
+size 3571807

models/text_rec/korean_mobile_v2.0_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6558500138b43b46a4941957fb8c918546dae5fb0e71718536f1883acc80faf
+size 3290650

requirements.txt

@@ -0,0 +1,7 @@
+Pillow
+onnxruntime
+rapidocr_onnxruntime
+streamlit_image_select
+python-multipart
+fastapi
+uvicorn

utils.py

@@ -0,0 +1,78 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]
+import math
+import random
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+
+def draw_ocr_box_txt(image, boxes, txts, font_path, scores=None, text_score=0.5):
+    h, w = image.height, image.width
+    img_left = image.copy()
+    img_right = Image.new("RGB", (w, h), (255, 255, 255))
+
+    random.seed(0)
+    draw_left = ImageDraw.Draw(img_left)
+    draw_right = ImageDraw.Draw(img_right)
+    for idx, (box, txt) in enumerate(zip(boxes, txts)):
+        if scores is not None and float(scores[idx]) < text_score:
+            continue
+
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        box = [tuple(v) for v in box]
+        draw_left.polygon(box, fill=color)
+        draw_right.text([box[3][0], box[3][1]], str(idx), fill=color)
+
+        draw_right.polygon(
+            [
+                box[0][0],
+                box[0][1],
+                box[1][0],
+                box[1][1],
+                box[2][0],
+                box[2][1],
+                box[3][0],
+                box[3][1],
+            ],
+            outline=color,
+        )
+
+        box_height = math.sqrt(
+            (box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2
+        )
+
+        box_width = math.sqrt(
+            (box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2
+        )
+
+        if box_height > 2 * box_width:
+            font_size = max(int(box_width * 0.9), 10)
+            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+            cur_y = box[0][1]
+            for c in txt:
+                char_size = font.getsize(c)
+                draw_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
+                cur_y += char_size[1]
+        else:
+            font_size = max(int(box_height * 0.8), 10)
+            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+            draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
+
+    img_left = Image.blend(image, img_left, 0.5)
+    img_show = Image.new("RGB", (w * 2, h), (255, 255, 255))
+    img_show.paste(img_left, (0, 0, w, h))
+    img_show.paste(img_right, (w, 0, w * 2, h))
+    return np.array(img_show)
+
+
+def visualize(image, boxes, txts, scores, font_path="./fonts/FZYTK.TTF"):
+    draw_img = draw_ocr_box_txt(image, boxes, txts, font_path, scores, text_score=0.5)
+
+    draw_img_save = Path("./inference_results/")
+    if not draw_img_save.exists():
+        draw_img_save.mkdir(parents=True, exist_ok=True)
+    return draw_img[:, :, ::-1]