Mapping Images to Scene Graphs with Permutation-Invariant Structured Prediction
Scene graph prediction is the task of mapping an image into a set of bounding boxes, along with their categories and relations.
We present a new architecture for graph inference that has the following structural property: on the one hand, the architecture is invariant to input permutations; on the other hand, every permutation-invariant function can be implemented via this architecture.
In this repository, we share our architecture implementation for the task of scene graph prediction.
Scene Graph Predictor (SGP) gets as an input inital confidience distributions per entity and relation and processes these to obtain new labels. SGP satisfies the graph permutation invariance property intoduced in the paper. The model is implemented in TensorFlow. For the initial confidence distributions per entity and relation, we simply re-use features learned by the baseline model from Zellers et al. (2017). (git repositiry https://github.com/rowanz/neural-motifs)
Our SGP implementation is using an iteratively RNN to process predictions. Each step outputs an improved predictions.
A schematic representation of the architecture. Given an image, a Label predictor outputs initial predictions . Then, our SGP model, computes each
element wise. Next, they are summed to create vector
, which is concatenated with
. Then,
is applied, and another summation creates the graph representation. Finally,
classifies objects and
classifies relation. The process of SGP could be repeated iteratively (in the paper we repeat it 2 times).
For more information, please look at the code (Module/Module.py file) and the paper.
Our SGP architecture uses attention at the feature-level for each node during inference. We weight the significance of each feature per node, such that the network can choose which features from adjacent nodes contributes the most information.
An example of attention per entities and global attention over all nodes. The size and location of objects provide a key signal to the attention mechanism. The model assigns higher confidence for the label "tie" when the label "shirt" is detected (third panel from the left). Similarly, the model assigns a higher confidence for the label "eye" when it is located near "hair".
To get started with the framework, install the following dependencies:
Run "pip install -r requirements.txt"
- to install all the requirements.
- Run
"python Run.py download"
to download and extract train, validation and test data. The data already contains the result of applying the baseline detecor over the VisualGenome data. - Run
"python Run.py eval gpi_linguistic_pretrained <gpu-number>"
to evaluate the pre-trained model of our best variant, linguistic with multi-head attention. (recall@100 SG Classification). - Run
"python Run.py train gpi_linguistic <gpu-number>"
to train a new model (linguistic with multi-head attention). - Run
"python Run.py eval gpi_linguistic_best <gpu-number>"
to evaluate the new model. (recall@100 SG Classification).