Probabilistic Detectron2

This repository contains the official implementation of Estimating and Evaluating Regression Predictive Uncertainty in Deep Object Detectors.

This code extends the detectron2 framework to estimate bounding box covariance matrices, and is meant to be a starter kit for entering the domain of probabilistic object detection.

Disclaimer

This research code was produced by one person with a single set of eyes, it may contain bugs and errors that I did not notice by the time of release.

Updates

Date | Change --- | --- | 30-September-2021 |Added pip frozen requirements (requirements_pip_freeze.txt). 10-October-2021 | Added ability to perform inference on images without passing through specific dataset handlers.

Requirements

Software Support:

Name | Supported Versions --- | --- | Ubuntu |20.04 Python |3.8 CUDA |11.0+ Cudnn |8.0.1+ PyTorch |1.8+

To install requirements choose between a python virtualenv or build a docker image using the provided Dockerfile.

# Clone repo
git clone https://github.com/asharakeh/probdet.git
cd probdet
git submodule update --init --recursive

1. Virtual Environment Creation:

# Create python virtual env
mkvirtualenv probdet

# Add library path to virtual env
add2virtualenv src

# Install requirements
cat requirements.txt | xargs -n 1 -L 1 pip install

2. Docker Image

# Clone repo
git clone https://github.com/asharakeh/probdet.git
cd probdet/Docker

# Build docker image
sh build.sh 

Datasets

COCO Dataset

Download the COCO Object Detection Dataset here. The COCO dataset folder should have the following structure: <br>

 └── COCO_DATASET_ROOT
     |
     ├── annotations
     ├── train2017
     └── val2017

To create the corrupted datasets using Imagenet-C corruptions, run the following code:

python src/core/datasets/generate_coco_corrupted_dataset.py --dataset-dir=COCO_DATASET_ROOT

OpenImages Datasets

Download our OpenImages validation splits here. We created a tarball that contains both shifted and out-of-distribution data splits used in our paper to make our repo easier to use. Do not modify or rename the internal folders as those paths are hard coded in the dataset reader. We will refer to the root folder extracted from the tarball as OPENIM_DATASET_ROOT.

Training

To train the models in the paper, use this command:

python src/train_net.py
--num-gpus xx
--dataset-dir COCO_DATASET_ROOT
--config-file COCO-Detection/architecture_name/config_name.yaml
--random-seed xx
--resume

For an explanation of all command line arguments, use python src/train_net.py -h

Evaluation

To run model inference after training, use this command:

python src/apply_net.py 
--dataset-dir TEST_DATASET_ROOT 
--test-dataset test_dataset_name 
--config-file path/to/config.yaml 
--inference-config /path/to/inference/config.yaml 
--random-seed xx
--image-corruption-level xx

For an explanation of all command line arguments, use python src/apply_net.py -h

--image-corruption-level can vary between 0-5, with 0 being the original COCO dataset with no corruption. In addition, --image-corruption-level has no effect when used with OpenImages dataset splits.

--test-dataset can be one of coco_2017_custom_val, openimages_val, or openimages_ood_val. --dataset-dir corresponds to the root directory of the dataset used. Evaluation code will run inference on the test dataset and then will generate mAP, Negative Log Likelihood, Brier Score, Energy Score, and Calibration Error results. If only evaluation of metrics is required, add --eval-only to the above code snippet.

Inference on new images

We provide a script to perform inference on new images without passing through dataset handlers.

python single_image_inference.py 
--image-dir /path/to/image/dir
--output-dir /path/to/output/dir
--config-file /path/to/config/file 
--inference-config /path/to/inference/config 
--model-ckpt /path/to/model.pth

image-dir is a folder containing all images to be used for inference. output-dir is a folder to write the output json file containing probabilistic detections. model-ckpt is the path to the model checkpoint to be used for inference. Look below to download model checkpoints.

Configurations in the paper

We provide a list of config combinations that generate the architectures used in our paper:

Method Name | Config File | Inference Config File | Model --- | --- | --- |--- Deterministic RetinaNet | retinanet_R_50_FPN_3x.yaml| standard_nms.yaml | retinanet_R_50_FPN_3x.pth RetinaNet NLL | retinanet_R_50_FPN_3x_reg_var_nll.yaml | standard_nms.yaml | retinanet_R_50_FPN_3x_reg_var_nll.pth RetinaNet DMM | retinanet_R_50_FPN_3x_reg_var_dmm.yaml | standard_nms.yaml | retinanet_R_50_FPN_3x_reg_var_dmm.pth RetinaNet ES | retinanet_R_50_FPN_3x_reg_var_es.yaml | standard_nms.yaml | retinanet_R_50_FPN_3x_reg_var_es.pth --- | --- | --- | --- Deterministic FasterRCNN | faster_rcnn_R_50_FPN_3x.yaml| standard_nms.yaml |faster_rcnn_R_50_FPN_3x.pth FasterRCNN NLL | faster_rcnn_R_50_FPN_3x_reg_covar_nll.yaml | standard_nms.yaml |faster_rcnn_R_50_FPN_3x_reg_covar_nll.pth FasterRCNN DMM | faster_rcnn_R_50_FPN_3x_reg_var_dmm.yaml | standard_nms.yaml |faster_rcnn_R_50_FPN_3x_reg_var_dmm.pth FasterRCNN ES | faster_rcnn_R_50_FPN_3x_reg_var_es.yaml | standard_nms.yaml |faster_rcnn_R_50_FPN_3x_reg_var_es.pth --- | --- | --- | --- Deterministic DETR | detr_R_50.yaml| standard_nms.yaml | detr_R_50.pth DETR NLL | detr_R_50_reg_var_nll.yaml | standard_nms.yaml | detr_R_50_reg_var_nll.pth DETR DMM| detr_R_50_reg_var_dmm.yaml | standard_nms.yaml | detr_R_50_reg_var_dmm.pth DETR ES| detr_R_50_reg_var_es.yaml | standard_nms.yaml | detr_R_50_reg_var_es.pth

Experiments in the paper were performed on 5 models trained and evaluated using random seeds [0, 1000, 2000, 3000, 4000]. The variance in performance between different seeds was seen to be negligible, and the results of the top performing seed were reported.

Additional Configurations

The repo supports many more variants including dropout and ensemble methods for estimating epistemic uncertainty. We provide a list of config combinations that generate the architectures used in our paper:

Method Name | Config File | Inference Config File --- | --- | --- RetinaNet Classification Loss Attenuation | retinanet_R_50_FPN_3x_cls_la.yaml | standard_nms.yaml RetinaNet Dropout Post-NMS Uncertainty Computation| retinanet_R_50_FPN_3x_dropout.yaml | mc_dropout_ensembles_post_nms_mixture_of_gaussians.yaml RetinaNet Dropout Pre-NMS Uncertainty Computation| retinanet_R_50_FPN_3x_dropout.yaml | mc_dropout_ensembles_pre_nms.yaml RetinaNet BayesOD with NLL loss| retinanet_R_50_FPN_3x_reg_var_nll.yaml | bayes_od.yaml RetinaNet BayesOD with ES loss| retinanet_R_50_FPN_3x_reg_var_es.yaml | bayes_od.yaml RetinaNet BayesOD with ES loss and Dropout| retinanet_R_50_FPN_3x_reg_var_es_dropout.yaml | bayes_od_mc_dropout.yaml RetinaNet Ensembles Post-NMS Uncertainty Estimation with NLL loss| retinanet_R_50_FPN_3x_reg_var_nll.yaml (Need to train 5 Models with different random seeds) | ensembles_post_nms_mixture_of_gaussians.yaml RetinaNet Ensembles Pre-NMS Uncertainty Estimation with NLL loss| retinanet_R_50_FPN_3x_reg_var_nll.yaml (Need to train 5 Models with different random seeds) | ensembles_pre_nms.yaml RetinaNet Ensembles Post-NMS Uncertainty Estimation with ES loss| retinanet_R_50_FPN_3x_reg_var_es.yaml (Need to train 5 Models with different random seeds) | ensembles_post_nms_mixture_of_gaussians.yaml RetinaNet Ensembles Pre-NMS Uncertainty Estimation with ES loss| retinanet_R_50_FPN_3x_reg_var_es.yaml (Need to train 5 Models with different random seeds) | ensembles_pre_nms.yaml --- | --- | --- FasterRCNN Classification Loss Attenuation | faster_rcnn_R_50_FPN_3x_cls_la.yaml | standard_nms.yaml FasterRCNN [Dropout](https://arxiv.org/pdf/1506.021