Optimizing parametric hand model and object mesh positions and orientations using ID masks

[friedOnions]

I have a scene containing a parametric hand model and a rigid object. I want to optimize the hand model parameters and object position and orientation using ground truth images and masks of the hand and object.

I first tried using the SoftSilhouetteShader to render a single silhouette for both objects and compare this to the combined ground truth masks. This approach works really well unless the objects overlap. Then, the hand model parameters do all kinds of weird stuff because there is basically no supervision anymore.

Thus, I wanted to try to render individual masks. For each of the two objects, I assigned different vertex colors, i.e (1, 0, 0) for one and (0, 1, 0) for the other objects, and rendered these using a custom shader taking advantage of the softmax_rgb_blend which is also used in the SoftPhongShader. I combined the ground truth masks using the same color coding, and because I don't really need any shading on the objects, I created the custom shader to save resources:

class IDShader(nn.Module):
    def __init__(self, blend_params=None):
        super().__init__()
        self.blend_params = blend_params if blend_params is not None else BlendParams()

    def forward(self, fragments, meshes, **kwargs):
        blend_params = kwargs.get("blend_params", self.blend_params)
        texels = meshes.sample_textures(fragments)
        colors = softmax_rgb_blend(texels, fragments, blend_params)  

        return colors

For some reason however, gradients are not stable and the optimization does not converge at all, even if adding the SoftShilhouette rendering and a respective loss to the pipeline. Does anyone have a suggestion why this might be and what I could do to solve the problem?