pix2tex (LaTeX OCR)

repository·main·Indexed 12 days ago

https://github.com/lukas-blecher/latex-ocr

A learning-based system that converts images of mathematical formulas into LaTeX code using a ViT encoder with a ResNet backbone and a Transformer decoder. It provides a CLI, a GUI (latexocr), a Streamlit-based API, and a Python API via the LatexOCR class. The system supports model training and fine-tuning using custom datasets and YAML configurations.

Tokens
4.9K
Snippets
24
Records
26
Agent score
94%

What's inside pix2tex

  1. Use the pix2tex CLI, GUI, and API packages

    main

    The pix2tex project provides several entry points for interacting with the LaTeX OCR model:

    • pix2tex.cli: Command-line interface for running OCR tasks from the terminal.
    • pix2tex.gui: A graphical user interface for interacting with the model visually.
    • pix2tex.api: Programmatic interfaces for integrating the OCR model into applications. This includes:
      • pix2tex.api.app: Core application logic.
      • pix2tex.api.streamlit: A Streamlit-based web interface for the model.
    • pix2tex.models: Contains the underlying model architectures, such as vit (Vision Transformer) and hybrid models.
    • pix2tex.dataset: Tools for managing and generating datasets, including modules for ArXiv scraping, LaTeX-to-PNG rendering, and data extraction.
  2. Use the pix2tex CLI and GUI

    main

    There are two primary ways to interact with the model via command line or interface:

    1. CLI Tool: Call pix2tex directly to parse images from your disk or from your clipboard.
    2. GUI: Call latexocr to launch a user interface. This allows you to take screenshots; the predicted LaTeX is rendered using MathJax and automatically copied to your clipboard.

    Linux Screenshot Tool Configuration: On Linux, the GUI uses gnome-screenshot by default (which supports multiple monitors). For other environments, you can set the SCREENSHOT_TOOL environment variable:

    • wlroots-based Wayland: Use grim and slurp.
    • KDE Plasma: Use spectacle.
    • Other: pil is also an available value.

    If gnome-screenshot is not compatible with your compositor, set the variable explicitly:

    export SCREENSHOT_TOOL=spectacle
    latexocr
  3. Install pix2tex via pip

    main

    To use pix2tex, you need Python 3.7 or higher. Ensure PyTorch is installed on your system before installing the package. You can install specific dependency sets using extras tags:

    • [gui]: Includes dependencies for the Graphical User Interface.
    • [train]: Includes dependencies for training the model.
    • [api]: Includes dependencies for the API.
    • [all]: Includes all dependencies.

    Model checkpoints are automatically downloaded upon the first execution of the script.

    pip install pix2tex[gui]
  4. Use pix2tex API via Streamlit or Docker

    main

    To use the model as an API, install the additional dependencies:

    pip install -U "pix2tex[api]"

    Running via Python

    Start a Streamlit demo that connects to the API at port 8502:

    python -m pix2tex.api.run

    Running via Docker

    API only:

    docker pull lukasblecher/pix2tex:api
    docker run --rm -p 8502:8502 lukasblecher/pix2tex:api

    API with Streamlit demo:

    docker run --rm -it -p 8501:8501 --entrypoint python lukasblecher/pix2tex:api pix2tex/api/run.py

    Then navigate to http://localhost:8501/.

    pip install -U "pix2tex[api]"
    python -m pix2tex.api.run
  5. Train the pix2tex model

    main

    To train the model, install the training dependencies:

    pip install "pix2tex[train]"

    1. Prepare the Dataset

    Combine images with their ground truth labels into a .pkl file using the dataset module:

    python -m pix2tex.dataset.dataset --equations path_to_textfile --images path_to_images --out dataset.pkl

    To create a custom tokenizer:

    python -m pix2tex.dataset.dataset --equations path_to_textfile --vocab-size 8000 --out tokenizer.json

    2. Configure Training

    Edit the data (and valdata) entries in your configuration file to point to your generated .pkl file. If using a custom tokenizer, update the path to the tokenizer and set num_tokens to your vocabulary size. Use pix2tex/model/settings/config.yaml as a template.

    3. Run Training

    Execute the training run with your config file:

    python -m pix2tex.train --config path_to_config_file
    python -m pix2tex.dataset.dataset --equations path_to_textfile --images path_to_images --out dataset.pkl
    python -m pix2tex.train --config path_to_config_file
  6. Set up LaTeX OCR in Google Colab

    main

    To use LaTeX OCR in a Colab environment, you must install the pix2tex and Pillow packages. Note that if your current Pillow version is older than version 9, the environment requires a mandatory kernel restart. If you see a message saying 'Mandatory restart: Execute this cell again!', simply run the setup cell a second time.

    Key dependencies installed:

    • pix2tex
    • Pillow
    • opencv-python-headless==4.1.2.30
    %reload_ext autoreload
    %autoreload
    import PIL
    !pip install Pillow -U -qq
    if int(PIL.__version__[0]) < 9:
        print('Mandatory restart: Execute this cell again!')
        import os
        os.kill(os.getpid(), 9)
    !pip install pix2tex -qq
    !pip install opencv-python-headless==4.1.2.30 -U -qq
  7. Prepare the training environment and datasets

    main

    Training requires setting up a specific directory structure and downloading datasets (such as handwritten CROHME data and PDF math data).

    1. Install supporting utilities: gpustat, opencv-python-headless, and gdown.
    2. Create directories: LaTeX-OCR, dataset/data, and images.
    3. Download datasets using gdown with specific Google Drive IDs.
    4. Unzip the downloaded archives.
    5. Split the data by moving a subset of images into a validation directory (e.g., valimages).
    # Example setup sequence
    !pip install gpustat -q
    !pip install opencv-python-headless==4.1.2.30 -U -q
    !pip install --upgrade --no-cache-dir gdown -q
    
    !mkdir -p dataset/data
    !mkdir images
    !gdown -O dataset/data/crohme.zip --id 13vjxGYrFCuYnwgDIUqkxsNGKk__D_sOM
    !gdown -O dataset/data/pdf.zip --id 176PKaCUDWmTJdQwc-OfkO0y8t4gLsIvQ
    !gdown -O dataset/data/pdfmath.txt --id 1QUjX6PFWPa-HBWdcY-7bA5TRVUnbyS1D
    
    # Split handwritten data into val set and train set
    os.chdir('images')
    !mkdir ../valimages
    !ls | shuf -n 1000 | xargs -i mv {} ../valimages
  8. Configure and run the training process

    main

    Training is initiated via the pix2tex.train module using a YAML configuration file.

    Key Configuration Options:

    • data: Path to the training .pkl file.
    • valdata: Path to the validation .pkl file.
    • load_chkpt: Path to the pretrained weights (e.g., weights.pth) for fine-tuning.
    • tokenizer: Path to the tokenizer.json file.
    • debug: Set to true if not using wandb (Weights & Biases).
    • model_path: Directory where checkpoints are saved.
    • output_path: Directory for training outputs.
    • batchsize, lr (learning rate), epochs, optimizer, scheduler: Standard hyperparameter controls.

    Execution:

    python -m pix2tex.train --config <your_config>.yaml
    # Example training command
    !python -m pix2tex.train --config colab.yaml