This project is part of the Rice D2K Capstone in collaboration with NASA Johnson Space Center. The goal is to develop a robust lunar crater detection model that supports autonomous spacecraft navigation by accurately identifying and fitting parameterized ellipses to crater rims in images of the Moon.
Spacecraft navigating beyond GPS coverage rely on costly ground-based tracking systems. To enable future autonomous lunar missions, this project leverages deep learning and computer vision to detect craters under challenging lighting and perspective conditions. Unlike previous approaches, our solution is designed to run efficiently on low-power spaceflight hardware.
- Python 3.12
- Conda or a virtual environment manager
-
Clone the repository:
git clone <repository-url> cd NASA_crater_Sp25
-
Install the required dependencies:
pip install "ellipse-rcnn[train,hf]"pip install -r requirements.txt
-
Set up the environment variables:
- Copy
example.envto.envand update the values as needed.
- Copy
-
Download the dataset. You can follow the section below for a detailed description.
-
(Optional) Install additional tools for fine-tuning:
- Follow the instructions in model/FINETUNE.md.
Use the LunarLens class to run end-to-end crater detection:
from pipeline import LunarLens
import cv2
# Load a grayscale image and convert to BGR
image = cv2.imread("path/to/image.png", cv2.IMREAD_GRAYSCALE)
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
# Initialize and run the pipeline
pipe = LunarLens(
yolo_model_path="model/yolo/weights/best.pt",
ellipse_choice="bbox_estimate",
yolo_confidence=0.7,
image=image
)
pipe.predict()
pipe.plot()For a detailed walkthrough, see example_run_pipeline.ipynb.
Use the download_data.py script to download and extract the required datasets into the data/ directory. This will download the
YOLO and EllipseRCNN respective datasets used during finetuning.
python script/download_data.pyYou can also find the final datasets used for training here:
For detailed information on the data augmentations for YOLO, please refer to the augmentation strategies documentation.
NASA_crater_Sp25/
├── .env # Environment variables
├── .gitignore # Git ignore rules
├── example.env # Template for environment variables
├── README.md # Project documentation
├── requirements.txt # Python dependencies
├── .vscode/ # VSCode configurations
│ └── launch.json # Debugging configuration
├── data/ # Directory for datasets
│ ├── moon_batch_0/ # Example dataset batch
│ └── ... # Additional dataset batches
├── data_generation/ # Data preprocessing scripts
│ ├── craterdetection/ # Crater detection utilities
│ │ ├── craterCatalog.py
│ │ ├── image.py
│ │ ├── camera.py
│ │ ├── generateMasks.py
│ │ ├── cameraConfig.json
│ │ ├── catalogConfig.json
│ │ └── maskConfig.json
│ └── yolo/ # YOLO-based data generation scripts
├── model/ # Model training and evaluation scripts
│ ├── __init__.py
│ ├── evaluate_ercnn.py
│ ├── FINETUNE.md # Fine-tuning guide
│ ├── ercnn/ # EllipseRCNN model implementation
│ │ ├── modules.py
│ │ ├── dataset.py
│ │ ├── modal_setup.py
│ │ └── finetune.py
│ ├── utils/ # Utility functions
│ │ ├── inference_utils.py
│ │ ├── data_handling.py
│ │ ├── metrics.py
│ │ └── tiling.py
│ └── yolo/ # YOLO-based model scripts
├── script/ # General scripts used for data processing
│ ├── json_converter.py # json converter for bounding box data to Label Studio annotation data
│ ├── download_data.py # Data download helper
│ ├── evaluate_pipeline.py # Pipeline evaluation
│ └── Label Studio Tutorial.pdf # Label Studio usage tutorial (conda environment)
└── results/ # Output results (ignored by Git)
.env: Contains environment variables for the project.example.env: Template for.envfile.
data_generation/craterdetection/: Contains scripts for preprocessing lunar crater datasets.craterCatalog.py: Defines theCraterCatalogclass for managing crater data.image.py: Contains image processing utilities.generateMasks.py: Script for generating masks for crater images.
data_generation/yolo/: Contains scripts for generating YOLO-compatible datasets.
Please refer to the finetuning tutorial for an in-depth guide of how to tune the models in this project.
The model/ directory contains training and evaluation code for both EllipseRCNN and YOLO-based pipelines:
FINETUNE.md: High-level guide for fine‑tuning both pipelines.evaluate_ercnn.py: CLI script to evaluate an EllipseRCNN checkpoint and generate metrics.finetune_ercnn.py: Entry point for training and fine‑tuning the EllipseRCNN model.finetune_yolo.py: Entry point for training and fine‑tuning the YOLO model.ercnn/: EllipseRCNN implementation:modules.py: DefinesEllipseRCNNModuleandCraterEllipseDataModule(PyTorch Lightning modules).dataset.py:CraterDatasetclass for loading and preprocessing crater images and ellipse annotations.modal_setup.py: Configuration and environment setup for training (e.g., Hydra, logging).finetune.py: Core training loop used byfinetune_ercnn.py.
utils/: Shared utilities:inference_utils.py: Functions for running inference and post-processing model outputs.data_handling.py: Data loaders, splits, and transforms (tiling, normalization).metrics.py: Evaluation metrics (IoU, precision, recall, F1).tiling.py: Utilities to tile large images and reconstruct full-image predictions.
yolo/: YOLO pipeline implementation and assets:augmentations/:AUGMENTATIONS.md: Documentation on augmentation strategies.augmentations.py: Implementation of custom data augmentations.
weights/: Directory for YOLO pretrained weights and checkpoints.
The notebooks/ directory contains pipeline usage and dataset description:
describe_data.ipynb: Data sample descriptions and examples.example_run_pipeline.ipynb: Pipeline usage example.
The results of the evaluation, including predictions and images, will be saved in the results directory.
- NASA Johnson Space Center for their collaboration.
- Rice University D2K Capstone program for their support.
- Weights & Biases for experiment tracking.