# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import hashlib
import io
import numpy as np
import pandas as pd
import pypdfium2
import streamlit as st
from PIL import Image
from rapid_latex_ocr import LatexOCR
from streamlit_drawable_canvas import st_canvas
MAX_WIDTH = 800
MAX_HEIGHT = 1000
st.set_page_config(layout="wide")
@st.cache_resource()
def load_model_cached():
return LatexOCR()
def get_canvas_hash(pil_image):
return hashlib.md5(pil_image.tobytes()).hexdigest()
def open_pdf(pdf_file):
stream = io.BytesIO(pdf_file.getvalue())
return pypdfium2.PdfDocument(stream)
@st.cache_data()
def page_count(pdf_file):
doc = open_pdf(pdf_file)
return len(doc)
@st.cache_data()
def get_page_image(pdf_file, page_num, dpi=96):
doc = open_pdf(pdf_file)
renderer = doc.render(
pypdfium2.PdfBitmap.to_pil,
page_indices=[page_num - 1],
scale=dpi / 72,
)
png = list(renderer)[0]
png_image = png.convert("RGB")
return png_image
@st.cache_data()
def get_uploaded_image(in_file):
if isinstance(in_file, Image.Image):
return in_file.convert("RGB")
return Image.open(in_file).convert("RGB")
def resize_image(pil_image):
if pil_image is None:
return
pil_image.thumbnail((MAX_WIDTH, MAX_HEIGHT), Image.Resampling.LANCZOS)
@st.cache_data()
def get_image_size(pil_image):
if pil_image is None:
return MAX_HEIGHT, MAX_WIDTH
height, width = pil_image.height, pil_image.width
return height, width
if __name__ == "__main__":
st.markdown(
"",
unsafe_allow_html=True,
)
st.markdown(
"""
""",
unsafe_allow_html=True,
)
in_file = st.file_uploader(
"PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"]
)
if in_file is None:
st.stop()
filetype = in_file.type
if "pdf" in filetype:
page_count = page_count(in_file)
page_number = st.number_input(
f"Page number out of {page_count}:",
min_value=1,
value=1,
max_value=page_count,
)
pil_image = get_page_image(in_file, page_number)
else:
pil_image = get_uploaded_image(in_file)
resize_image(pil_image)
canvas_hash = get_canvas_hash(pil_image) if pil_image else "canvas"
model = load_model_cached()
canvas_result = st_canvas(
fill_color="rgba(255, 165, 0, 0.1)",
stroke_width=1,
stroke_color="#FFAA00",
background_color="#FFF",
background_image=pil_image,
update_streamlit=True,
height=get_image_size(pil_image)[0],
width=get_image_size(pil_image)[1],
drawing_mode="rect",
point_display_radius=0,
key=canvas_hash,
)
if canvas_result.json_data is not None:
objects = pd.json_normalize(canvas_result.json_data["objects"])
bbox_list = None
if objects.shape[0] > 0:
boxes = objects[objects["type"] == "rect"][
["left", "top", "width", "height"]
]
boxes["right"] = boxes["left"] + boxes["width"]
boxes["bottom"] = boxes["top"] + boxes["height"]
bbox_list = boxes[["left", "top", "right", "bottom"]].values.tolist()
if bbox_list:
bbox_nums = len(bbox_list)
for i, bbox in enumerate(bbox_list):
input_img = pil_image.crop(bbox)
rec_res, elapse = model(np.array(input_img))
st.markdown(f"#### {i + 1}. (cost: {elapse:.3f}s)")
st.latex(rec_res)
st.code(rec_res)