Pix2Struct

repository·main·Indexed 20 days ago

https://github.com/google-research/pix2struct

A framework for screenshot parsing and visual language understanding. It provides pretrained Base and Large models and tools for finetuning on downstream tasks such as DocVQA, ChartQA, and Screen2Words. The framework includes preprocessing modules for various datasets and implements training and evaluation as a wrapper around the T5X library.

Tokens: 6.3K
Snippets: 17
Records: 17
Agent score: 22%

What's inside Pix2Struct

Preprocess InfographicVQA dataset

main

Download InfographicVQA Task 1 (requires registration). Unzip the image files and JSON files into their respective split directories (train, val, test), upload to GCS, and run pix2struct.preprocessing.convert_docvqa (note: the command uses convert_docvqa for this task) using the DataflowRunner.

mkdir -p data/infographicvqa
cd data/infographicvqa
# (Download infographicVQA files from official source)
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

Preprocess TextCaps dataset

main

To preprocess the TextCaps dataset, download the JSON annotations and images from Facebook's public servers, unzip them, and upload the data to your PIX2STRUCT_DIR. Then, run the pix2struct.preprocessing.convert_textcaps module using Apache Beam with the DataflowRunner to process the data in the cloud.

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

Configure environment variables for Pix2Struct
main
Pix2Struct uses Google Cloud Storage (GCS) for data and model storage. You must set the PIX2STRUCT_DIR environment variable to point to your GCS bucket and directory. Additionally, for data preprocessing via Dataflow, set your GCP project and region:
- PIX2STRUCT_DIR: The GCS path for data and model storage (e.g., gs://<your_bucket>/<path>).
- GCP_PROJECT: Your Google Cloud Project ID.
- GCP_REGION: Your Google Cloud Region.
```
export PIX2STRUCT_DIR="gs://<your_bucket>/<path_to_pix2struct_dir>"
export GCP_PROJECT=<your_project_id>
export GCP_REGION=<your_region>
```
https://github.com/google-research/pix2struct/blob/main/README.md

Preprocess ChartQA dataset

main

To preprocess ChartQA, clone the ChartQA repository, move the dataset files into your local directory, upload them to GCS, and run the pix2struct.preprocessing.convert_chartqa module using the DataflowRunner.

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

Preprocess DocVQA dataset

main

Download DocVQA from the official source (requires registration). Extract the train, val, and test tarballs, remove the ocr_results directories, upload to GCS, and run pix2struct.preprocessing.convert_docvqa using the DataflowRunner.

mkdir -p data/docvqa
cd data/docvqa
# (Download train.tar.gz, val.tar.gz, test.tar.gz from official source)
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

Preprocess Widget Captioning dataset

main

Widget Captioning requires the RICO images to be prepared first. Download the widget-caption dataset, move the CSV and split files into your local directory, upload to GCS, and run pix2struct.preprocessing.convert_widget_captioning using the DataflowRunner.

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

Prepare RICO images for downstream tasks
main
The datasets Screen2Words, RefExp, and Widget Captioning all require images from the RICO dataset. To prepare them, download the unique_uis.tar.gz file, extract it, and upload the combined directory to your $PIX2STRUCT_DIR/data/rico_images path.
```
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
```
https://github.com/google-research/pix2struct/blob/main/README.md

Preprocess OCR-VQA dataset

main

Follow the instructions on the OCR-VQA website to download the data into data/ocrvqa. Ensure the directory contains an images folder and a dataset.json file. Upload to GCS and run pix2struct.preprocessing.convert_ocrvqa using the DataflowRunner.

mkdir -p data/ocrvqa
cd data/ocrvqa
# (Download data via OCR-VQA website)
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

Preprocess RefExp dataset

main

RefExp requires RICO images. Download the tfrecord files for train, dev, and test, rename dev.tfrecord to val.tfrecord, upload to GCS, and run pix2struct.preprocessing.convert_refexp using the DataflowRunner.

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

Preprocess AI2D dataset

main

Download the AI2D dataset, unzip it, upload to GCS, and run pix2struct.preprocessing.convert_ai2d using the DataflowRunner. You must provide the --test_ids_path pointing to the AI2D test IDs CSV on GCS.

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

Install Pix2Struct
main
To install Pix2Struct, clone the repository, create a Conda environment with Python 3.9, and install the package in editable mode along with development dependencies. You will need to provide a URL for the libtpu_releases to ensure compatibility with TPU environments. It is also recommended to have Java installed (sudo apt install default-jre).
After installation, run pytest to verify that all dependencies are correctly configured.
```
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
```
https://github.com/google-research/pix2struct/blob/main/README.md

Preprocess Screen2Words dataset

main

Screen2Words requires RICO images. Clone the screen2words repository, upload it to GCS, and run pix2struct.preprocessing.convert_screen2words using the DataflowRunner.

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