DiffICP: Fully-Differentiable ICP in PyTorch

DiffICP is a fast, extensively tested, and fully-differentiable ICP implemented in PyTorch. You can use it to build end-to-end solutions for complex graphics tasks where you want to backpropagate gradients through your ICP.

For theoretical details, check my post on how we used DiffICP for end-to-end 3D reconstruction.

Requirements

Linux or macOS
Python 3.8.10
(Optional) NVIDIA GPU + CUDA 10.2

Required Packages

torch 1.9.1
kornia 0.5.11
pytorch3d 0.3.0

Optional Requirements for `test_icp.py`

tqdm 4.62.3

Optional Requirements for `io_utils.py`

numpy 1.21.2
imageio 2.9.0
pandas 1.3.3
opencv-python 4.5.3.56
pyntcloud 0.1.5

Setup

Installation

Make sure you have all required packages listed above or install with pip install -r requirements.txt.
Run pip install -e . to install the DiffICP code. This will allow you to import DiffICP in your own projects with e.g. from difficp import ICP6DoF.

Full Install:

pip install -r requirements.txt
pip install -e .

Usage

Rigid ICP

If you want to perform rigid ICP, just import our ICP6DoF, initialize an object, and run it on your point clouds or depth maps:

from difficp import ICP6DoF

icp = ICP6DoF(differentiable=True, ...)
rigid_pose = icp(sources, targets, ...)

See arguments of ICP.__init__() and ICP.__call__() in difficp/icp/icp.py for a list of possible hyperparamters/inputs.

Non-rigid ICP

If you need non-rigid ICP, you can build it on your own by subclassing difficp.icp.icp.ICP and overwriting it's abstract methods:

from difficp import ICP

class MyNonRigidICP(ICP):

    def define_num_params(self):
        return ...

    def define_transform_fn(self):
        def my_fn(points, params):
          point_cloud = ...
          return point_cloud
        return my_fn

    def define_solver(self):
        return ...

In define_solver(), you can for instance return a non-rigid LM optimizer that you can build by subclassing difficp.icp.optimizer.LMOptimizer.

As an example, see difficp.icp.icp.ICPSMPL, which performs non-rigid ICP based on a given SMPL-X model and internally uses a non-rigid LM optimizer defined in difficp.icp.optimizer.LMOptimizerSMPL.

Repository Structure

difficp/ is the sources root.
difficp/icp/ contains the differentiable ICP PyTorch implementation, split into several modules.
- difficp/icp/icp.py contains the main ICP logic.
  - difficp.icp.icp.ICP6DoF is the most important class, which performs rigid ICP.
  - difficp.icp.icp.ICP is an abstract class which ICP6DoF inherits from. You can subclass it to build custom non-rigid ICPs.
- difficp/icp/correspondence_function.py contains functions for correspondence finding, weighting, and rejection.
- difficp/icp/distance_function.py contains ICP distance functions (point-to-point, point-to-plane, Symmetric ICP).
- difficp/icp/optimizer.py contains a custom LM optimizer than can be used to minimize the distance functions.
- difficp/icp/procrustes_solver.py contains both an analytical procrustes solver for point-to-point, based on SVD, and a linear solver that can be used to minimize any of the three distance functions (or a combination of them).
- difficp/icp/linear_solver.py contains a differentiable solver for linear equations based on LU-decomposition. This is used internally by both the LM optimizer and the linear solver.
difficp/tests/ contains unit tests.
difficp/utils/ contains all kinds of util functions for:
- 3D geometry (geometry_utils.py),
- Depth map reprojection (depthmap_utils.py),
- Normal calculation (normal_utils.py),
- Saving/Loading point clouds, depth maps, or other data (io_utils.py).

Note for difficp/utils/: Some util functions are used by the ICP, but most are just intended as supplementary code that you can use to load/preprocess your data before you give it to the ICP.

Citation

If you use this code in your research, please cite it as follows:

@unpublished{DiffICP2021,
  title={DiffICP: Fully-Differentiable ICP in PyTorch},
  author={Altenberger, Felix and Niessner, Matthias},
  year={2021},
}

fa9r/DiffICP