История

Jonathan Leban 2ad64cc505 removed useless line		2022-04-04 01:35:46 -07:00
..
documentation	Updating docs again	2021-09-23 14:27:10 -07:00
pose_estimation	removed useless line	2022-04-04 01:35:46 -07:00
tests	change the name of the folder	2021-09-10 14:05:14 -07:00
Dockerfile	change the name of the folder	2021-09-10 14:05:14 -07:00
README.md	change the name of the folder	2021-09-10 14:05:14 -07:00
__init__.py	change the name of the folder	2021-09-10 14:05:14 -07:00
base_dockerfile	change the name of the folder	2021-09-10 14:05:14 -07:00
config.yaml.sample	Docs updates	2021-09-23 00:14:51 -07:00
environment-gpu.yml	change the name of the folder	2021-09-10 14:05:14 -07:00
environment.yml	change the name of the folder	2021-09-10 14:05:14 -07:00
setup.py	change the name of the folder	2021-09-10 14:05:14 -07:00

README.md

Drone Pose Estimation Model

This section contains code for training and evaluating a deep neural network to predict the pose of a single object from RGB images. We provide support for running both locally and with Docker.

This model is a modified implementation of Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World, by Tobin et. al. It is based on the classic VGG-16 backbone architecture, and initialized with weights pre-trained on the ImageNet dataset. The head of the network is replaced with a 3D position prediction head that outputs (x, y, z), and an orientation predicton head that outputs a quaternion (q_x, q_y, q_z, q_w).

Quick Start (Recommended)

We trained this model on sythetic data collected in Unity. To learn how to collect this data and train the model yourself, see our data collection and training tutorial.

Pre-Trained Model

Setup

CLI

This model supports a train and an evaluate command. Both of these have many arguments, which you can examine in cli.py. They will default to the values in config.yaml for convenience, but can be overridden via the command line.

The most important train arguments to be aware of are:

--data_root: Path to the directory containing your data folders. These directory should include UR3_single_cube_training and UR3_single_cube_validation, containing the training and validation data, respectively.
--log-dir-system: Path to directory where you'd like to save Tensorboard log files and model checkpoint files.

The most important evaluate arguments to be aware of are:

--load-dir-checkpoint: Path to model to be evaluated.
--data_root: Path to the directory containing your data folders. These directory should include UR3_single_cube_training and UR3_single_cube_validation, containing the training and validation data, respectively.

Performance

Unit Testing

We use pytest to run tests located under tests/. You can run the tests after following the setup instructions in Running on Local with Conda commands.

You can run the entire test suite with:

python -m pytest

or run individual test files with:

python -m pytest tests/test_average_translation_mean_square_error.py

Resources

Documentation describing structure of the model code
Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World
Pose Estimation Tutorial