Upload app.py

app.py

@@ -0,0 +1,183 @@
+# coding: utf-8
+# Copyright (C) 2021, [Breezedeus](https://github.com/breezedeus).
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+from collections import OrderedDict
+
+import cv2
+import numpy as np
+from PIL import Image
+import streamlit as st
+from cnstd.utils import pil_to_numpy, imsave
+
+from cnocr import CnOcr, DET_AVAILABLE_MODELS, REC_AVAILABLE_MODELS
+from cnocr.utils import set_logger, draw_ocr_results, download
+
+
+logger = set_logger()
+st.set_page_config(layout="wide")
+
+
+def plot_for_debugging(rotated_img, one_out, box_score_thresh, crop_ncols, prefix_fp):
+    import matplotlib.pyplot as plt
+    import math
+
+    rotated_img = rotated_img.copy()
+    crops = [info['cropped_img'] for info in one_out]
+    print('%d boxes are found' % len(crops))
+    ncols = crop_ncols
+    nrows = math.ceil(len(crops) / ncols)
+    fig, ax = plt.subplots(nrows=nrows, ncols=ncols)
+    for i, axi in enumerate(ax.flat):
+        if i >= len(crops):
+            break
+        axi.imshow(crops[i])
+    crop_fp = '%s-crops.png' % prefix_fp
+    plt.savefig(crop_fp)
+    print('cropped results are save to file %s' % crop_fp)
+
+    for info in one_out:
+        box, score = info.get('position'), info['score']
+        if score < box_score_thresh:  # score < 0.5
+            continue
+        if box is not None:
+            box = box.astype(int).reshape(-1, 2)
+            cv2.polylines(rotated_img, [box], True, color=(255, 0, 0), thickness=2)
+    result_fp = '%s-result.png' % prefix_fp
+    imsave(rotated_img, result_fp, normalized=False)
+    print('boxes results are save to file %s' % result_fp)
+
+
+@st.cache(allow_output_mutation=True)
+def get_ocr_model(det_model_name, rec_model_name, det_more_configs):
+    det_model_name, det_model_backend = det_model_name
+    rec_model_name, rec_model_backend = rec_model_name
+    return CnOcr(
+        det_model_name=det_model_name,
+        det_model_backend=det_model_backend,
+        rec_model_name=rec_model_name,
+        rec_model_backend=rec_model_backend,
+        det_more_configs=det_more_configs,
+    )
+
+
+def visualize_naive_result(img, det_model_name, std_out, box_score_thresh):
+    img = pil_to_numpy(img).transpose((1, 2, 0)).astype(np.uint8)
+
+    plot_for_debugging(img, std_out, box_score_thresh, 2, './streamlit-app')
+    st.subheader('Detection Result')
+    if det_model_name == 'default_det':
+        st.warning('⚠️ Warning: "default_det" 检测模型不返回文本框位置！')
+    cols = st.columns([1, 7, 1])
+    cols[1].image('./streamlit-app-result.png')
+
+    st.subheader('Recognition Result')
+    cols = st.columns([1, 7, 1])
+    cols[1].image('./streamlit-app-crops.png')
+
+    _visualize_ocr(std_out)
+
+
+def _visualize_ocr(ocr_outs):
+    st.empty()
+    ocr_res = OrderedDict({'文本': []})
+    ocr_res['得分'] = []
+    for out in ocr_outs:
+        # cropped_img = out['cropped_img']  # 检测出的文本框
+        ocr_res['得分'].append(out['score'])
+        ocr_res['文本'].append(out['text'])
+    st.table(ocr_res)
+
+
+def visualize_result(img, ocr_outs):
+    out_draw_fp = './streamlit-app-det-result.png'
+    font_path = 'docs/fonts/simfang.ttf'
+    if not os.path.exists(font_path):
+        url = 'https://huggingface.co/datasets/breezedeus/cnocr-wx-qr-code/resolve/main/fonts/simfang.ttf'
+        os.makedirs(os.path.dirname(font_path), exist_ok=True)
+        download(url, path=font_path, overwrite=True)
+    draw_ocr_results(img, ocr_outs, out_draw_fp, font_path)
+    st.image(out_draw_fp)
+
+
+def main():
+    st.sidebar.header('模型设置')
+    det_models = list(DET_AVAILABLE_MODELS.all_models())
+    det_models.append(('naive_det', 'onnx'))
+    det_models.sort()
+    det_model_name = st.sidebar.selectbox(
+        '选择检测模型', det_models, index=det_models.index(('ch_PP-OCRv3_det', 'onnx'))
+    )
+
+    all_models = list(REC_AVAILABLE_MODELS.all_models())
+    all_models.sort()
+    idx = all_models.index(('densenet_lite_136-fc', 'onnx'))
+    rec_model_name = st.sidebar.selectbox('选择识别模型', all_models, index=idx)
+
+    st.sidebar.subheader('检测参数')
+    rotated_bbox = st.sidebar.checkbox('是否检测带角度文本框', value=True)
+    use_angle_clf = st.sidebar.checkbox('是否使用角度预测模型校正文��框', value=False)
+    new_size = st.sidebar.slider(
+        'resize 后图片（长边）大小', min_value=124, max_value=4096, value=768
+    )
+    box_score_thresh = st.sidebar.slider(
+        '得分阈值（低于阈值的结果会被过滤掉）', min_value=0.05, max_value=0.95, value=0.3
+    )
+    min_box_size = st.sidebar.slider(
+        '框大小阈值（更小的文本框会被过滤掉）', min_value=4, max_value=50, value=10
+    )
+    # std = get_std_model(det_model_name, rotated_bbox, use_angle_clf)
+
+    # st.sidebar.markdown("""---""")
+    # st.sidebar.header('CnOcr 设置')
+    det_more_configs = dict(rotated_bbox=rotated_bbox, use_angle_clf=use_angle_clf)
+    ocr = get_ocr_model(det_model_name, rec_model_name, det_more_configs)
+
+    st.markdown('# 开源Python OCR工具 ' '[CnOCR](https://github.com/breezedeus/cnocr)')
+    st.markdown('> 详细说明参见：[CnOCR 文档](https://cnocr.readthedocs.io/) ；'
+                '欢迎加入 [交流群](https://cnocr.readthedocs.io/zh/latest/contact/) ；'
+                '作者：[breezedeus](https://github.com/breezedeus) 。')
+    st.markdown('')
+    st.subheader('选择待检测图片')
+    content_file = st.file_uploader('', type=["png", "jpg", "jpeg", "webp"])
+    if content_file is None:
+        st.stop()
+
+    try:
+        img = Image.open(content_file).convert('RGB')
+
+        ocr_out = ocr.ocr(
+            img,
+            return_cropped_image=True,
+            resized_shape=new_size,
+            preserve_aspect_ratio=True,
+            box_score_thresh=box_score_thresh,
+            min_box_size=min_box_size,
+        )
+        if det_model_name[0] == 'naive_det':
+            visualize_naive_result(img, det_model_name[0], ocr_out, box_score_thresh)
+        else:
+            visualize_result(img, ocr_out)
+
+    except Exception as e:
+        st.error(e)
+
+
+if __name__ == '__main__':
+    main()