Skip to content

🚀 A neural network for monocular spacecraft 6D pose estimation: 3rd place solution of the European Space Agency’s pose estimation challenge

License

Notifications You must be signed in to change notification settings

pedropro/UrsoNet

Repository files navigation

UrsoNet for Monocular Spacecraft 6D Pose estimation

Tensorflow/Keras implementation of UrsoNet, the network used on the European Space Agency's Pose Estimation Challenge that achieved 3th place.

UrsoNet is a simple architecture developed for object pose estimation, which uses Resnet as a backbone. While one branch simply regresses location, the other performs orientation soft classification. Other options, e.g., quaternion regression, are also provided by this implementation.

For more details check our Arxiv preprint.

@article{proenca2019deep,
  title={Deep Learning for Spacecraft Pose Estimation from Photorealistic Rendering},
  author={Proenca, Pedro F and Gao, Yang},
  journal={arXiv preprint arXiv:1907.04298},
  year={2019}
}

Getting started

Data

First, you will need data. This framework is ready to use either our URSO datasets or the SPEED dataset released by Stanford SLAB and ESA.

Just download a dataset and copy it to the directory datasets.

Then you will need the dataset split into train/val/test subsets. This was already done for URSO datasets by retaining 10% for the test set and other 10% for the validation set - notice the CSV files. However you can easily regenerate these random splits by using the script split_dataset.py. Say you want instead 15% each for the soyuz_hard, then run:

python3 split_dataset.py --dataset_dir datasets/soyuz_hard --test_percentage 15 --val_percentage 15

If you are using SPEED then in a Python console, type:

import utils
split_speed('datasets/speed', 10)

This will split 90% of the dataset for training and 10% for the validation set by generating the respective JSON files. N.b. At the time of writing, SPEED test set is not open.

Dependencies

Python>=3.4, TensorFlow=>1.9, Keras>=2.1.6 and other common packages listed in requirements.txt. You can use the following to install all dependencies.

pip3 install -r requirements.txt

Usage and Configuration

While the actual network implementation is in net.py. The main script to train and test UrsoNet is the pose_estimator.py which takes the following general arguments:

Argument Required Type Description
backbone False string Backbones currently supported {'resnet18', 'resnet34', 'resnet50', 'resnet101'}. Default='resnet50'
dataset True string Name of dataset folder. If you are using SPEED set this to 'speed'
weights True string A path to weights .h5 file or 'coco' for coco pre-trained weights or 'last' for last trained model
image_scale False float Scale used to resize input image, default=1.0
square_image False bool Pads input images with zeros to get a square. By default, images are only resized and padded enough to get dimensions multiples of 64
bottleneck False int Bottleneck width, default=32
branch_size False int Size of branch input layers, default=1024
f16 False bool If you are using modern GPUs (RTX), this tells TF/Keras to use half precison: float16. By default this is False.
regress_ori False bool Sets orientation branch to regression. By default this branch does classification
ori_resolution False int Number of bins assigned to each Euler angle, default=16
ori_param False string Regression orientation parameterization: {'quaternion', 'euler_angles', 'angle_axis'}, default='quaternion'
classify_loc False bool Sets location branch to classifcation. By default this branch does regression
regress_keypoints False bool Experimental. Overrides the branch configuration above and sets this to regression of 3 3D keypoints. By default this is false.

Training

The training configuration takes also the following arguments:

Argument Required Type Description
batch_size False int Number of images used per GPU, default=4
epochs False int Number of epochs, default=100
clr False bool Option to use cyclical learning rate, default=False, You need to adjust other clr params in config.py
learn_rate False float Fixed learning rate, in theory this should depend on the batch size, default: 0.001
rot_aug False bool Option to use camera orientation perturbations as augmentation, default=False
rot_image_aug False bool Option to use in-plane image rotation as augmentation, default=False
sim2real False bool Enables the image augmentation pipeline proposed in the paper. This includes converting to grayscale. You can change this in net.py, default=False
ori_weight False float Orientation loss weight, default = 1.0
loc_weight False float Location loss weight, default = 1.0

To train UrsoNet on 'soyuz_easy' with ResNet-50, pre-trained backbone weights from COCO and rotation augmentation use:

python3 pose_estimator.py train --dataset soyuz_easy --weights coco --image_scale 0.5 --ori_resolution 24 --rot_aug --rot_image_aug

This will generate a new checkpoint directory inside logs, where weights will be stored and accumulated along with the configuration parameters. To suppress annoying warnings just add -W ignore to the command above before the script name.

You can stop and resume training whenever you want. To continue training the last trained model from the last checkpoint, with a lower learning_rate simply run:

python3 pose_estimator.py train --dataset soyuz_easy --weights last --image_scale 0.5 --ori_resolution 24 --rot_aug --rot_image_aug --learning_rate 0.0001

Just make sure you are using consistent network configurations.

Inference

To test and visualize results on the test-set using the weights of a specific model inside logs, e.g., 'soyuz_20191001T1207' (feel free to rename it), run

python3 pose_estimator.py test --dataset soyuz_easy --weights soyuz_easy20191001T1207 --image_scale 0.5 --ori_resolution 24

This will show several windows per image, which you will need to close to load the next image.

To evaluate the network on the full test set (val set for 'SPEED'). Use:

python3 pose_estimator.py evaluate --dataset soyuz_easy --weights soyuz_easy20191001T1207 --image_scale 0.5 --ori_resolution 24

Pre-trained weights on URSO and SPEED

Our weights are available here. If you want to use directly our weights, all you have to do is change the argument weights to one of these {'soyuz_hard', 'dragon_hard', 'speed'} and make sure you pass the right arguments like this:

python3 -W ignore  pose_estimator.py test --dataset soyuz_easy --weights soyuz_hard --image_scale 0.5 --ori_resolution 24 --bottleneck 128 --square_image

or this:

python3 -W ignore  pose_estimator.py test --backbone resnet101 --dataset speed --weights speed --image_scale 0.5 --ori_resolution 32 --bottleneck 528 --f16

About

🚀 A neural network for monocular spacecraft 6D pose estimation: 3rd place solution of the European Space Agency’s pose estimation challenge

Topics

Resources

License

Stars

Watchers

Forks

Packages

No packages published