This repository contains code for training and evaluating CNN-based local feature extractors targeted for optical navigation purposes in the vicinity of Solar System small bodies. In addition, the code base supports hyperparameter optimization and fetching and preprocessing image data from original sources. Images from several missions are supported, including Rosetta, OSIRIS-REx, Hayabusa, and NEAR Shoemaker.
The work is described in detail in the following article:
@article{
title={{CNN}-based local features for navigation near an asteroid},
author={Knuuttila, Olli and Kestil{\"a}, Antti and Kallio, Esa},
journal={arXiv preprint},
year={2023},
doi={10.48550/arXiv.2309.11156}
}
Related data and fully trained feature extractors can be found in Zenodo: https://dx.doi.org/10.5281/zenodo.8319117.
The system works best on Linux, but can be run on Windows as well, except for hyperparameter optimization.
The following instructions assume that you are using Linux.
Setup environment using conda:
conda create -n navex -c pytorch -c nvidia -c conda-forge --override-channels python==3.10 pip pytorch torchvision \
pytorch-cuda==11.6 opencv pytorch-lightning==1.1.7 quaternion scipy paramiko scikit-optimize pandas numba
conda activate navex
pip install pytorch-lightning-bolts
pip install git+https://github.com/jatentaki/unets
pip install git+https://github.com/jatentaki/torch-localize
pip install git+https://github.com/jatentaki/torch-dimcheck
# to enable hyperparameter optimization:
pip install ray==2.2.0 ray[default]
# additionally, modify functions node_ip_address_from_perspective and get_node_ip_address to always return "127.0.0.1":
vim <env-path>/lib/python3.10/site-packages/ray/_private/services.py
# to enable testing the original R2D2 network
pip install git+https://github.com/oknuutti/r2d2.git
# to enable relative pose estimation during feature extractor evaluation:
pip install git+https://github.com/oknuutti/featstat.git
# to enable experimental MAC operation counting
pip install thop
The rest of the instructions assume that you are in the same folder as this README.md
file and that there are
additional data
, models
and results
folders. In addition, hyperparameter optimization uses output
and cache
folders. The data
folder is thought to contain the datasets (.tar
files) downloaded from Zenodo and the models
folder is thought to contain the trained models (.ckpt
files), also downloaded from Zenodo. Please note that the
file paths in e.g. rot-cg67p-osinac.tar
and cg67p-osinac-d.tar
archives are the same, so you need to either rename
the extracted folder, extract them to different folders, or only extract the archive that you need.
E.g., assuming data/hpatches/image_list_hpatches_sequences.txt
contains
a list of
hpatches images, the following command can be used to extract
features from the images and save them to .lafe
files:
python -m navex.extract --images data/hpatches/image_list_hpatches_sequences.txt \
--model models/lafe.ckpt --tag lafe --top-k 2000
Alternatively, to extract features of all images in a folder (incl. subfolders), run:
python -m navex.extract --images data/cg67p/navcam \
--model models/lafe.ckpt --tag lafe --top-k 2000
The example above extracts LAFE features of all images belonging to the CG67P Navcam dataset (file
cg67p-navcam.tar
in Zenodo).
The meaning of the command line arguments can be found by running python -m navex.extract --help
.
The following command can be used to extract and evaluate LAFE features on the Eros MSI, Itokawa AMICA,
C-G/67P OSINAC and synthetic datasets. Note that the evaluation code needs the contents of the *-d.tar
archives
instead of the rot-*.tar
archives.
python -m navex.evaluate --root=./data --model=./models/extractor-lafe.ckpt \
--output=./results/lafe-eval-results.csv \
-d=eros -d=ito -d=67p -d=synth \
--gpu=1 --min-size=256 --max-size=512 --success-px-limit=5 \
--det-lim=0.5 --qlt-lim=0.5 --max-scale=1.0 --est-ori
The results are saved to ./results/lafe-eval-results.csv
. If --est-ori
is given, the
featstat
package is needed, see section "Getting started" above. The --model
argument also accepts akaze
, orb
,
sift
, rsift
(for RootSIFT), and surf
as values. Also .onnx
models are accepted. However, in that case,
--min-size
and --max-size
has to be equal as different models for different scales is not supported.
Giving --help
as an argument prints out a list of possible arguments and their meanings.
Converting a .ckpt
model into an ONNX model can be done with the following commands:
import torch
from navex.dataset.tools import load_model
path = "models/extractor-lafe.ckpt"
model = load_model(path, 'cpu')
torch.onnx.export(model, torch.randn(1, 1, 512, 512), path[:-5] + ".512x512.onnx",
input_names=['input'], output_names=['des', 'det', 'qlt'])
To produce the plots in the article based on the evaluation output, run the following command varying the dataset
argument to be each of 67p
, eros
, ito
, and synth
:
python -m navex.visualizations.evaluation --path=./results/lafe-eval-results.csv --dataset=67p
To train a model from scratch, copy a suitable config file from navex/experiments
to a new file, e.g. myconfig.yaml
,
and modify the parameters as needed. The config file format is defined in navex/experiments/definition.yaml
.
The config file ast_lafe_2.yaml
is the one used to train the LAFE-ME model in the article.
The config file ast_r2d2_v2_best.yaml
corresponds to the HAFE-R2D2 model in the article.
For example, to retrain the LAFE-ME model, you could use the following command:
python -m navex.train --config=./navex/experiments/ast_lafe_2.yaml -d=./data --cache=./cache -o=./output \
--loss--teacher=./models/extractor-hafe-r2d2.ckpt \
--training--batch-size=32 --data--image-size=448 \
--id=my-v1 -j=6
The meaning of the command line arguments can be found by running python -m navex.train --help
. Most of the arguments
are the same as in the config file, but can be overridden from the command line. However, due to the experimental nature of the
code, the help text may not be up-to-date and inspection of the source code might be necessary to understand the meaning
of some arguments. The arguments starting with --search
are only used for hyperparameter optimization, see below.
Training expects the contents of the rot-*.tar
archives. The contents of the *-d.tar
archives are only used when
evaluating fully trained feature extractors.
Hyperparameter optimization is done using Ray Tune customized for the HPC environments that we used: Triton of Aalto Science-IT project, and Puhti of CSC - IT Center for Science, Finland.
This functionality is very experimental and will likely need a lot of tweaking to get it to work in your environment.
To use the hyperparameter optimization functionality, you need to have access to a cluster with a SLURM job scheduler.
You also need to create a new worker .sbatch
file with your settings at navex/ray
and modify the method
raytune.ScheduledWorkerNode.schedule_slurm_node
to use your custom .sbatch
file. The current worker-csc.sbatch
and worker-triton.sbatch
files are for Puhti and Triton, respectively.
After managing to set up everything, you can run the hyperparameter optimization that produced the HAFE-R2D2 model
with the following command, probably embedded within a custom .sbatch
file and scheduled on a long-running,
non-GPU node:
srun python -m navex.raytune --config=./navex/experiments/ast_r2d2_ashabo_2.yaml \
--out=./output --cache=./cache --host=`hostname` \
-b=8 --epochs=24001 --bm=32000 --tf=1500 \
--data--image-size=224 --lr=1e-3 -j=6 --nodes=6
The command line arguments are the same as for navex.train
. The --nodes=6
argument specifies that
navex.raytune
should attempt to schedule work on 6 GPU nodes.
Tensorboard logs can be viewed during and after optimization by running
tensorboard --logdir=./output/ast_r2d2_ashabo_2 --port=16006 > tensorboard.log 2>&1 &
and then opening localhost:16006
in a browser.
To plot partial dependency results of a hyperparameter search, you can use the following command:
python -m navex.visualizations.raytune --path=./output/ast_r2d2_ashabo_2
The data used in the article can be downloaded from Zenodo. However,
if you wish to fetch the data from the original sources, you can use scripts in the navex/datasets/preproc
folder.
Each script has a --help
argument that explains its usage. The --start
and --end
arguments can be used to
limit the amount of data to be fetched (the range is from 0.0 to 1.0). The image indices are randomized once
and persisted so that the same indices are used for consequtive runs, provided that the .sqlite
database is preserved.
However, before you begin, create a separate conda environment for it. You can do it e.g. like this:
conda create -n data_io -c conda-forge "python==3.8" pip opencv matplotlib gdal geos \
tqdm scipy pvl quaternion requests bs4 spiceypy
conda activate data_io
pip install pds4-tools
To fetch the Bennu TAGCAMS dataset (100% of suitable images):
python -m navex.datasets.preproc.bennu_tagcams --dst=./data/bennu/tagcams \
--start=0.0 --end=1.0
For C-G/67P NAVCAM (50% of suitable images):
python -m navex.datasets.preproc.cg67p_rosetta --dst=./data/cg67p/navcam --instr=navcam \
--start=0.0 --end=0.5
For C-G/67P OSINAC (20% of suitable images):
python -m navex.datasets.preproc.cg67p_rosetta --dst=./data/cg67p/osinac --instr=osinac \
--min-matches=90000 --min-angle=10 --max-angle=30 \
--start=0.0 --end=0.2
For Eros MSI (10% of suitable images):
python -m navex.datasets.preproc.eros_msi --dst=./data/eros --pairs-only \
--min-matches=90000 --min-angle=10 --max-angle=30 \
--start=0.0 --end=0.1
For Itokawa AMICA (100% of suitable images):
python -m navex.datasets.preproc.itokawa_amica --dst=./data/itokawa \
--min-matches=10000 --min-angle=10 --max-angle=30 \
--start=0.0 --end=1.0
The datasets obtained using the above commands contain apart from PNG images, image pair pixel correspondences
in the aflow
folder, detailed metadata *.lbl
, pixel depths *.d
, fixed frame 3D coordinates *.xyz
,
and a shadow mask *.sdw
. The database dataset_all.sqlite
contains structured metadata of the images, including poses
and light direction. However, these values were not used for the article as they were not reliable
enough. Instead, all relative poses are calculated based on the *.xyz
files, derived from the available geometry
backplanes. Better values should be obtained from e.g. the mission SPICE kernels. Apart from the file names of
the pixel correspondences, pairing is contained in the pairs.txt
file.
The C-G/67P NAVCAM and Bennu TAGCAMS datasets do not contain pairing, pixel depths or 3D coordinates, as geometry backplanes were not available at the time when the data preprocessing related to the article was done.
Feature extractor training uses only the image files and the pixel correspondences at the aflow
folder. However,
to conserve network capacity, the images of geo-referenced datasets are rotated so that the target body rotation axis
points upwards in each image (resulting in the rot-*.tar
archives). This can be done for all geo-referenced datasets
used by an experiment by issuing the following command:
python -m navex.train --config=./navex/experiments/ast_lafe_2.yaml \
-d=./data --preproc-path=./data/rotated
The depth information in .d
files are used during evaluation if the --est-ori
argument is given. In addition,
evaluation expects non-rotated images as input. Thus, the need for both rot-*.tar
and *-d.tar
archives.