Pix2Struct
This repository contains code for Pix2Struct: Screenshot Parsing as Pretraining for Visual Language Understanding.
We release pretrained checkpoints for the Base and Large models and code for finetuning them on the nine downstream tasks discussed in the paper. We are unable to release the pretraining data, but they can be replicated using the publicly available URLs released in the C4 dataset.
Getting Started
Clone the github repository, install the pix2struct package, and run
the tests to ensure that all dependencies were successfully installed.
git clone https://github.com/google-research/pix2struct.git
cd pix2struct
conda create -n pix2struct python=3.9
conda activate pix2struct
pip install -e ."[dev]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pytest
You may first need to install Java (sudo apt install default-jre) and
conda
if not already installed.
We will be using Google Cloud Storage (GCS) for data and model storage. For the
remaining documentation we will assume that the path to your own bucket and
directory is in the PIX2STRUCT_DIR environment variable:
export PIX2STRUCT_DIR="gs://<your_bucket>/<path_to_pix2struct_dir>"
The code for running experiments assumes this environment variable when looking for the preprocessed data.
Data Preprocessing
Our data preprocessing scripts are run with Dataflow
by default using the Apache Beam library.
They can also be run locally by turning off flags appearing after --.
For the remaining documentation we will assume that GCP project information is in the following environment variables:
export GCP_PROJECT=<your_project_id>
export GCP_REGION=<your_region>
Below are the commands required to preprocess each dataset. The results will
be written to $PIX2STRUCT_DIR/data/<task_name>/preprocessed/, which is the
file structure assumed in tasks.py.
TextCaps
mkdir -p data/textcaps
cd data/textcaps
curl -O https://dl.fbaipublicfiles.com/textvqa/data/textcaps/TextCaps_0.1_train.json
curl -O https://dl.fbaipublicfiles.com/textvqa/data/textcaps/TextCaps_0.1_val.json
curl -O https://dl.fbaipublicfiles.com/textvqa/data/textcaps/TextCaps_0.1_test.json
curl -O https://dl.fbaipublicfiles.com/textvqa/images/train_val_images.zip
curl -O https://dl.fbaipublicfiles.com/textvqa/images/test_images.zip
unzip train_val_images.zip
rm train_val_images.zip
unzip test_images.zip
rm test_images.zip
cd ..
gsutil -m cp -r textcaps_data $PIX2STRUCT_DIR/data/textcaps
python -m pix2struct.preprocessing.convert_textcaps \
--textcaps_dir=$PIX2STRUCT_DIR/data/textcaps \
--output_dir=$PIX2STRUCT_DIR/data/textcaps/processed \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
ChartQA
mkdir -p data/chartqa
cd data/chartqa
git clone https://github.com/vis-nlp/ChartQA.git
cp -r ChartQA/ChartQA\ Dataset/* ./
rm -rf ChartQA
cd ..
gsutil -m cp -r chartqa $PIX2STRUCT_DIR/data/chartqa
python -m pix2struct.preprocessing.convert_chartqa \
--data_dir=$PIX2STRUCT_DIR/data/chartqa \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
RICO Images
Screen2Words, RefExp, and Widget Captioning all require images from the RICO dataset. If you'd like to use any of these datasets, please process RICO images before proceeding.
cd data
wget https://storage.googleapis.com/crowdstf-rico-uiuc-4540/rico_dataset_v0.1/unique_uis.tar.gz
tar xvfz unique_uis.tar.gz
rm unique_uis.tar.gz
gsutil -m cp -r combined $PIX2STRUCT_DIR/data/rico_images
Widget Captioning
If you haven't already setup RICO, please do so before you proceed.
mkdir -p data/widget_captioning
cd data/widget_captioning
git clone https://github.com/google-research-datasets/widget-caption.git
cp widget-caption/widget_captions.csv ./
cp widget-caption/split/*.txt ./
mv dev.txt val.txt
rm -rf widget-caption
cd ..
gsutil -m cp -r widget_captioning $PIX2STRUCT_DIR/data/widget_captioning
python -m pix2struct.preprocessing.convert_widget_captioning \
--data_dir=$PIX2STRUCT_DIR/data/widget_captioning \
--image_dir=$PIX2STRUCT_DIR/data/rico_images \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
Screen2Words
If you haven't already setup RICO, please do so before you proceed.
cd data
git clone https://github.com/google-research-datasets/screen2words.git
gsutil -m cp -r screen2words $PIX2STRUCT_DIR/data/screen2words
python -m pix2struct.preprocessing.convert_screen2words \
--screen2words_dir=$PIX2STRUCT_DIR/data/screen2words \
--rico_dir=$PIX2STRUCT_DIR/data/rico_images \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
RefExp
If you haven't already setup RICO, please do so before you proceed.
mkdir -p data/refexp
cd data/refexp
wget https://github.com/google-research-datasets/uibert/raw/main/ref_exp/train.tfrecord
wget https://github.com/google-research-datasets/uibert/raw/main/ref_exp/dev.tfrecord
wget https://github.com/google-research-datasets/uibert/raw/main/ref_exp/test.tfrecord
mv dev.tfrecord val.tfrecord
cd ..
gsutil -m cp -r refexp $PIX2STRUCT_DIR/data/refexp
python -m pix2struct.preprocessing.convert_refexp \
--data_dir=$PIX2STRUCT_DIR/data/refexp \
--image_dir=$PIX2STRUCT_DIR/data/rico_images \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
DocVQA
mkdir -p data/docvqa
cd data/docvqa
Download DocVQA (Single Document Visual Question Answering) from
the official source (requires
registration). The following steps assume that the train/val/test.tar.gz files
are in data/docvqa.
tar xvf train.tar.gz
tar xvf val.tar.gz
tar xvf test.tar.gz
rm -r *.tar.gz */ocr_results
cd ..
gsutil -m cp -r docvqa $PIX2STRUCT_DIR/data/docvqa
python -m pix2struct.preprocessing.convert_docvqa \
--data_dir=$PIX2STRUCT_DIR/data/docvqa \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
InfographicVQA
mkdir -p data/infographicvqa
cd data/infographicvqa
Download InfographicVQA Task 1 from this
website (requires registration). The following steps assume that the
train/val/test.json and the zip files are in data/infographicvqa.
for split in train val test
do
unzip infographicVQA_${split}_v1.0_images.zip
mv infographicVQA_${split}_v1.0_images $split
mv infographicVQA_${split}_v1.0.json $split/${split}_v1.0.json
done
rm *.zip
cd ..
gsutil -m cp -r infographicvqa $PIX2STRUCT_DIR/data/infographicvqa
python -m pix2struct.preprocessing.convert_docvqa \
--data_dir=$PIX2STRUCT_DIR/data/infographicvqa \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
OCR-VQA
mkdir -p data/ocrvqa
cd data/ocrvqa
Follow instructions on the OCR-VQA website to
download the data into data/ocrvqa (requires crawling). The following steps
assume that data/ocrvqa contains a directory called images and a file called
dataset.json.
cd ..
gsutil -m cp -r ocrvqa $PIX2STRUCT_DIR/data/ocrvqa
python -m pix2struct.preprocessing.convert_ocrvqa \
--data_dir=$PIX2STRUCT_DIR/data/ocrvqa \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
AI2D
mkdir -p data/
cd data/
wget https://ai2-public-datasets.s3.amazonaws.com/diagrams/ai2d-all.zip
unzip ai2d-all.zip
rm ai2d-all.zip
gsutil -m cp -r ai2d $PIX2STRUCT_DIR/data/ai2d
python -m pix2struct.preprocessing.convert_ai2d \
--data_dir=$PIX2STRUCT_DIR/data/ai2d \
--test_ids_path=gs://pix2struct-data/ai2d_test_ids.csv \
-- \
--runner=DataflowRunner \
--save_main_session \
--project=$GCP_PROJECT \
--region=$GCP_REGION \
--temp_location=$PIX2STRUCT_DIR/data/temp \
--staging_location=$PIX2STRUCT_DIR/data/staging \
--setup_file=./setup.py
Running experiments
The main experiments are implemented as a light wrapper around the T5X library. For brevity, we illustrate an example workflow of finetuning the pretrained base Pix2Struct model on the Screen2Words dataset. To scale up to larger setups, please see to the T5X documentation.
Setting up the TPU
Following official instructions
for running JAX on a Cloud TPU VM, which allows you to directly ssh into the
TPU host.
In this example, we are using a v3-8 TPU:
TPU_TYPE=v3-8
TPU_NAME=pix2struct-$TPU_TYPE
TPU_ZONE=europe-west4-a
gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$TPU_ZONE \
--accelerator-type=$TPU_TYPE \
--version=tpu-vm-base
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$TPU_ZONE
Once you have sshed into the TPU host, follow the "Getting Started"
instructions to install the pix2struct package.
Training
The following command will initiate the training loop, which consists of train steps interleaved with evaluations on the validation set.
python -m t5x.train \
--gin_search_paths="pix2struct/configs" \
--gin_file="models/pix2struct.gin" \
--gin_file="runs/train.gin" \
--gin_file="sizes/base.gin" \
--gin_file="optimizers/adafactor.gin" \
--gin_file="schedules/screen2words.gin" \
--gin_file="init/pix2struct_base_init.gin" \
--gin.MIXTURE_OR_TASK_NAME="'screen2words'" \
--gin.MODEL_DIR="'$PIX2STRUCT_DIR/experiments/screen2words_base'" \
--gin.TASK_FEATURE_LENGTHS="{'inputs': 4096, 'targets': 128}" \
--gin.BATCH_SIZE=32
Evaluation
The following command evaluates the model on the test set. You will need to replace the checkpoint path with the one that was actually selected based on the validation performance.
python -m t5x.eval \
--gin_search_paths="pix2struct/configs" \
--gin_file="models/pix2struct.gin" \
--gin_file="runs/eval.gin" \
--gin_file="sizes/base.gin" \
--gin.MIXTURE_OR_TASK_NAME="'screen2words'" \
--gin.CHECKPOINT_PATH="'$PIX2STRUCT_DIR/experiments/screen2words_base/checkpoint_286600'" \
--gin.EVAL_OUTPUT_DIR="'$PIX2STRUCT_DIR/experiments/test_exp/test_eval'" \
--gin.EVAL_SPLIT="'test'" \
--gin.TASK_FEATURE_LENGTHS="{'inputs': 4096, 'targets': 128}" \
--gin.BATCH_SIZE=32
Finetuned Checkpoints
In addition to the pretrained checkpoints released and specified in the
configs/init directory. We also release checkpoints for the finetuned models
on all tasks below.
| Task | GCS Path (Base) | GCS Path (Large) |
|---|---|---|
| TextCaps | gs://pix2struct-data/textcaps_base/checkpoint_280400 |
gs://pix2struct-data/textcaps_large/checkpoint_180600 |
| ChartQA | gs://pix2struct-data/chartqa_base/checkpoint_287600 |
gs://pix2struct-data/charqa_large/checkpoint_182600 |
| WidgetCaptioning | gs://pix2struct-data/widget_captioning_base/checkpoint_281600 |
gs://pix2struct-data/widget_captioning_large/checkpoint_181600 |
| Screen2Words | gs://pix2struct-data/screen2words_base/checkpoint_282600 |
gs://pix2struct-data/screen2words_large/checkpoint_183000 |
| RefExp | gs://pix2struct-data/refexp_base/checkpoint_290000 |
gs://pix2struct-data/refexp_large/checkpoint_187800 |
| DocVQA | gs://pix2struct-data/docvqa_base/checkpoint_284400 |
gs://pix2struct-data/docvqa_large/checkpoint_184000 |
| InfographicVQA | gs://pix2struct-data/infographicvqa_base/checkpoint_284000 |
gs://pix2struct-data/infographicvqa_large/checkpoint_182000 |
| OCR-VQA | gs://pix2struct-data/ocrvqa_base/checkpoint_290000 |
gs://pix2struct-data/ocrvqa_large/checkpoint_188400 |
| AI2D | gs://pix2struct-data/ai2d_base/checkpoint_284400 |
gs://pix2struct-data/ai2d_large/checkpoint_184000 |
These checkpoints are compatible with the eval command documented above and the
two ways of performing inference mentioned below. Please ensure that the config
file under configs/sizes is set to be consistent with the checkpoint.
Inference
We provide two ways of performing inference. For testing and demoing purposes,
these may be run on CPU. In that case, please set the JAX_PLATFORMS
environment variable to cpu.
Command-line example
We provide a minimal script for performing inference on a single example. This
path has only been tested at extremely small scale and is not meant for
larger-scale inference. For large-scale inference, we recommend setting a custom
task with placeholder labels and running the evaluation script (t5x.eval) as
documented above.
In the following example, we show the command for predicting the caption of an
image using a base-sized checkpoint finetuned on the TextCaps task. For a task
that also accepts textual prompts such as questions in VQA, you can also supply
the question via the text flag (in addition to specifying the image with the
image flag).
python -m pix2struct.example_inference \
--gin_search_paths="pix2struct/configs" \
--gin_file=models/pix2struct.gin \
--gin_file=runs/inference.gin \
--gin_file=sizes/base.gin \
--gin.MIXTURE_OR_TASK_NAME="'placeholder_pix2struct'" \
--gin.TASK_FEATURE_LENGTHS="{'inputs': 2048, 'targets': 128}" \
--gin.BATCH_SIZE=1 \
--gin.CHECKPOINT_PATH="'gs://pix2struct-data/textcaps_base/checkpoint_280400'" \
--image=$HOME/test_image.jpg
Web Demo
For a more user-friendly demo, we also provide a web-based alternative of
inference script above. While running this command, the web demo can be accessed
at localhost:8080 (or any port specified via the port flag), assuming you
are running the demo locally. You can then upload your custom image and optional
prompt instead of specifying it via the command line.
python -m pix2struct.demo \
--gin_search_paths="pix2struct/configs" \
--gin_file=models/pix2struct.gin \
--gin_file=runs/inference.gin \
--gin_file=sizes/base.gin \
--gin.MIXTURE_OR_TASK_NAME="'placeholder_pix2struct'" \
--gin.TASK_FEATURE_LENGTHS="{'inputs': 2048, 'targets': 128}" \
--gin.BATCH_SIZE=1 \
--gin.CHECKPOINT_PATH="'gs://pix2struct-data/textcaps_base/checkpoint_280400'"
Clean up
When you are done with your TPU VM, remember to delete the instance:
gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$TPU_ZONE
Note
This is not an officially supported Google product.