Spaces:
Running
Running
| # -*- encoding: utf-8 -*- | |
| # @Author: SWHL | |
| # @Contact: liekkaskono@163.com | |
| from pathlib import Path | |
| import numpy as np | |
| import streamlit as st | |
| from PIL import Image | |
| from rapid_latex_ocr import LatexOCR | |
| from streamlit_cropper import st_cropper | |
| from streamlit_image_select import image_select | |
| st.set_option("deprecation.showfileUploaderEncoding", False) | |
| class RecEquation: | |
| def __init__(self, model_dir: str): | |
| model_dir = Path(model_dir) | |
| image_resizer_path = model_dir / "image_resizer.onnx" | |
| encoder_path = model_dir / "encoder.onnx" | |
| decoder_path = model_dir / "decoder.onnx" | |
| tokenizer_json = model_dir / "tokenizer.json" | |
| self.model = LatexOCR( | |
| image_resizer_path=str(image_resizer_path), | |
| encoder_path=str(encoder_path), | |
| decoder_path=str(decoder_path), | |
| tokenizer_json=str(tokenizer_json), | |
| ) | |
| def __call__(self, img: np.ndarray): | |
| result, elapse = self.model(img) | |
| return result, elapse | |
| if __name__ == "__main__": | |
| st.markdown( | |
| "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidLatexOCR' style='text-decoration: none'>Rapid Latex OCR</a></h1>", | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| """ | |
| <p align="left"> | |
| <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a> | |
| <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a> | |
| <a href="https://pepy.tech/project/rapid_latex_ocr"><img src="https://static.pepy.tech/personalized-badge/rapid_latex_ocr?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a> | |
| <a href="https://pypi.org/project/rapid_latex_ocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid_latex_ocr"></a> | |
| <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a> | |
| <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a> | |
| </p> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| # Upload an image and set some options for demo purposes | |
| img_file = st.sidebar.file_uploader(label="Upload a file", type=["png", "jpg"]) | |
| realtime_update = st.sidebar.checkbox(label="Update in Real Time", value=False) | |
| box_color = st.sidebar.color_picker(label="Box Color", value="#0000FF") | |
| aspect_choice = st.sidebar.radio( | |
| label="Aspect Ratio", options=["Free", "1:1", "16:9", "4:3", "2:3"] | |
| ) | |
| aspect_dict = { | |
| "Free": None, | |
| "1:1": (1, 1), | |
| "16:9": (16, 9), | |
| "4:3": (4, 3), | |
| "2:3": (2, 3), | |
| } | |
| aspect_ratio = aspect_dict[aspect_choice] | |
| with st.sidebar.container(): | |
| img = image_select( | |
| label="Examples（click to select）:", | |
| images=[ | |
| "images/equation.png", | |
| "images/eq_2.png", | |
| "images/eq_3.png", | |
| "images/eq_4.png", | |
| ], | |
| key="equation_default", | |
| use_container_width=False, | |
| ) | |
| rec_eq_sys = RecEquation(model_dir="models") | |
| select_img_container = st.container() | |
| st.markdown("#### Select image:") | |
| img_empty = st.empty() | |
| img_empty.image(img, use_column_width=False) | |
| rec_res, elapse = rec_eq_sys(img) | |
| if img_file: | |
| img = Image.open(img_file) | |
| # Get a cropped image from the frontend | |
| with select_img_container: | |
| if not realtime_update: | |
| select_img_container.markdown("#### Double click to save crop") | |
| img = st_cropper( | |
| img, | |
| realtime_update=realtime_update, | |
| box_color=box_color, | |
| aspect_ratio=aspect_ratio, | |
| ) | |
| img_empty.image(img, use_column_width=False) | |
| rec_res, elapse = rec_eq_sys(np.array(img)) | |
| st.markdown(f"#### Rec result (cost: {elapse:.4f}s):") | |
| st.latex(rec_res) | |
| st.markdown("#### Latex source code:") | |
| st.code(rec_res) | |