Official Implementation of "SGTR: End-to-end Scene Graph Generation with Transformer"

Our paper SGTR: End-to-end Scene Graph Generation with Transformer has been accepted by CVPR 2022.

Installation

Check INSTALL.md for installation instructions.

Dataset

Check DATASET.md for instructions of dataset preprocessing.

Model Zoo

SGTR performance:

The methods implemented in our toolkit and reported results are given in Model Zoo.md

Training (IMPORTANT)

Our codebase is developed based on cvpods, you can refer to usage of cvpods and tutorial.

Prepare Faster-RCNN Detector

• You can download the pretrained DETR on ResNet-101 we used in the paper:

  • VG,
  • OIv6,

• Unzip the checkpoint into the folder

• Then, you need to modify the pre-trained weight parameter MODEL.WEIGHTS in config.py playground/experiment/path/config_xx.py to the path of corresponding pre-trained detector weight to make sure you load the detection weight parameter correctly.

• Besides, you can also train your own detector weight by the provide configs in Model Zoo.md

Scene Graph Generation Model

• You can follow the following instructions to train your own, which takes 4 GPUs for train each SGG model. The results should be very close to the reported results given in paper.

You can run training by:

# activate the environments
conda activate sgtr
cvpods_env

# move to experiment directory
cd playground/sgg/detr.res101.c5.one_stage_rel_tfmer
# link the config file
rm config.py; ln -s config_vg.py config.py 

pods_train --num-gpus 4

• We also provide the trained model .pth and config.json of SGTR(vg),SGTR(oiv6)

More Training args:

Single Machine 4-GPU Run

pods_train --num-gpus 4

If you want check the machine with available gpus:

pods_train --num-gpus 4 --gpu-check

Test

• First, you need download and unzip the provided trained model parameters.
• By replacing the parameter of MODEL.TEST_WEIGHTS in config.py to the trained model weight you can directly eval the model on validation or test set.
  • The program will load the config file (config.json) in same directory of model weights.
• You can use the following scripts for directly produce the results from the checkpoint provide by us.

# visual genome dataset
pods_test --num-gpus 4 DATASETS.TEST "(vg_test,)"
pods_test --num-gpus 4 DATASETS.TEST "(vg_val,)"

# openimage dataset
pods_test --num-gpus 4 DATASETS.TEST "(oiv6_test,)"
pods_test --num-gpus 4 DATASETS.TEST "(oiv6_val,)"

Citations

If you find this project helps your research, please kindly consider citing our papers in your publications.

@InProceedings{Li_2022_CVPR,
    author    = {Li, Rongjie and Zhang, Songyang and He, Xuming},
    title     = {SGTR: End-to-end Scene Graph Generation with Transformer},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {19486-19496}
}

Q&A

• Import ipdb in anywhere in your code will cause the multi-process initialization error, try pdb when you debug in multi-process mode.

Acknowledgment

This repository borrows code from scene graph benchmarking framework developed by KaihuaTang and Rongjie Li