ellipse-rcnn

KO | EN

gitlite — search

#javascript #python #hacktoberfest #react #ai #typescript #llm #go #golang #android #machine-learning #rust #deep-learning #linux

★ 49 Open GitHub ↗

A PyTorch implementation of Ellipse R-CNN

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Ellipse R-CNN

</div>

A PyTorch (Lightning) implementation of Ellipse R-CNN. Extracted from [another project](https://github.com/wdoppenberg/crater-detection).
The methodology is based on [Ellipse R-CNN: Learning to Infer Elliptical Object from Clustering and Occlusion](https://arxiv.org/abs/2001.11584), 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](https://vis-www.cs.umass.edu/fddb/) (FDDB) dataset as well.

![Sample FDDB predictions](https://raw.githubusercontent.com/wdoppenberg/ellipse-rcnn/main/docs/fddb_sample.png)
![Sample crater predictions](https://raw.githubusercontent.com/wdoppenberg/ellipse-rcnn/main/docs/craters_sample.png)

</div>

## Installation

```shell
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](https://huggingface.co/wdoppenberg).

Run the following:

```python
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:

</div>

## Training Setup

Make sure you have [uv](https://docs.astral.sh/uv/getting-started/installation/) 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:

```shell
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:

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

If you want to do experiment tracking, use tensorboard:

```shell
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:

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

## FDDB Data

The [Face Detection Dataset & Benchmark](https://vis-www.cs.umass.edu/fddb/) (FDDB) [module](ellipse_rcnn/data/fddb.py) 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](https://github.com/wdoppenberg/crater-detection/blob/main/src/common/data.py), and ensure
you have a license to run [SurRender](https://www.airbus.com/en/products-services/space/space-customer-support/surrendersoftware).

# Citations

If you use this code in your work, please consider citing the [original paper](https://arxiv.org/abs/2001.11584) 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.}
}
```

ellipse-rcnn

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README