Update files

app.py

@@ -84,7 +84,24 @@ def visualize(image_path, boxes, txts, scores,
     return image_save
 
 
-def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5):
+def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5,
+              text_det=None, text_rec=None):
+    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]
+
+    print('Init Class')
+    s = time.time()
+    rapid_ocr = RapidOCR(det_model_path=det_model_path,
+                         rec_model_path=rec_model_path,
+                         rec_img_shape=rec_image_shape)
+    print(det_model_path, rec_model_path, rec_image_shape)
+    elapse = time.time() - s
+    print(elapse)
+
     img = cv2.imread(img_path)
     ocr_result, _ = rapid_ocr(img, box_thresh=box_thresh,
                               unclip_ratio=unclip_ratio,
@@ -99,8 +116,7 @@ def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5):
     return img_save_path, output_text
 
 
-rapid_ocr = RapidOCR()
-examples = [['images/1.jpg']]
+examples = [['images/1.jpg'], ['images/ch_en_num.jpg']]
 
 with gr.Blocks(title='RapidOCR') as demo:
     gr.Markdown("""
@@ -112,10 +128,8 @@ with gr.Blocks(title='RapidOCR') as demo:
         - **unclip_ratio**: 控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]
         - **text_score**: 文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]
         ### 运行环境：
-        - Python: 3.8
-        - onnxruntime: 1.14.1
-        - rapidocr_onnxruntime: 1.2.5""")
-    gr.Markdown('### 超参数调节')
+        Python: 3.8 | onnxruntime: 1.14.1 | rapidocr_onnxruntime: 1.2.5""")
+    gr.Markdown('**超参数调节**')
     with gr.Row():
         box_thresh = gr.Slider(minimum=0, maximum=1.0, value=0.5,
                                label='box_thresh', step=0.1,
@@ -130,7 +144,7 @@ with gr.Blocks(title='RapidOCR') as demo:
                                interactive=True,
                                info='[0, 1.0]')
 
-    gr.Markdown('### 模型选择')
+    gr.Markdown('**模型选择**')
     with gr.Row():
         text_det = gr.Dropdown(['ch_PP-OCRv3_det_infer.onnx',
                                 'ch_PP-OCRv2_det_infer.onnx',
@@ -138,6 +152,7 @@ with gr.Blocks(title='RapidOCR') as demo:
                                label='选择文本检测模型',
                                value='ch_PP-OCRv3_det_infer.onnx',
                                interactive=True)
+
         text_rec = gr.Dropdown(['ch_PP-OCRv3_rec_infer.onnx',
                                 'ch_PP-OCRv2_rec_infer.onnx',
                                 'ch_ppocr_server_v2.0_rec_infer.onnx'],
@@ -151,9 +166,11 @@ with gr.Blocks(title='RapidOCR') as demo:
     out_txt = gr.outputs.Textbox(type='text', label='RecText')
     button = gr.Button('Submit')
     button.click(fn=inference,
-                 inputs=[input_img, box_thresh, unclip_ratio, text_score],
+                 inputs=[input_img, box_thresh, unclip_ratio, text_score,
+                         text_det, text_rec],
                  outputs=[out_img, out_txt])
     gr.Examples(examples=examples,
-                inputs=[input_img, box_thresh, unclip_ratio, text_score],
+                inputs=[input_img, box_thresh, unclip_ratio, text_score,
+                        text_det, text_rec],
                 outputs=[out_img, out_txt], fn=inference)
 demo.launch(debug=True, enable_queue=True)