zebiao1998 / latex-ocr

pix2tex: Using a ViT to convert images of equations into LaTeX code.

pix2tex - LaTeX OCR

The goal of this project is to create a learning based system that takes an image of a math formula and returns corresponding LaTeX code.

Requirements

Model

• PyTorch (tested on v1.7.1)
• Python 3.7+ & dependencies (requirements.txt)

  pip install -r requirements.txt
  

Dataset

In order to render the math in many different fonts we use XeLaTeX, generate a PDF and finally convert it to a PNG. For the last step we need to use some third party tools:

• XeLaTeX
• ImageMagick with Ghostscript. (for converting pdf to png)
• Node.js to run KaTeX (for normalizing Latex code)
• de-macro >= 1.4 (only for parsing arxiv papers)
• Python 3.7+ & dependencies (requirements.txt)

Using the model

1. Download/Clone this repository
2. For now you need to install the Python dependencies specified in requirements.txt (look above)
3. Download the weights.pth (and optionally image_resizer.pth) file from the Releases->Assets section and place it in the checkpoints directory

Thanks to @katie-lim, you can use a nice user interface as a quick way to get the model prediction. Just call the GUI with python gui.py. From here you can take a screenshot and the predicted latex code is rendered using MathJax and copied to your clipboard.

If the model is unsure about the what's in the image it might output a different prediction every time you click "Retry". With the temperature parameter you can control this behavior (low temperature will produce the same result).

Alternatively you can use pix2tex.py with similar functionality as gui.py, only as command line tool. In this case you don't need to install PyQt5. Using this script you can also parse already existing images from the disk.

Note: As of right now it works best with images of smaller resolution. Don't zoom in all the way before taking a picture. Double check the result carefully. You can try to redo the prediction with an other resolution if the answer was wrong.

Update: I have trained an image classifier on randomly scaled images of the training data to predict the original size. This model will automatically resize the custom image to best resemble the training data and thus increase performance of images found in the wild. To use this preprocessing step, all you have to do is download the second weights file mentioned above. You should be able to take bigger (or smaller) images of the formula and still get a satisfying result

Training the model

1. First we need to combine the images with their ground truth labels. I wrote a dataset class (which needs further improving) that saves the relative paths to the images with the LaTeX code they were rendered with. To generate the dataset pickle file run

python dataset/dataset.py --equations path_to_textfile --images path_to_images --tokenizer dataset/tokenizer.json --out dataset.pkl

You can find my generated training data on the Google Drive as well (formulae.zip - images, math.txt - labels). Repeat the step for the validation and test data. All use the same label text file.

2. Edit the data (and valdata) entry in the config file to the newly generated .pkl file. Change other hyperparameters if you want to. See settings/config.yaml for a template.
3. Now for the actual training run

python train.py --config path_to_config_file

If you want to use your own data you might be interested in creating your own tokenizer with

python dataset/dataset.py --equations path_to_textfile --vocab-size 8000 --out tokenizer.json

Don't forget to update the path to the tokenizer in the config file and set num_tokens to your vocabulary size.

Model

The model consist of a ViT [1] encoder with a ResNet backbone and a Transformer [2] decoder.

Performance

BLEU score	normed edit distance
0.88	0.10

Data

We need paired data for the network to learn. Luckily there is a lot of LaTeX code on the internet, e.g. wikipedia, arXiv. We also use the formulae from the im2latex-100k dataset. All of it can be found here

Fonts

Latin Modern Math, GFSNeohellenicMath.otf, Asana Math, XITS Math, Cambria Math

TODO

• add more evaluation metrics
• create a GUI
• add beam search
• support handwritten formulae
• reduce model size (distillation)
• find optimal hyperparameters
• tweak model structure
• fix data scraping and scrape more data
• trace the model

Contribution

Contributions of any kind are welcome.

Acknowledgment

Code taken and modified from lucidrains, rwightman, im2markup, arxiv_leaks, pkra: Mathjax, harupy: snipping tool

References

[1] An Image is Worth 16x16 Words

[2] Attention Is All You Need