import io import os import sys import argparse import numpy as np import torch import hashlib import pypdfium2 import pandas as pd import streamlit as st from PIL import Image from streamlit_drawable_canvas import st_canvas import unimernet.tasks as tasks from unimernet.common.config import Config from unimernet.processors import load_processor MAX_WIDTH = 872 MAX_HEIGHT = 1024 class ImageProcessor: """ImageProcessor class handles the loading of the model and processing of images.""" def __init__(self, cfg_path): self.cfg_path = cfg_path self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model, self.vis_processor = self.load_model_and_processor() def load_model_and_processor(self): # Load the model and visual processor from the configuration args = argparse.Namespace(cfg_path=self.cfg_path, options=None) cfg = Config(args) task = tasks.setup_task(cfg) model = task.build_model(cfg).to(self.device) vis_processor = load_processor( "formula_image_eval", cfg.config.datasets.formula_rec_eval.vis_processor.eval, ) return model, vis_processor def process_single_image(self, pil_image): # Process an image and return the LaTeX string image = self.vis_processor(pil_image).unsqueeze(0).to(self.device) output = self.model.generate({"image": image}) pred = output["pred_str"][0] return pred @st.cache_data(show_spinner=False) def read_markdown(path): with open(path, "r", encoding="utf-8") as f: data = f.read() return data def open_pdf(pdf_file): stream = io.BytesIO(pdf_file.getvalue()) return pypdfium2.PdfDocument(stream) @st.cache_data() def get_page_image(pdf_file, page_num, dpi=300): # Extract an image from a PDF page 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): # Load an uploaded image file return Image.open(in_file).convert("RGB") def resize_image(pil_image): # Resize an image to fit within the MAX_WIDTH and MAX_HEIGHT if pil_image is None: return pil_image.thumbnail((MAX_WIDTH, MAX_HEIGHT), Image.Resampling.LANCZOS) def display_image_cropped(pil_image, bbox): # Display a cropped portion of an image cropped_image = pil_image.crop(bbox) st.image(cropped_image, use_column_width=True) @st.cache_data() def page_count_fn(pdf_file): # Return the number of pages in a PDF doc = open_pdf(pdf_file) return len(doc) def get_canvas_hash(pil_image): return hashlib.md5(pil_image.tobytes()).hexdigest() @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 @st.cache_data(hash_funcs={ImageProcessor: id}) def infer_image(processor, pil_image, bbox): # Perform inference on a cropped image input_img = pil_image.crop(bbox) pred = processor.process_single_image(input_img) return pred @st.cache_resource() def load_image_processor(cfg_path): processor = ImageProcessor(cfg_path) return processor def run_mode1(): """Direct Recognition mode: recognize formulas directly from an image """ col1, col2 = st.columns([0.5, 0.5]) in_file = st.sidebar.file_uploader( "Input Image:", type=["png", "jpg", "jpeg", "gif", "webp"] ) if in_file is None: st.stop() filetype = in_file.type pil_image = get_uploaded_image(in_file) resize_image(pil_image) with col1: st.image(pil_image, use_column_width=True) st.markdown( "

[Input: Image]

", unsafe_allow_html=True, ) bbox_list = [(0, 0, pil_image.width, pil_image.height)] with col2: inferences = [infer_image(processor, pil_image, bbox) for bbox in bbox_list] for idx, (bbox, inference) in enumerate( zip(reversed(bbox_list), reversed(inferences)) ): st.latex(inference) st.markdown( "

[Prediction: Rendered Image]

", unsafe_allow_html=True, ) st.divider() st.code(inference) st.markdown( "

[Prediction: LaTeX Code]

", unsafe_allow_html=True, ) def run_mode2(): """Manual Selection mode: allows users to select formulas in an image or PDF for recognition. """ col1, col2 = st.columns([0.7, 0.3]) in_file = st.sidebar.file_uploader( "PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"] ) if in_file is None: st.stop() # Determine if the uploaded file is a PDF or an image whole_image = False if in_file.type == "application/pdf": page_count = page_count_fn(in_file) page_number = st.sidebar.number_input( "Page number:", 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) whole_image = st.sidebar.button("Formula Recognition") resize_image(pil_image) canvas_hash = get_canvas_hash(pil_image) if pil_image else "canvas" with col1: # Create a canvas component where users can draw rectangles to select formulas canvas_result = st_canvas( fill_color="rgba(255, 165, 0, 0.1)", # Fixed fill color with some opacity 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, ) # Process the drawn rectangles or the whole image if 'whole_image' is True if canvas_result.json_data is not None or whole_image: objects = pd.json_normalize(canvas_result.json_data["objects"]) bbox_list = [] 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 whole_image: bbox_list = [(0, 0, pil_image.width, pil_image.height)] if bbox_list: with col2: # Perform inference on each selected area and display results inferences = [infer_image(processor, pil_image, bbox) for bbox in bbox_list] for idx, (bbox, inference) in enumerate(zip(reversed(bbox_list), reversed(inferences))): st.markdown(f"### Result {len(inferences) - idx}") st.markdown( "
[Input: Image]
", unsafe_allow_html=True, ) display_image_cropped(pil_image, bbox) st.markdown( "
[Prediction: Rendered Image]
", unsafe_allow_html=True, ) st.latex(inference) st.markdown( "
[Prediction: LaTeX Code]
", unsafe_allow_html=True, ) st.code(inference) st.divider() with col2: tips = """ ### Usage tips - Draw a box around the equation to get the prediction.""" st.markdown(tips) def run_mode3(): st.markdown("Coming Soon!") if __name__ == "__main__": st.set_page_config(layout="wide") html_code = """

UniMERNet Online Demo

This App is based on UniMERNet. There are three optional modes for mathematical expression recognition:
""" readme_text = st.markdown(html_code, unsafe_allow_html=True) root_path = os.path.abspath(os.getcwd()) config_path = os.path.join(root_path, "configs/demo.yaml") processor = load_image_processor(config_path) app_mode = st.sidebar.selectbox( "Switch Mode:", ["Direct Recognition", "Manual Selection", "Auto Detection"] ) # Direct Recognition: Input an image containing formulas and output the recognition results. if app_mode == "Direct Recognition": st.markdown("---") st.markdown( "

Direct Recognition

", unsafe_allow_html=True, ) run_mode1() # Manual Selection: Input a document or webpage screenshot, detect all formulas, then recognize each one. elif app_mode == "Manual Selection": st.markdown("---") st.markdown( "

Manual Selection and Recognition

", unsafe_allow_html=True, ) run_mode2() # Auto Detection: Input an image or document, and the model automatically detects and recognizes all formulas. elif app_mode == "Auto Detection": st.markdown("---") st.markdown( "

Auto Detection and Recognition (Coming Soon)

", unsafe_allow_html=True, ) run_mode3()