Spaces:
Running
Running
File size: 4,080 Bytes
6936493 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: [email protected]
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)
|