latex-ocr / app.py
Young Ho Shin
Add examples and article.md
36bccd1
raw
history blame
2.38 kB
import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import requests
from PIL import Image
url = 'https://huggingface.co/yhshin/latex-ocr/raw/main/tokenizer-wordlevel.json'
r = requests.get(url)
open('tokenizer-wordlevel.json' , 'wb').write(r.content)
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-small-printed")
model = VisionEncoderDecoderModel.from_pretrained("yhshin/latex-ocr")
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file("tokenizer-wordlevel.json")
# load image examples
def process_image(image):
# prepare image
pixel_values = processor(image, return_tensors="pt").pixel_values
# generate (no beam search)
generated_ids = model.generate(pixel_values)
# decode
generated_text = tokenizer.decode_batch(generated_ids.tolist(), skip_special_tokens=True)[0]
# Strip spaces
generated_text = generated_text.replace(" ", "")
return generated_text
# !ls examples | grep png
# +
title = "Convert an image of an equation to LaTeX source code"
with open('article.md',mode='r') as file:
article = file.read()
description = """
This is a demo of machine learning model trained to parse an image and reconstruct the LaTeX source code of an equation.
To use it, simply upload an image or use one of the example images below and click 'submit'.
Results will show up in a few seconds.
Try rendering the equation [here](https://quicklatex.com/) to compare with the original image.
(The model is not perfect yet, so you may need to edit the resulting LaTeX a bit to get it to render a good match.)
"""
examples = [
[ "examples/1d32874f02.png" ],
[ "examples/1e466b180d.png" ],
[ "examples/2d3503f427.png" ],
[ "examples/2f9d3c4e43.png" ],
[ "examples/51c5cc2ff5.png" ],
[ "examples/545a492388.png" ],
[ "examples/6a51a30502.png" ],
[ "examples/6bf6832adb.png" ],
[ "examples/7afdeff0e6.png" ],
[ "examples/b8f1e64b1f.png" ],
]
#examples =[["examples/image_0.png"], ["image_1.png"], ["image_2.png"]]
# -
iface = gr.Interface(fn=process_image,
inputs=[gr.inputs.Image(type="pil")],
outputs=gr.outputs.Textbox(),
title=title,
description=description,
article=article,
examples=examples)
iface.launch()