SDFEst is a package for pose, scale, and shape estimation using discretized signed distance fields. It includes and combines three main components (generative model, discriminative initialization, and differentiable renderer) to enable pose and shape estimation in an analysis-by-synthesis framework.
Installation | Dataset Preparation | Paper Reproduction | Code Structure | Development | Citation | Docs
You can install everything with pip (either in your global Python environment, or a virtual environment).
pip install sdfest
In principle, this allows to train novel networks, apply the pipeline, use the individual modules, and reproduce our experiments.
git clone [email protected]:roym899/sdfest.git
cd sdfest
pip install -e .
This will get the latest version of all dependencies, which might break if a breaking change was introduced in a dependency. To reproduce our tested environment you can use the provided requirements.txt.
Execute the following command to install the packages and their dependencies. Note that this will likely downgrade / upgrade versions in your current environment, so it's better to use this in a virtual environment only.
git clone [email protected]:roym899/sdfest.git
cd sdfest
pip install -r requirements.txt -e .
You need to install Detectron2 manually to run the pipeline with automatic instance segmentation.
Follow the detectron2 installation guide from there. Tested with detectron2 0.6 + torch 1.12.0 (and various older versions).
See below for expected folder structure for each dataset.
ShapeNet (Website)
./data/shapenet/02876657/...
./data/shapenet/02880940/...
./data/shapenet/03797390/...
ModelNet (Website)
./data/modelnet/bottle/...
./data/modelnet/bowl/...
./data/modelnet/cup/...
Redwood (Website)
./data/redwood/bottle/rgbd/...
./data/redwood/bowl/rgbd/...
./data/redwood/mug/rgbd/...
RGB-D Object Dataset (Website)
./data/rgbd_objects_uw/{bowl,coffee_mug,shampoo,water_bottle}/{bowl,coffee_mug,shampoo,water_bottle}_1/
./data/rgbd_objects_uw/{bowl,coffee_mug,shampoo,water_bottle}/{bowl,cofee_mug,shampoo,water_bottle}_2/
...
First, make sure the datasets are in the right format.
See the bash script in reproduction_scripts/reproduce_paper.sh
.
Evaluation code for REAL275 and REDWOOD75 experiments will be integrated in cpas_toolbox soon.
Non-cleaned up version of evaluation code can be found in icaps_eval
branch.
To train a network for a specific category you need to first train a per-category VAE, and afterwards an initialization network.
First we need to convert the ShapeNet meshes to SDFs and optionally filter the dataset. To reproduce the preprocessing of the paper run
source preprocess_shapenet.sh
Then run
source train_vaes.sh
to train the models using the same configuration as used for the paper.
To train the initialization network we used in our paper, run
source train_init_networks.sh
If you want to train the initialization network based on a previously trained object model, you need to create a new config linking to the newly trained VAE.
See, for example, sdfest/initialization/configs/discretized_mug.yaml
, which links to sdfest/initialization/vae_models/mug.yaml
).
Code is structured into 4 sub-packages:
- vae: variational auto-encoder for shapes
- initialization: pose and shape estimation from partial pointset
- differentiable_renderer: differentiable renderer for discretized SDF
- estimation: integration of VAE, differentiable renderer and single shot network into initialization + render and compare pipeline
Differentiable rendering of depth image for signed distance fields.
The signed distance field is assumed to be voxelized and its pose is given by a x, y, z in the camera frame, a quaternion describing its orientation and a scale parameter describing its size. This module provides the derivative with respect to the signed distance values, and the full pose description (position, orientation, scale).
General workflow for PyTorch extensions (only tested for JIT, probably similar otherwise)
If you develop PyTorch extensions and want to get correct code checking with ccls / etc. you can do so by going to the ninja build directory (normally home_directory/.cache/torch_extensions/sdf_renderer_cpp
, or set load(..., verbose=True)
in sdf_renderer.py
and check the output), running
ninja -t compdb > compile_commands.json
and moving compile_commands.json
to the projects root directory.
Architecture to learn a low-dimensional representation of signed-distance fields (i.e., an explicit voxel representation of a signed distance function)
You need to preprocess the mesh data prior to running the training script, like this:
python -m sdfest.vae.scripts.process_shapenet --inpath shapenet_subfolder --outpath output_path --resolution 64 --padding 2
You can control the resolution and added padding so that there is some empty space left in the signed distance field on all borders. If you are running this script via ssh you need to run export PYOPENGL_PLATFORM=egl
prior to running the script and use the --all
option which will disable any filtering. Otherwise mesh selection will proceed in two steps: first you see one mesh after another and need to decide which to keep. Pressing left will remove a mesh, pressing right will keep it. After a decision has been made, the conversion will run. Finally another manual selection process is started, where you can remove SDFs in which the mesh to SDF conversion has failed.
To train the network you can now use either python -m sdfest.vae.scripts.train
or python sdfest/vae/scripts/train.py
.
Configuration can be provided through command-line arguments and hierarchical yaml files. The config files are read in depth first order and later specifications will overwrite previous specifications. To summarize:
- command line parameters will take precedence over all config files
- when specifying multiple config files, the second config file will overwrite params set by the first config file
- parent config files overwrite the params set in the contained config files
Architectures for single-shot SDF shape and pose estimation from a single (in the future possibly also multiple) depth views.
To train a new model run
python -m sdfest.initialization.scripts.train --config CONFIG_FILE
See the sdfest/initialization/configs
or reproduction_scripts
folder.
Modular architecture and experiments for SDF shape and pose estimation.
Combines initialization, VAE, and differentiable renderer for full analysis-by-synthesis.
- Use
pip install -e .
to install the package in editable mode - Use
pip install -r requirements-dev.txt
to install dev tools - Use
pytest --cov=sdfest --cov-report term-missing tests/
to run tests and check code coverage
If you find this library useful in your research, consider citing our publication:
@article{bruns2022sdfest,
title={{SDFEst}: Categorical Pose and Shape Estimation of Objects From {RGB-D} Using Signed Distance Fields},
author={Bruns, Leonard and Jensfelt, Patric},
journal={IEEE Robotics and Automation Letters},
volume={7},
number={4},
pages={9597--9604},
year={2022}
}