aci11/F12
0
1# coding: utf-82# [Pix2Text](https://github.com/breezedeus/pix2text): an Open-Source Alternative to Mathpix.3# Copyright (C) 2022-2024, [Breezedeus](https://www.breezedeus.com).4 5import os6import json7import functools8import random9import shutil10import string11import tempfile12import time13import zipfile14from pathlib import Path15 16import yaml17 18import gradio as gr19import numpy as np20from huggingface_hub import hf_hub_download21 22# from cnstd.utils import pil_to_numpy, imsave23 24from pix2text import Pix2Text25from pix2text.utils import set_logger, merge_line_texts26 27logger = set_logger()28 29LANGUAGES = yaml.safe_load(open('languages.yaml', 'r', encoding='utf-8'))['languages']30OUTPUT_RESULT_DIR = Path('./output-results')31OUTPUT_RESULT_DIR.mkdir(exist_ok=True)32 33 34def prepare_mfd_model():35 target_fp = './yolov7-model/mfd-yolov7-epoch224-20230613.pt'36 if os.path.exists(target_fp):37 return target_fp38 HF_TOKEN = os.environ.get('HF_TOKEN')39 local_path = hf_hub_download(40 repo_id='breezedeus/paid-models',41 subfolder='cnstd/1.2',42 filename='yolov7-model-20230613.zip',43 repo_type="model",44 cache_dir='./',45 token=HF_TOKEN,46 )47 with zipfile.ZipFile(local_path) as zf:48 zf.extractall('./')49 return target_fp50 51 52def get_p2t_model(lan_list: list, mfd_model_name: str, mfr_model_name: str):53 mfd_config = {}54 if 'yolov7_tiny' not in mfd_model_name:55 mfd_fp = prepare_mfd_model()56 mfd_config = dict( # 声明 LayoutAnalyzer 的初始化参数57 model_type='yolov7', # 表示使用的是 YoloV7 模型，而不是 YoloV7_Tiny 模型58 model_fp=mfd_fp, # 注：修改成你的模型文件所存储的路径59 )60 formula_config = {}61 if 'mfr-pro' in mfr_model_name:62 formula_config = dict( # 声明 LayoutAnalyzer 的初始化参数63 model_name='mfr-pro', model_backend='onnx',64 )65 text_formula_config = dict(66 languages=lan_list, mfd=mfd_config, formula=formula_config,67 )68 total_config = {69 'layout': {'scores_thresh': 0.45},70 'text_formula': text_formula_config,71 }72 p2t = Pix2Text.from_config(total_configs=total_config,)73 return p2t74 75 76def latex_render(latex_str):77 return f"$$\n{latex_str}\n$$"78 # return latex_str79 80 81def recognize(82 lang_list, mfd_model_name, mfr_model_name, rec_type, resized_shape, image_file83):84 lang_list = [LANGUAGES[l] for l in lang_list]85 p2t = get_p2t_model(lang_list, mfd_model_name, mfr_model_name)86 87 # 如果 OUTPUT_RESULT_DIR 文件数量超过 100，按时间删除最早的 100 个文件88 if len(os.listdir(OUTPUT_RESULT_DIR)) > 100:89 shutil.rmtree(OUTPUT_RESULT_DIR)90 OUTPUT_RESULT_DIR.mkdir(exist_ok=True)91 92 out_det_fp = './docs/no-det-res.jpg'93 kwargs = dict(94 resized_shape=resized_shape,95 return_text = True,96 auto_line_break = True,97 )98 if rec_type == 'page':99 suffix = list(string.ascii_letters)100 random.shuffle(suffix)101 suffix = ''.join(suffix[:6])102 fp_suffix = f'{time.time()}-{suffix}'103 out_debug_dir = f'out-debug-{fp_suffix}'104 output_dir = OUTPUT_RESULT_DIR / f'output-{fp_suffix}'105 kwargs['save_debug_res'] = OUTPUT_RESULT_DIR / out_debug_dir106 elif rec_type == 'text_formula':107 suffix = list(string.ascii_letters)108 random.shuffle(suffix)109 suffix = ''.join(suffix[:6])110 out_det_fp = f'out-det-{time.time()}-{suffix}.jpg'111 kwargs['save_analysis_res'] = str(OUTPUT_RESULT_DIR / out_det_fp)112 113 out = p2t.recognize(image_file, file_type=rec_type, **kwargs)114 out_text = out115 if rec_type == 'page':116 out_text = out.to_markdown(output_dir)117 out_det_fp =kwargs['save_debug_res'] / 'layout_res.jpg'118 elif rec_type == 'text_formula':119 out_det_fp = kwargs['save_analysis_res']120 121 return out_text, out_det_fp122 123 124def example_func(lang_list, rec_type, image_file):125 return recognize(126 lang_list,127 mfd_model_name='yolov7 (paid)',128 mfr_model_name='mfr-pro (paid)',129 rec_type=rec_type,130 resized_shape=768,131 image_file=image_file,132 )133 134 135def main():136 langs = list(LANGUAGES.keys())137 langs.sort(key=lambda x: x.lower())138 139 title = ': a Free Alternative to Mathpix'140 examples = [141 [['English'], 'page', 'docs/examples/page.png',],142 [['English'], 'text_formula', 'docs/examples/mixed-en.jpg',],143 [['English', 'Chinese Simplified'], 'text_formula', 'docs/examples/mixed-ch_sim.jpg',],144 [145 ['English', 'Chinese Traditional'],146 'text_formula',147 'docs/examples/mixed-ch_tra.jpg',148 ],149 [['English', 'Vietnamese'], 'text_formula', 'docs/examples/mixed-vietnamese.jpg',],150 [['English'], 'formula', 'docs/examples/formula1.png'],151 [['English'], 'formula', 'docs/examples/formula2.jpg'],152 [['English'], 'formula', 'docs/examples/hw-formula.png'],153 [['English', 'Chinese Simplified'], 'text', 'docs/examples/pure-text.jpg',],154 ]155 156 table_desc = """157<div align="center">158<img src="https://pix2text.readthedocs.io/zh/latest/figs/p2t-logo.png" width="120px"/>159 160[](https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fbreezedeus%2Fpix2text-demo)161 162[](https://discord.gg/GgD87WM8Tf)163 164| | |165| ------------------------------- | --------------------------------------- |166| 🏄 **Online Service** | [p2t.breezedeus.com](https://p2t.breezedeus.com) |167| 📖 **Doc** | [Online Doc](https://pix2text.readthedocs.io) |168| 📀 **Code** | [Github](https://github.com/breezedeus/pix2text) |169| 🤗 **Models** | [breezedeus/Models](https://huggingface.co/breezedeus) |170| 📄 **More Infos** | [Pix2Text Infos](https://www.breezedeus.com/article/pix2text) |171 172If useful, please help to **star 🌟 [Pix2Text](https://github.com/breezedeus/pix2text)** 🙏173</div>174 """175 176 with gr.Blocks() as demo:177 gr.HTML(178 f'<h1 style="text-align: center; margin-bottom: 1rem;"><a href="https://github.com/breezedeus/pix2text" target="_blank">Pix2Text V1.1</a>{title}</h1>'179 )180 with gr.Row(equal_height=False):181 with gr.Column(min_width=200, variant='panel', scale=3):182 gr.Markdown('### Settings')183 lang_list = gr.Dropdown(184 label='Text Languages',185 choices=langs,186 value=['English', 'Chinese Simplified'],187 multiselect=True,188 # info='Which languages to be recognized as Texts.',189 )190 mfd_model_name = gr.Dropdown(191 label='MFD Models',192 choices=['yolov7_tiny (free)', 'yolov7 (paid)'],193 value='yolov7 (paid)',194 )195 mfr_model_name = gr.Dropdown(196 label='MFR Models',197 choices=['mfr (free)', 'mfr-pro (paid)'],198 value='mfr-pro (paid)',199 )200 rec_type = gr.Dropdown(201 label='File Type',202 choices=['page', 'text_formula', 'formula', 'text'],203 value='text_formula',204 # info='Which type of image to be recognized.',205 )206 with gr.Accordion('More Options', open=False):207 resized_shape = gr.Slider(208 label='resized_shape',209 minimum=512,210 maximum=2048,211 value=768,212 step=32,213 )214 215 with gr.Column(scale=6, variant='compact'):216 gr.Markdown('### Upload Image to be Recognized')217 image_file = gr.Image(218 label='Image', type="pil", image_mode='RGB', show_label=False219 )220 sub_btn = gr.Button("Submit", variant="primary")221 222 with gr.Column(scale=2, variant='compact'):223 gr.Markdown(table_desc)224 with gr.Row(equal_height=False):225 with gr.Column(scale=1, variant='compact'):226 gr.Markdown('**Detection Result**')227 det_result = gr.Image(228 label='Detection Result', scale=1, show_label=False229 )230 with gr.Column(scale=1, variant='compact'):231 gr.Markdown(232 '**Recognition Results (Paste them into the [P2T Online Service](https://p2t.breezedeus.com) to view rendered outcomes)**'233 )234 rec_result = gr.Textbox(235 label=f'Recognition Result ',236 lines=5,237 value='',238 scale=1,239 show_label=False,240 show_copy_button=True,241 )242 # render_result = gr.Markdown(label=f'After Rendering', value='')243 # rec_result.change(latex_render, rec_result, render_result)244 sub_btn.click(245 recognize,246 inputs=[247 lang_list,248 mfd_model_name,249 mfr_model_name,250 rec_type,251 resized_shape,252 image_file,253 ],254 outputs=[rec_result, det_result],255 )256 257 gr.Examples(258 label='Examples',259 examples=examples,260 inputs=[lang_list, rec_type, image_file,],261 outputs=[rec_result, det_result],262 fn=example_func,263 cache_examples=os.getenv('CACHE_EXAMPLES') == '1',264 )265 266 demo.queue(max_size=10)267 demo.launch()268 269 270if __name__ == '__main__':271 main()272