A PyTorch (Lightning) implementation of Ellipse R-CNN. Extracted from another project. The methodology is based on Ellipse R-CNN: Learning to Infer Elliptical Object from Clustering and Occlusion, albeit with slight changes. Primarily this implementation was made to enable Crater detection from Moon orbiter sensors, but works with the Face Detection Dataset & Benchmark (FDDB) dataset as well.

Installation

pip install ellipse-rcnn

Optional extras

Enable a feature with the ellipse-rcnn[<FEATURE>, ...] pattern.

• train: Installs all dependencies necessary to train this model.
• hf: Installs huggingface-hub and safetensors for easy weights saving & loading through the Huggingface platform.

Quickstart

• Install with all extras through pip install "ellipse-rcnn[hf,train]"
• Select the weights you need from my Huggingface profile.

Run the following:

import torch
from ellipse_rcnn.hf import EllipseRCNN  # This is the model with HF functionality included through PyTorchModelHubMixin
from PIL import Image
from torchvision.transforms.functional import to_tensor
from ellipse_rcnn.utils.viz import plot_single_pred


model = EllipseRCNN.from_pretrained("wdoppenberg/crater-rcnn")  # For the crater detection model
model.eval()

png = Image.open("docs/example_craters.png").convert("L")
img = to_tensor(png)
with torch.no_grad():
    pred = model([img])

plot_single_pred(img, pred)

This should output the following:

Training Setup

Make sure you have uv installed. This library's dependencies have mostly been made optional to make it easy to import the base EllipseRCNN model. For training, more dependencies, most importantly pytorch-lightning, are under the train group.

For a training setup, run the following from the project root:

uv sync --extra train

This sets up a new virtual environment and installs all packages into it.

Get info on either the training or sample (prediction) script using:

uv run scripts/train.py --help
# or
uv run scripts/sample.py --help

If you want to do experiment tracking, use tensorboard:

uvx tensorboard --logdir=lightning_logs

Data

Currently the training script only supports training with FDDB. See the required steps for getting & structuring the data below. More datasets can be supported if needed. If you want to download a dataset directly, use the following script:

uv run scripts/download.py <DATASET_NAME> [Optional: --root <ROOT_FOLDER>]

FDDB Data

The Face Detection Dataset & Benchmark (FDDB) module contains the Dataset subclass that does all the data loading and transformations. It will download and unpack the necessary files to ./data/FDDB. Simply running the training script will download the necessary files.

Craters Data

Unfortunately, the dataset used to create a Crater Detector is not releasable. To generate it, check out this module, and ensure you have a license to run SurRender.

Citations

If you use this code in your work, please consider citing the original paper and this repository. For this repository, you can use the following BibTex entry:

@misc{EllipseRCNNPyTorch2025,
    author = {Doppenberg, Wouter},
    title = {Ellipse {R-CNN}: {PyTorch} {Implementation}},
    year = {2025},
    publisher = {GitHub},
    journal = {GitHub repository},
    howpublished = {\url{https://github.com/wdoppenberg/ellipse_rcnn}},
    note = {An implementation of the Ellipse R-CNN object detection model in PyTorch, based on "Ellipse R-CNN: Learning to Infer Elliptical Object from Clustering and Occlusion" by Dong et al.}
}