SWHL/RapidLaTeXOCRDemo

1# -*- encoding: utf-8 -*-2# @Author: SWHL3# @Contact: liekkaskono@163.com4import hashlib5import io6 7import numpy as np8import pandas as pd9import pypdfium210import streamlit as st11from PIL import Image12from rapid_latex_ocr import LaTeXOCR13from streamlit_drawable_canvas import st_canvas14 15MAX_WIDTH = 80016MAX_HEIGHT = 100017 18st.set_page_config(layout="wide")19 20 21@st.cache_resource()22def load_model_cached():23    return LaTeXOCR()24 25 26def get_canvas_hash(pil_image):27    return hashlib.md5(pil_image.tobytes()).hexdigest()28 29 30def open_pdf(pdf_file):31    stream = io.BytesIO(pdf_file.getvalue())32    return pypdfium2.PdfDocument(stream)33 34 35@st.cache_data()36def page_count(pdf_file):37    doc = open_pdf(pdf_file)38    return len(doc)39 40 41@st.cache_data()42def get_page_image(pdf_file, page_num, dpi=96):43    doc = open_pdf(pdf_file)44    renderer = doc.render(45        pypdfium2.PdfBitmap.to_pil,46        page_indices=[page_num - 1],47        scale=dpi / 72,48    )49    png = list(renderer)[0]50    png_image = png.convert("RGB")51    return png_image52 53 54@st.cache_data()55def get_uploaded_image(in_file):56    if isinstance(in_file, Image.Image):57        return in_file.convert("RGB")58    return Image.open(in_file).convert("RGB")59 60 61def resize_image(pil_image):62    if pil_image is None:63        return64    pil_image.thumbnail((MAX_WIDTH, MAX_HEIGHT), Image.Resampling.LANCZOS)65 66 67@st.cache_data()68def get_image_size(pil_image):69    if pil_image is None:70        return MAX_HEIGHT, MAX_WIDTH71    height, width = pil_image.height, pil_image.width72    return height, width73 74 75if __name__ == "__main__":76    st.markdown(77        "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidLaTeXOCR' style='text-decoration: none'>Rapid ⚡︎ LaTeX OCR</a></h1>",78        unsafe_allow_html=True,79    )80    st.markdown(81        """82    <p align="center">83        <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>84        <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>85        <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>86        <a href="https://pypi.org/project/rapid_latex_ocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid_latex_ocr"></a>87        <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>88        <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>89        <a href="https://github.com/RapidAI/RapidLaTeXOCR"><img src="https://img.shields.io/badge/Github-link-brightgreen.svg"></a>90    </p>91    """,92        unsafe_allow_html=True,93    )94 95    in_file = st.file_uploader(96        "PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"]97    )98 99    if in_file is None:100        st.stop()101 102    filetype = in_file.type103    if "pdf" in filetype:104        page_count = page_count(in_file)105        page_number = st.number_input(106            f"Page number out of {page_count}:",107            min_value=1,108            value=1,109            max_value=page_count,110        )111        pil_image = get_page_image(in_file, page_number)112    else:113        pil_image = get_uploaded_image(in_file)114 115    resize_image(pil_image)116    canvas_hash = get_canvas_hash(pil_image) if pil_image else "canvas"117 118    model = load_model_cached()119    canvas_result = st_canvas(120        fill_color="rgba(255, 165, 0, 0.1)",121        stroke_width=1,122        stroke_color="#FFAA00",123        background_color="#FFF",124        background_image=pil_image,125        update_streamlit=True,126        height=get_image_size(pil_image)[0],127        width=get_image_size(pil_image)[1],128        drawing_mode="rect",129        point_display_radius=0,130        key=canvas_hash,131    )132 133    if canvas_result.json_data is not None:134        objects = pd.json_normalize(canvas_result.json_data["objects"])135        bbox_list = None136        if objects.shape[0] > 0:137            boxes = objects[objects["type"] == "rect"][138                ["left", "top", "width", "height"]139            ]140            boxes["right"] = boxes["left"] + boxes["width"]141            boxes["bottom"] = boxes["top"] + boxes["height"]142            bbox_list = boxes[["left", "top", "right", "bottom"]].values.tolist()143 144        if bbox_list:145            bbox_nums = len(bbox_list)146            for i, bbox in enumerate(bbox_list):147                input_img = pil_image.crop(bbox)148                rec_res, elapse = model(np.array(input_img))149                st.markdown(f"#### {i + 1}. (cost: {elapse:.3f}s)")150                st.latex(rec_res)151                st.code(rec_res)152