To setup the environment, we use conda to manage our dependencies.

Our developers use CUDA 10.1 to do experiments.

You can specify the appropriate cudatoolkit version to install on your machine in the environment.yml file, and then run the following to create the conda environment:

conda env create -f environment.yml

You shall manually install the following dependencies.

# Install mmcv
## CAUTION: The latest versions of mmcv 1.5.3, mmdet 2.25.0 are not well supported, due to bugs in mmdet.
pip install mmcv-full==1.4.3 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.7.0/index.html

# Install mmdet
pip install openmim
mim install mmdet==2.20.0

# Install coco panopticapi
pip install git+https://github.com/cocodataset/panopticapi.git

# For visualization
conda install -c conda-forge pycocotools
pip install detectron2==0.5 -f \
  https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.7/index.html

# If you're using wandb for logging
pip install wandb
wandb login

# If you develop and run openpsg directly, install it from source:
pip install -v -e .
# "-v" means verbose, or more output
# "-e" means installing a project in editable mode,
# thus any local modifications made to the code will take effect without reinstallation.

To explore and download their original Dataset: Original Dataset Website

Or you can directly download their original dataset using the link they provided: Original Dataset Link

And some pre-trained models that possibly be used: Pre-trained Weights

Our codebase accesses the datasets from ./data/ and pre-trained models from ./work_dirs/checkpoints/ by default.

Our checkpoint for PSGTR is available at: Enhanced PSGTR

The enhanced training dataset (enhanced by BAI) is also provided: Enhanced training dataset

1. Download the Original Dataset to ./data/.
2. Unzip the coco.zip.
3. Download some pre-trained models that possibly be used.

Run the scripts below:

PYTHONPATH='.':$PYTHONPATH  \
python tools/test.py  \
configs/psgtr/psgtr_r50_psg.py \
work_dirs/PATH_TO_PSGTR_CHECKPOINT \
--eval sgdet

Tips on how to evaluate our PSGTR checkpoint:

1. Modify the variable "ann_file" in configs/base/datasets/psg.py (Line 4-5) to './data/psg/psg.json' (psg.json includes the test set).
2. Download our PSGTR checkpoint and modify the PATH_TO_PSGTR_CHECKPOINT in the script..
3. Run the script above.

Run the scripts below:

PYTHONPATH='.':$PYTHONPATH  \
python tools/train.py  \
configs/psgtr/psgtr_r50_psg.py \
work_dirs/PATH_TO_PSGTR_CHECKPOINT \
--launcher pytorch

Tips on how to do the fine-tuning (like for PSGTR):

1. Download our enhanced training dataset, and place the enhanced training dataset in ./data/psg/.
2. Modify the variable "ann_file" in configs/base/datasets/psg.py (Line 4-5) to the path of the enhanced training dataset.
3. Download the original PSGTR checkpoint and modify the PATH_TO_PSGTR_CHECKPOINT in the script.
4. Run the script above.