app.py

# coding: utf-8
# Copyright (C) 2023, [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.
# Ref: https://huggingface.co/spaces/hysts/Manga-OCR/blob/main/app.py

import os
import json
import functools

import gradio as gr
import cv2
import numpy as np
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()
MODELS = {}


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))
    if len(crops) < 1:
        return
    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)


def get_ocr_model(det_model_name, rec_model_name, det_more_configs):
    global MODELS
    config_str = json.dumps(det_more_configs)
    if (det_model_name, rec_model_name, config_str) in MODELS:
        return MODELS[(det_model_name, rec_model_name, config_str)]

    det_model_name, det_model_backend = det_model_name.split('::')
    rec_model_name, rec_model_backend = rec_model_name.split('::')
    model = 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,
    )
    if len(MODELS) > 50:
        MODELS = {}
    MODELS[(det_model_name, rec_model_name, config_str)] = model
    return model


def visualize_naive_result(img, det_model_name, std_out, box_score_thresh):
    if len(std_out) < 1:
        # gr.Warning(f'未检测到文本！')
        return []
    img = pil_to_numpy(img).transpose((1, 2, 0)).astype(np.uint8)

    # plot_for_debugging(img, std_out, box_score_thresh, 2, './streamlit-app')
    # gr.Markdown('## Detection Result')
    # if det_model_name == 'naive_det':
    #     gr.Warning('⚠️ Warning: "naive_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')

    return _visualize_ocr(std_out)


def _visualize_ocr(ocr_outs):
    if len(ocr_outs) < 1:
        return
    ocr_res = []
    for out in ocr_outs:
        # cropped_img = out['cropped_img']  # 检测出的文本框
        ocr_res.append([out['score'], out['text']])
    return 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)
    return out_draw_fp


def recognize(
    det_model_name,
    is_single_line,
    rec_model_name,
    rotated_bbox,
    use_angle_clf,
    new_size,
    box_score_thresh,
    min_box_size,
    image_file,
):
    img = image_file.convert('RGB')
    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)

    if is_single_line:
        ocr_out = [ocr.ocr_for_single_line(np.array(img))]
    else:
        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,
        )

    det_model_name, det_model_backend = det_model_name.split('::')
    if is_single_line or det_model_name == 'naive_det':
        out_texts = visualize_naive_result(
            img, det_model_name, ocr_out, box_score_thresh
        )
        if is_single_line:
            return [
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(value=out_texts, visible=True),
            ]
        return [
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(value=out_texts, visible=True),
        ]
    else:
        out_img_path = visualize_result(img, ocr_out)
        return [
            gr.update(value=out_img_path, visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
        ]


def main():
    det_models = list(DET_AVAILABLE_MODELS.all_models())
    det_models.append(('naive_det', 'onnx'))
    det_models.sort()
    det_models = [f'{m}::{b}' for m, b in det_models]

    all_models = list(REC_AVAILABLE_MODELS.all_models())
    all_models.sort()
    all_models = [f'{m}::{b}' for m, b in all_models]

    title = '开源Python OCR工具：'
    desc = (
        '<p style="text-align: center">详细说明参见：<a href="https://github.com/breezedeus/CnOCR" target="_blank">Github</a>；'
        '<a href="https://cnocr.readthedocs.io" target="_blank">在线文档</a>；'
        '欢迎加入 <a href="https://www.breezedeus.com/join-group" target="_blank">交流群</a>；'
        '作者：<a href="https://www.breezedeus.com" target="_blank">Breezedeus</a> ，'
        '<a href="https://github.com/breezedeus" target="_blank">Github</a> 。</p>'
    )
    example_func = functools.partial(
        recognize,
        # det_model_name='ch_PP-OCRv3_det::onnx',
        rotated_bbox=True,
        # use_angle_clf=False,
        new_size=768,
        box_score_thresh=0.3,
        min_box_size=10,
    )
    examples = [
        [
            'ch_PP-OCRv3_det::onnx',
            True,
            'number-densenet_lite_136-fc::onnx',
            False,
            'docs/examples/card1-s.jpg',
        ],
        [
            'ch_PP-OCRv3_det::onnx',
            True,
            'number-densenet_lite_136-fc::onnx',
            False,
            'docs/examples/card2-s.jpg',
        ],
        [
            'ch_PP-OCRv3_det::onnx',
            True,
            'number-densenet_lite_136-fc::onnx',
            False,
            'docs/examples/cy1-s.jpg',
        ],
        [
            'ch_PP-OCRv3_det::onnx',
            False,
            'densenet_lite_136-gru::onnx',
            False,
            'docs/examples/huochepiao.jpeg',
        ],
        [
            'ch_PP-OCRv3_det::onnx',
            False,
            'densenet_lite_136-gru::onnx',
            False,
            'docs/examples/1_res.jpg',
        ],
        [
            'db_shufflenet_v2::pytorch',
            False,
            'en_number_mobile_v2.0::onnx',
            False,
            'docs/examples/en_book1.jpeg',
        ],
        [
            'db_shufflenet_v2::pytorch',
            False,
            'densenet_lite_136-gru::onnx',
            True,
            'docs/examples/beauty0.jpg',
        ],
    ]

    with gr.Blocks() as demo:
        gr.Markdown(
            f'<h1 style="text-align: center; margin-bottom: 1rem;">{title} <a href="https://github.com/breezedeus/cnocr" target="_blank">CnOCR</a></h1>'
        )
        gr.Markdown(desc)
        with gr.Row(equal_height=False):
            with gr.Column(min_width=200, variant='panel', scale=1):
                gr.Markdown('### 模型设置')
                det_model_name = gr.Dropdown(
                    label='选择检测模型', choices=det_models, value='ch_PP-OCRv3_det::onnx',
                )
                is_single_line = gr.Checkbox(label='单行文字模式（不使用检测模型）', value=False)

                rec_model_name = gr.Dropdown(
                    label='选择识别模型',
                    choices=all_models,
                    value='densenet_lite_136-fc::onnx',
                )

                gr.Markdown('### 检测参数')
                rotated_bbox = gr.Checkbox(label='检测带角度文本框', value=True)
                use_angle_clf = gr.Checkbox(label='使用角度预测模型校正文本框', value=False)
                new_size = gr.Slider(
                    label='resize 后图片（长边）大小', minimum=124, maximum=4096, value=768
                )
                box_score_thresh = gr.Slider(
                    label='得分阈值（低于阈值的结果会被过滤掉）', minimum=0.05, maximum=0.95, value=0.3
                )
                min_box_size = gr.Slider(
                    label='框大小阈值（更小的文本框会被过滤掉）', minimum=4, maximum=50, value=10
                )

            with gr.Column(scale=3, variant='compact'):
                gr.Markdown('### 选择待识别图片')
                image_file = gr.Image(label='待识别图片', type="pil", image_mode='RGB')
                sub_btn = gr.Button("Submit", variant="primary")
                out_image = gr.Image(label='识别结果', interactive=False, visible=False)
                naive_warn = gr.Markdown(
                    '**⚠️ Warning**: "naive_det" 检测模型不返回文本框位置！', visible=False
                )
                out_texts = gr.Dataframe(
                    headers=['得分', '文本'], label='识别结果', interactive=False, visible=False
                )
                sub_btn.click(
                    recognize,
                    inputs=[
                        det_model_name,
                        is_single_line,
                        rec_model_name,
                        rotated_bbox,
                        use_angle_clf,
                        new_size,
                        box_score_thresh,
                        min_box_size,
                        image_file,
                    ],
                    outputs=[out_image, naive_warn, out_texts],
                )

        gr.Examples(
            label='示例',
            examples=examples,
            inputs=[
                det_model_name,
                is_single_line,
                rec_model_name,
                use_angle_clf,
                image_file,
            ],
            outputs=[out_image, naive_warn, out_texts],
            fn=example_func,
            cache_examples=os.getenv('CACHE_EXAMPLES') == '1',
        )

    demo.queue(concurrency_count=4)
    demo.launch()


if __name__ == '__main__':
    main()