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