marsvision

marsvision is an open-source library for automatically detecting surface features on Mars.

Overview

MarsVision is an open source computer vision package containing tools for training and evaluating Scikit-Learn models and PyTorch CNN models. We implemented a sliding window pipeline, where we can run a sliding window across images for classification of HiRISE images of the surface of Mars.

This package makes use of the PDSC: Planetary Data System Coincidences package to retrieve metadata (e.g. latitudes, longitudes) for classified windows.

See PDSC here

Definitions for Pytorch datasets and models are in the vision directory. Notebooks from our model training and evaluations are in the prototyping directory.

Documentation

The documentation for marsvision is available here.

Installation

Clone the directory and run this command to set up a conda environment and the PDSC package:

./install_env.sh

Training a CNN

This package was written to work with the Deep Mars dataset from this paper, though the pytorch code in the model class can be extended to work with any dataset.

The code to train a model looks like this:

from marsvision.pipeline import Model
from marsvision.vision import ModelDefinitions
alex_model = ModelDefinitions.alexnet_grayscale()
model = Model(alex_model, "pytorch", 
              dataset_root_directory=r"/content/hirise-map-proj"
             )
test_results = model.train_and_test_pytorchcnn()

Where the root directory is the path to the DeepMars dataset. The definitions for the Dataset class and the model class are the vision folder.

A dictionary containing evaluation results indexed by epoch (including loss, accuracy, precisions, and recalls) will be returned by the training method.

Hyparameters such as number of epochs and learning rate are in the config file.

To write the trained CNN and evaluation results to a file, run the following code:

model.save_model("MarsVisionSampleCNN.pt")
model.save_results("MarsVisionSampleResults.p")

Running the sliding window pipeline with a trained model

We wrote a script that runs the sliding window pipeline using a model and writes the results to a database.

The output directory of the cropped windows is specified in the configuration file: marsvision/config.yml.

To run the sample pipeline, run the following script:

bin/run_pipeline_random -n 10 -f MarsVisionSampleCNN.pt -m "pytorch"

The arguments are as follows:

-n: The number of random sample images to process.

-f: Path to the model file. Can either be a pickled SKLearn file (.p) or a Pytorch file (.pt).

-m: Set this to either "pytorch" or "sklearn" depending on the type of model being passed to the script.