# 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()