Update files

.gitattributes

@@ -34,3 +34,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.TTF filter=lfs diff=lfs merge=lfs -text
 images/car_plate.jpeg filter=lfs diff=lfs merge=lfs -text
 *.ttc filter=lfs diff=lfs merge=lfs -text
+fonts/simfang.ttf filter=lfs diff=lfs merge=lfs -text
+models/text_det/ch_PP-OCRv4_det_infer.onnx filter=lfs diff=lfs merge=lfs -text
+models/text_rec/ch_PP-OCRv4_rec_infer.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,3 +1,5 @@
 *.pyc
 
 __pycache__/
+.vscode
+.DS_Store

README.md

@@ -3,8 +3,8 @@ title: RapidOCR
 emoji: ⚡
 colorFrom: blue
 colorTo: blue
-sdk: gradio
-sdk_version: 3.6
+sdk: streamlit
+sdk_version: 1.25.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -1,138 +1,202 @@
 # -*- encoding: utf-8 -*-
-import json
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
 import time
 from pathlib import Path
 
 import cv2
-import gradio as gr
+import numpy as np
+import pandas as pd
+import streamlit as st
+from PIL import Image
 from rapidocr_onnxruntime import RapidOCR
+from streamlit_image_select import image_select
 
 from utils import visualize
 
 font_dict = {
-    'ch': 'FZYTK.TTF',
-    'japan': 'japan.ttc',
-    'korean': 'korean.ttf',
-    'en': 'FZYTK.TTF',
+    "ch": "chinese_cht.ttf",
+    "japan": "japan.ttc",
+    "korean": "korean.ttf",
+    "en": "chinese_cht.ttf",
 }
 
 
-def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5,
-              text_det=None, text_rec=None):
-    out_log_list = []
-
-    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:
+def init_sidebar():
+    st.session_state["params"] = {}
+
+    st.sidebar.markdown(
+        "### [🛠️ Parameter Settings](https://github.com/RapidAI/RapidOCR/wiki/config_parameter)"
+    )
+    box_thresh = st.sidebar.slider(
+        "box_thresh",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.5,
+        step=0.1,
+        help="检测到的框是文本的概率，值越大，框中是文本的概率就越大。存在漏检时，调低该值。取值范围：[0, 1.0]，默认值为0.5",
+    )
+    st.session_state["params"]["box_thresh"] = box_thresh
+
+    unclip_ratio = st.sidebar.slider(
+        "unclip_ratio",
+        min_value=1.5,
+        max_value=2.0,
+        value=1.6,
+        step=0.1,
+        help="控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]，默认值为1.6",
+    )
+    st.session_state["params"]["unclip_ratio"] = unclip_ratio
+
+    text_score = st.sidebar.slider(
+        "text_score",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.5,
+        step=0.1,
+        help="文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]，默认值为0.5",
+    )
+    st.session_state["params"]["text_score"] = text_score
+
+    img_file_buffer = st.sidebar.file_uploader(
+        "Upload an image",
+        accept_multiple_files=False,
+        label_visibility="visible",
+        type=["png", "jpg", "jpeg", "bmp"],
+    )
+    if img_file_buffer:
+        image = Image.open(img_file_buffer)
+        img = np.array(image)
+
+    with st.sidebar.container():
+        img_path = image_select(
+            label="Examples（click to select）:",
+            images=examples,
+            key="equation_default",
+            use_container_width=True,
+        )
+        img = cv2.imread(img_path)
+    st.session_state["img"] = img
+
+
+def inference(
+    text_det=None,
+    text_rec=None,
+):
+    img = st.session_state.get("img")
+    box_thresh = st.session_state["params"].get("box_thresh")
+    unclip_ratio = st.session_state["params"].get("unclip_ratio")
+    text_score = st.session_state["params"].get("text_score")
+
+    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:
         rec_image_shape = [3, 32, 320]
     else:
         rec_image_shape = [3, 48, 320]
 
-    out_log_list.append('Init Model')
-    s = time.time()
-    rapid_ocr = RapidOCR(det_model_path=det_model_path,
-                         rec_model_path=rec_model_path,
-                         rec_img_shape=rec_image_shape)
-    elapse = time.time() - s
-
-    if 'ch' in rec_model_path or 'en' in rec_model_path:
-        lan_name = 'ch'
-    elif 'japan' in rec_model_path:
-        lan_name = 'japan'
-    elif 'korean' in rec_model_path:
-        lan_name = 'korean'
+    rapid_ocr = RapidOCR(
+        det_model_path=det_model_path,
+        rec_model_path=rec_model_path,
+        rec_img_shape=rec_image_shape,
+    )
+
+    if "ch" in rec_model_path or "en" in rec_model_path:
+        lan_name = "ch"
+    elif "japan" in rec_model_path:
+        lan_name = "japan"
+    elif "korean" in rec_model_path:
+        lan_name = "korean"
     else:
-        lan_name = 'ch'
+        lan_name = "ch"
 
-    out_log_list.append(f'Init Model cost: {elapse:.5f}')
-    out_log_list.extend([f'det_model: {det_model_path}',
-                         f'rec_model: {rec_model_path}',
-                         f'rec_image_shape: {rec_image_shape}'])
+    ocr_result, infer_elapse = rapid_ocr(
+        img, box_thresh=box_thresh, unclip_ratio=unclip_ratio, text_score=text_score
+    )
+    if not ocr_result or not infer_elapse:
+        return None, None, None, None
 
-    img = cv2.imread(img_path)
-    ocr_result, infer_elapse = rapid_ocr(img, box_thresh=box_thresh,
-                                         unclip_ratio=unclip_ratio,
-                                         text_score=text_score)
     det_cost, cls_cost, rec_cost = infer_elapse
-    out_log_list.extend([f'det cost: {det_cost:.5f}',
-                         f'cls cost: {cls_cost:.5f}',
-                         f'rec cost: {rec_cost:.5f}'])
-    out_log = '\n'.join([str(v) for v in out_log_list])
-    if not ocr_result:
-        return img_path, '未识别到有效文本', out_log
-
+    elapse = f"- `det cost`: {det_cost:.5f}\n - `cls cost`: {cls_cost:.5f}\n - `rec cost`: {rec_cost:.5f}"
     dt_boxes, rec_res, scores = list(zip(*ocr_result))
-    font_path = Path('fonts') / font_dict.get(lan_name)
-    img_save_path = visualize(img_path, dt_boxes, rec_res, scores,
-                              font_path=str(font_path))
-    out_dict = {str(i): {'rec_txt': rec, 'score': score} 
-                for i, (rec, score) in enumerate(zip(rec_res, scores))}
-    return img_save_path, out_dict, out_log
-
-
-if __name__ == '__main__':
-    examples = [['images/1.jpg'],
-                ['images/ch_en_num.jpg'],
-                ['images/air_ticket.jpg'],
-                ['images/car_plate.jpeg'],
-                ['images/idcard.jpg'],
-                ['images/train_ticket.jpeg'],
-                ['images/japan_2.jpg'],
-                ['images/korean_1.jpg']]
-
-    with gr.Blocks(title='RapidOCR') as demo:
-        gr.Markdown("""
-            <h1><center><a href="https://github.com/RapidAI/RapidOCR" target="_blank">Rapid⚡OCR</a></center></h1>
-
-            ### Docs: [Docs](https://rapidocr.rtfd.io/)
-            ### 运行环境：
-            Python: 3.8 | onnxruntime: 1.14.1 | rapidocr_onnxruntime: 1.2.5""")
-        gr.Markdown(
-            '''**[超参数调节](https://github.com/RapidAI/RapidOCR/tree/main/python#configyaml%E4%B8%AD%E5%B8%B8%E7%94%A8%E5%8F%82%E6%95%B0%E4%BB%8B%E7%BB%8D)**
-            - **box_thresh**: 检测到的框是文本的概率，值越大，框中是文本的概率就越大。存在漏检时，调低该值。取值范围：[0, 1.0]
-            - **unclip_ratio**: 控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]
-            - **text_score**: 文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]
-            ''')
-        with gr.Row():
-            box_thresh = gr.Slider(minimum=0, maximum=1.0, value=0.5,
-                                   label='box_thresh', step=0.1,
-                                   interactive=True,
-                                   info='[0, 1.0]')
-            unclip_ratio = gr.Slider(minimum=1.5, maximum=2.0, value=1.6,
-                                     label='unclip_ratio', step=0.1,
-                                     interactive=True,
-                                     info='[1.5, 2.0]')
-            text_score = gr.Slider(minimum=0, maximum=1.0, value=0.5,
-                                   label='text_score', step=0.1,
-                                   interactive=True,
-                                   info='[0, 1.0]')
-
-        gr.Markdown('**[模型选择](https://github.com/RapidAI/RapidOCR/blob/main/docs/models.md)** (模型转换→[PaddleOCRModelConverter](https://github.com/RapidAI/PaddleOCRModelConverter))')
-        with gr.Row():
-            text_det = gr.Dropdown(['ch_PP-OCRv3_det_infer.onnx',
-                                    'ch_PP-OCRv2_det_infer.onnx',
-                                    'ch_ppocr_server_v2.0_det_infer.onnx'],
-                                   label='选择文本检测模型',
-                                   value='ch_PP-OCRv3_det_infer.onnx',
-                                   interactive=True)
-            rec_model_list = [v.name for v in Path('models/text_rec').iterdir()]
-            text_rec = gr.Dropdown(rec_model_list,
-                                   label='选择文本识别模型(包括中英文和多语言，欢迎提交更多模型)',
-                                   value='ch_PP-OCRv3_rec_infer.onnx',
-                                   interactive=True)
-
-        with gr.Row():
-            input_img = gr.Image(type='filepath', label='Input')
-            out_img = gr.Image(type='filepath', label='Output')
-        out_json = gr.JSON(label='Rec Res')
-        out_log = gr.outputs.Textbox(type='text', label='Run Log')
-        button = gr.Button('Submit')
-        button.click(fn=inference,
-                     inputs=[input_img, box_thresh, unclip_ratio, text_score,
-                             text_det, text_rec],
-                     outputs=[out_img, out_json, out_log])
-        gr.Examples(examples=examples,
-                    inputs=[input_img, box_thresh, unclip_ratio, text_score,
-                            text_det, text_rec],
-                    outputs=[out_img, out_json, out_log], fn=inference)
-    demo.launch(debug=True, enable_queue=True)
+    font_path = Path("fonts") / font_dict.get(lan_name)
+    vis_img = visualize(
+        Image.fromarray(img), dt_boxes, rec_res, scores, font_path=str(font_path)
+    )
+    out_df = pd.DataFrame(
+        [[rec, score] for rec, score in zip(rec_res, scores)],
+        columns=("Rec", "Score"),
+    )
+    return vis_img, out_df, elapse, rec_res
+
+
+def tips(txt: str, wait_time: int = 2, icon: str = "🎉"):
+    st.toast(txt, icon=icon)
+    time.sleep(wait_time)
+
+
+if __name__ == "__main__":
+    st.markdown(
+        "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidOCR' style='text-decoration: none'>Rapid⚡OCR</a></h1>",
+        unsafe_allow_html=True,
+    )
+    st.markdown(
+        """
+    <p align="left">
+        <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>
+        <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
+        <a href="https://pepy.tech/project/rapidocr_onnxruntime"><img src="https://static.pepy.tech/personalized-badge/rapidocr_onnxruntime?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20Ort"></a>
+        <a href="https://pypi.org/project/rapidocr-onnxruntime/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapidocr-onnxruntime"></a>
+        <a href='https://rapidocr.readthedocs.io/en/latest/?badge=latest'>
+            <img src='https://readthedocs.org/projects/rapidocr/badge/?version=latest' alt='Documentation Status' />
+    </p>
+    """,
+        unsafe_allow_html=True,
+    )
+
+    examples = [
+        "images/1.jpg",
+        "images/ch_en_num.jpg",
+        "images/air_ticket.jpg",
+        "images/car_plate.jpeg",
+        "images/train_ticket.jpeg",
+        "images/japan_2.jpg",
+        "images/korean_1.jpg",
+    ]
+
+    init_sidebar()
+
+    menu_det, menu_rec = st.columns([1, 1])
+    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",
+    ]
+    select_det = menu_det.selectbox("Det model:", det_models)
+
+    rec_models = [
+        "ch_PP-OCRv4_rec_infer.onnx",
+        "ch_PP-OCRv3_rec_infer.onnx",
+        "ch_PP-OCRv2_rec_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",
+    ]
+    select_rec = menu_rec.selectbox("Rec model:", rec_models)
+
+    out_img, out_json, elapse, only_txts = inference(select_det, select_rec)
+    if all(v is not None for v in [out_img, out_json, elapse]):
+        st.markdown("#### Visualize:")
+        st.image(out_img)
+
+        st.markdown("### Rec Result:")
+        st.markdown(elapse)
+        st.dataframe(out_json, use_container_width=True)
+
+        st.markdown("### Only Txts")
+        st.code(only_txts)
+    else:
+        tips("识别结果为空", wait_time=5, icon="⚠️")

requirements.txt

@@ -1,4 +1,3 @@
-Gradio
 Pillow
 onnxruntime==1.14.1
 rapidocr_onnxruntime==1.2.5

utils.py

@@ -3,19 +3,16 @@
 # @Contact: liekkaskono@163.com
 import math
 import random
-import time
 from pathlib import Path
 
-import cv2
 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):
+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))
+    img_right = Image.new("RGB", (w, h), (255, 255, 255))
 
     random.seed(0)
     draw_left = ImageDraw.Draw(img_left)
@@ -24,62 +21,58 @@ def draw_ocr_box_txt(image, boxes, txts, font_path,
         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))
+        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.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)
+        draw_right.text([box[3][0], box[3][1]], str(idx), fill=color)
 
-        box_height = math.sqrt((box[0][0] - box[3][0])**2
-                               + (box[0][1] - box[3][1])**2)
+        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_width = math.sqrt((box[0][0] - box[1][0])**2
-                              + (box[0][1] - box[1][1])**2)
+        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")
+            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)
+                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)
+            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 = 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_path, boxes, txts, scores,
-              font_path="./fonts/FZYTK.TTF"):
-    image = Image.open(image_path)
-
-    draw_img = draw_ocr_box_txt(image, boxes,
-                                txts, font_path,
-                                scores,
-                                text_score=0.5)
+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)
-
-    time_stamp = time.strftime(
-        '%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
-    image_save = str(draw_img_save / f'{time_stamp}_{Path(image_path).name}')
-    cv2.imwrite(image_save, draw_img[:, :, ::-1])
-    return image_save
+    return draw_img[:, :, ::-1]