Home

Awesome

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

LICENSE Python PyTorch

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

Scene Graph Generation Model

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

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

# 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

Acknowledgment

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