SoftRasterizer enhancement to learn 3D from RGB only

[monniert]

🚀 Feature

Add the layered blending function of https://arxiv.org/abs/2204.10310 to make SoftRasterizer work using RGB loss only

Motivation

SoftRasterizer does not work without silhouettes (see #359, #507, #839, #840, #1004). The work (appendix A) analyses why (TLDR: the original softmax-like blending function annihilates the differentiability wrt the opacity maps - and thus wrt the vertex positions - because opacity maps appear in both the numerator and denominator of the softmax) and presents a simple modification to SoftRas based on a new layered aggregation function. The resulting Layered SoftRasterizer can successfully learn from RGB loss only.

Pitch

A version of this layered blending function is implemented here. Similar to the original softmax_rgb_blend, this function layered_rgb_blend could be placed in pytorch3d/renderer/blending.py and handled by the different shaders given a dedicated argument.

Using the resulting shader in the fit_textured_mesh.ipynb tutorial enables to learn 3D meshes using RGB only. This zip file contains the ipynb + html versions of the tutorial showing that:

1. SoftRas with silhouette works

2. SoftRas with silhouette + RGB works

3. SoftRas with RGB diverges

4. Layered SoftRas with RGB works

[ykzzyk]

Hello @monniert , thank you for the amazing work! Can u elaborate more on how to layered_rgb_blend function? I am currently trying to perform camera pose estimation with only RGB channels, and I noticed they cannot converge.

[monniert]

Hi @ykzzyk, sure you should have a look at this file for an example on how to use this layered_rgb_blend function.

Specifically, you should redefine a pytorch3d Shader class (called LayeredShader) and reuse it inside the classical MeshRenderer class. I also noticed it do not converge for pose estimation, and in my case this fix makes the optimization for pose estimation work