Code for "Image Generation from Scene Graphs", Johnson et al, CVPR 2018
Switch branches/tags
Nothing to show
Clone or download
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
images Initial commit Jun 29, 2018
scene_graphs Initial commit Jun 29, 2018
scripts Add script to sample images Oct 28, 2018
sg2im Fix bug in VG data loader by converting PyTorch scalars to Python sca… Oct 16, 2018
.gitignore Initial commit Jun 29, 2018
CONTRIBUTING.md Initial commit Jun 29, 2018
LICENSE Initial commit Jun 29, 2018
README.md Initial commit Jun 29, 2018
TRAINING.md Initial commit Jun 29, 2018
requirements.txt Delete pkg-resources dependency Jul 4, 2018

README.md

sg2im

This is the code for the paper

Image Generation from Scene Graphs
Justin Johnson, Agrim Gupta, Li Fei-Fei
Presented at CVPR 2018

Please note that this is not an officially supported Google product.

A scene graph is a structured representation of a visual scene where nodes represent objects in the scene and edges represent relationships between objects. In this paper we present and end-to-end neural network model that inputs a scene graph and outputs an image.

Below we show some example scene graphs along with images generated from those scene graphs using our model. By modifying the input scene graph we can exercise fine-grained control over the objects in the generated image.


If you find this code useful in your research then please cite

@inproceedings{johnson2018image,
  title={Image Generation from Scene Graphs},
  author={Johnson, Justin and Gupta, Agrim and Fei-Fei, Li},
  booktitle={CVPR},
  year={2018}
}

Model

The input scene graph is processed with a graph convolution network which passes information along edges to compute embedding vectors for all objects. These vectors are used to predict bounding boxes and segmentation masks for all objects, which are combined to form a coarse scene layout. The layout is passed to a cascaded refinement network (Chen an Koltun, ICCV 2017) which generates an output image at increasing spatial scales. The model is trained adversarially against a pair of discriminator networks which ensure that output images look realistic.

Setup

All code was developed and tested on Ubuntu 16.04 with Python 3.5 and PyTorch 0.4.

You can setup a virtual environment to run the code like this:

python3 -m venv env               # Create a virtual environment
source env/bin/activate           # Activate virtual environment
pip install -r requirements.txt   # Install dependencies
echo $PWD > env/lib/python3.5/site-packages/sg2im.pth  # Add current directory to python path
# Work for a while ...
deactivate  # Exit virtual environment

Pretrained Models

You can download pretrained models by running the script bash scripts/download_models.sh. This will download the following models, and will require about 355 MB of disk space:

  • sg2im-models/coco64.pt: Trained to generate 64 x 64 images on the COCO-Stuff dataset. This model was used to generate the COCO images in Figure 5 from the paper.
  • sg2im-models/vg64.pt: Trained to generate 64 x 64 images on the Visual Genome dataset. This model was used to generate the Visual Genome images in Figure 5 from the paper.
  • sg2im-models/vg128.pt: Trained to generate 128 x 128 images on the Visual Genome dataset. This model was used to generate the images in Figure 6 from the paper.

Table 1 in the paper presents an ablation study where we disable various components of the full model. You can download the additional models used in this ablation study by running the script bash scripts/download_ablated_models.sh. This will download 12 additional models, requiring and additional 1.25 GB of disk space.

Running Models

You can use the script scripts/run_model.py to easily run any of the pretrained models on new scene graphs using a simple human-readable JSON format. For example you can replicate the sheep images above like this:

python scripts/run_model.py \
  --checkpoint sg2im-models/vg128.pt \
  --scene_graphs scene_graphs/figure_6_sheep.json \
  --output_dir outputs

The generated images will be saved to the directory specified by the --output_dir flag. You can control whether the model runs on CPU or GPU using py passing the flag --device cpu or --device gpu.

We provide JSON files and pretrained models allowing you to recreate all images from Figures 5 and 6 from the paper.

(Optional): GraphViz

This script can also draw images for the scene graphs themselves using GraphViz; to enable this option just add the flag --draw_scene_graphs 1 and the scene graph images will also be saved in the output directory. For this option to work you must install GraphViz; on Ubuntu 16.04 you can simply run sudo apt-get install graphviz.

Training new models

Instructions for training new models can be found here.