app.py

# coding: utf-8
# [Pix2Text](https://github.com/breezedeus/pix2text): an Open-Source Alternative to Mathpix.
# Copyright (C) 2022-2024, [Breezedeus](https://www.breezedeus.com).

import os
import json
import functools
import random
import shutil
import string
import tempfile
import time
import zipfile
from pathlib import Path

import yaml

import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download

# from cnstd.utils import pil_to_numpy, imsave

from pix2text import Pix2Text
from pix2text.utils import set_logger, merge_line_texts

logger = set_logger()

LANGUAGES = yaml.safe_load(open('languages.yaml', 'r', encoding='utf-8'))['languages']
OUTPUT_RESULT_DIR = Path('./output-results')
OUTPUT_RESULT_DIR.mkdir(exist_ok=True)


def prepare_mfd_model():
    target_fp = './yolov7-model/mfd-yolov7-epoch224-20230613.pt'
    if os.path.exists(target_fp):
        return target_fp
    HF_TOKEN = os.environ.get('HF_TOKEN')
    local_path = hf_hub_download(
        repo_id='breezedeus/paid-models',
        subfolder='cnstd/1.2',
        filename='yolov7-model-20230613.zip',
        repo_type="model",
        cache_dir='./',
        token=HF_TOKEN,
    )
    with zipfile.ZipFile(local_path) as zf:
        zf.extractall('./')
    return target_fp


def get_p2t_model(lan_list: list, mfd_model_name: str, mfr_model_name: str):
    mfd_config = {}
    if 'yolov7_tiny' not in mfd_model_name:
        mfd_fp = prepare_mfd_model()
        mfd_config = dict(  # 声明 LayoutAnalyzer 的初始化参数
            model_type='yolov7',  # 表示使用的是 YoloV7 模型，而不是 YoloV7_Tiny 模型
            model_fp=mfd_fp,  # 注：修改成你的模型文件所存储的路径
        )
    formula_config = {}
    if 'mfr-pro' in mfr_model_name:
        formula_config = dict(  # 声明 LayoutAnalyzer 的初始化参数
            model_name='mfr-pro', model_backend='onnx',
        )
    text_formula_config = dict(
        languages=lan_list, mfd=mfd_config, formula=formula_config,
    )
    total_config = {
        'layout': {'scores_thresh': 0.45},
        'text_formula': text_formula_config,
    }
    p2t = Pix2Text.from_config(total_configs=total_config,)
    return p2t


def latex_render(latex_str):
    return f"$$\n{latex_str}\n$$"
    # return latex_str


def recognize(
    lang_list, mfd_model_name, mfr_model_name, rec_type, resized_shape, image_file
):
    lang_list = [LANGUAGES[l] for l in lang_list]
    p2t = get_p2t_model(lang_list, mfd_model_name, mfr_model_name)

    # 如果 OUTPUT_RESULT_DIR 文件数量超过 100，按时间删除最早的 100 个文件
    if len(os.listdir(OUTPUT_RESULT_DIR)) > 100:
        shutil.rmtree(OUTPUT_RESULT_DIR)
        OUTPUT_RESULT_DIR.mkdir(exist_ok=True)

    out_det_fp = './docs/no-det-res.jpg'
    kwargs = dict(
        resized_shape=resized_shape,
        return_text = True,
        auto_line_break = True,
    )
    if rec_type == 'page':
        suffix = list(string.ascii_letters)
        random.shuffle(suffix)
        suffix = ''.join(suffix[:6])
        fp_suffix = f'{time.time()}-{suffix}'
        out_debug_dir = f'out-debug-{fp_suffix}'
        output_dir = OUTPUT_RESULT_DIR / f'output-{fp_suffix}'
        kwargs['save_debug_res'] = OUTPUT_RESULT_DIR / out_debug_dir
    elif rec_type == 'text_formula':
        suffix = list(string.ascii_letters)
        random.shuffle(suffix)
        suffix = ''.join(suffix[:6])
        out_det_fp = f'out-det-{time.time()}-{suffix}.jpg'
        kwargs['save_analysis_res'] = str(OUTPUT_RESULT_DIR / out_det_fp)

    out = p2t.recognize(image_file, file_type=rec_type, **kwargs)
    out_text = out
    if rec_type == 'page':
        out_text = out.to_markdown(output_dir)
        out_det_fp =kwargs['save_debug_res'] / 'layout_res.jpg'
    elif rec_type == 'text_formula':
        out_det_fp = kwargs['save_analysis_res']

    return out_text, out_det_fp


def example_func(lang_list, rec_type, image_file):
    return recognize(
        lang_list,
        mfd_model_name='yolov7 (paid)',
        mfr_model_name='mfr-pro (paid)',
        rec_type=rec_type,
        resized_shape=768,
        image_file=image_file,
    )


def main():
    langs = list(LANGUAGES.keys())
    langs.sort(key=lambda x: x.lower())

    title = ': a Free Alternative to Mathpix'
    examples = [
        [['English'], 'page', 'docs/examples/page.png',],
        [['English'], 'text_formula', 'docs/examples/mixed-en.jpg',],
        [['English', 'Chinese Simplified'], 'text_formula', 'docs/examples/mixed-ch_sim.jpg',],
        [
            ['English', 'Chinese Traditional'],
            'text_formula',
            'docs/examples/mixed-ch_tra.jpg',
        ],
        [['English', 'Vietnamese'], 'text_formula', 'docs/examples/mixed-vietnamese.jpg',],
        [['English'], 'formula', 'docs/examples/formula1.png'],
        [['English'], 'formula', 'docs/examples/formula2.jpg'],
        [['English'], 'formula', 'docs/examples/hw-formula.png'],
        [['English', 'Chinese Simplified'], 'text', 'docs/examples/pure-text.jpg',],
    ]

    table_desc = """
<div align="center">
<img src="https://pix2text.readthedocs.io/zh/latest/figs/p2t-logo.png" width="120px"/>

[![Visitors](https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fbreezedeus%2Fpix2text-demo&labelColor=%23697689&countColor=%23f5c791&style=flat&labelStyle=upper)](https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fbreezedeus%2Fpix2text-demo)

[![Discord](https://img.shields.io/discord/1200765964434821260?logo=discord&label=Discord)](https://discord.gg/GgD87WM8Tf)

|                                 |                                         |
| ------------------------------- | --------------------------------------- |
| 🏄 **Online Service**             | [p2t.breezedeus.com](https://p2t.breezedeus.com) |
| 📖 **Doc**              | [Online Doc](https://pix2text.readthedocs.io) |
| 📀 **Code**              | [Github](https://github.com/breezedeus/pix2text) |
| 🤗 **Models**              | [breezedeus/Models](https://huggingface.co/breezedeus) |
| 📄 **More Infos**              | [Pix2Text Infos](https://www.breezedeus.com/article/pix2text) |

If useful, please help to **star 🌟 [Pix2Text](https://github.com/breezedeus/pix2text)** 🙏
</div>
    """

    with gr.Blocks() as demo:
        gr.HTML(
            f'<h1 style="text-align: center; margin-bottom: 1rem;"><a href="https://github.com/breezedeus/pix2text" target="_blank">Pix2Text V1.1</a>{title}</h1>'
        )
        with gr.Row(equal_height=False):
            with gr.Column(min_width=200, variant='panel', scale=3):
                gr.Markdown('### Settings')
                lang_list = gr.Dropdown(
                    label='Text Languages',
                    choices=langs,
                    value=['English', 'Chinese Simplified'],
                    multiselect=True,
                    # info='Which languages to be recognized as Texts.',
                )
                mfd_model_name = gr.Dropdown(
                    label='MFD Models',
                    choices=['yolov7_tiny (free)', 'yolov7 (paid)'],
                    value='yolov7 (paid)',
                )
                mfr_model_name = gr.Dropdown(
                    label='MFR Models',
                    choices=['mfr (free)', 'mfr-pro (paid)'],
                    value='mfr-pro (paid)',
                )
                rec_type = gr.Dropdown(
                    label='File Type',
                    choices=['page', 'text_formula', 'formula', 'text'],
                    value='text_formula',
                    # info='Which type of image to be recognized.',
                )
                with gr.Accordion('More Options', open=False):
                    resized_shape = gr.Slider(
                        label='resized_shape',
                        minimum=512,
                        maximum=2048,
                        value=768,
                        step=32,
                    )

            with gr.Column(scale=6, variant='compact'):
                gr.Markdown('### Upload Image to be Recognized')
                image_file = gr.Image(
                    label='Image', type="pil", image_mode='RGB', show_label=False
                )
                sub_btn = gr.Button("Submit", variant="primary")

            with gr.Column(scale=2, variant='compact'):
                gr.Markdown(table_desc)
        with gr.Row(equal_height=False):
            with gr.Column(scale=1, variant='compact'):
                gr.Markdown('**Detection Result**')
                det_result = gr.Image(
                    label='Detection Result', scale=1, show_label=False
                )
            with gr.Column(scale=1, variant='compact'):
                gr.Markdown(
                    '**Recognition Results (Paste them into the [P2T Online Service](https://p2t.breezedeus.com) to view rendered outcomes)**'
                )
                rec_result = gr.Textbox(
                    label=f'Recognition Result ',
                    lines=5,
                    value='',
                    scale=1,
                    show_label=False,
                    show_copy_button=True,
                )
            # render_result = gr.Markdown(label=f'After Rendering', value='')
            # rec_result.change(latex_render, rec_result, render_result)
        sub_btn.click(
            recognize,
            inputs=[
                lang_list,
                mfd_model_name,
                mfr_model_name,
                rec_type,
                resized_shape,
                image_file,
            ],
            outputs=[rec_result, det_result],
        )

        gr.Examples(
            label='Examples',
            examples=examples,
            inputs=[lang_list, rec_type, image_file,],
            outputs=[rec_result, det_result],
            fn=example_func,
            cache_examples=os.getenv('CACHE_EXAMPLES') == '1',
        )

    demo.queue(max_size=10)
    demo.launch()


if __name__ == '__main__':
    main()