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}
}

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.

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

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:

The training configuration takes also the following arguments:

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.

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

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