CenterNet w/ PyTorchLightning

Description

My attempt at a cleaner implementation of the glorious CenterNet.

Features

• Decoupled backbones and heads for easier backbone integration
• Split sample creation into image augmentation (with imgaug) and actual sample creation
• Comes shipped with Lightning modules but can also be used with good ol' plain PyTorch
• Stripped all code not used to reproduce the results in the paper
• Smaller code base with more meaningful variable names
• Requires significantly less memory
• Same or slightly better results than the original implementation

ToDos

Some features of the original repository are not implemented yet but pull requests are welcome!

• 3D bounding box detection
• ExtremeNet detection
• Pascal VOC dataset

How to run

First, install dependencies

# Install ninja for DCNv2 JIT compilation
sudo apt-get install ninja-build

# clone CenterNet
git clone https://github.com/tteepe/CenterNet-pytorch-lightning

# install CenterNet
cd CenterNet-pytorch-lightning
pip install -e .   
pip install -r requirements.txt

Next, navigate to any file and run it.

# module folder
cd CenterNet

# run module
python centernet_detection.py    
python centernet_multi_pose.py    

Imports

This project is setup as a package which means you can now easily import any file into any other file like so:

from pytorch_lightning import Trainer
from torchvision.datasets import CocoDetection
from CenterNet import CenterNetDetection

# model
model = CenterNetDetection("dla_34")

# data
train = CocoDetection("train2017", "instances_train2017.json")
val = CocoDetection("val2017", "instances_val2017.json")

# train
trainer = Trainer()
trainer.fit(model, train, val)

# test using the best backbone!
test = CocoDetection("test2017", "image_info_test2017.json")
trainer.test(test_dataloaders=test)

BibTeX

If you want to cite the implementation feel free to use this or zenodo:

@article{teepe2021centernet,
  title={CenterNet PyTorch Lightning},
  author={Teepe, Torben and Gilg, Johannes},
  journal={GitHub. Note: https://github.com/tteepe/CenterNet-pytorch-lightning},
  volume={1},
  year={2021}
}