How to generate normal map as ControlNet conditioning

[buaacyw]

Hi, I'm trying to use the normal map rendered by NeRF as the condition of ControlNet like the below image: To get the normal map, we need to transform comp_normal from world space to camera space with c2w.

                comp_normal = nerfacc.accumulate_along_rays(
                    weights[..., 0],
                    values=geo_out["normal"],
                    ray_indices=ray_indices,
                    n_rays=n_rays,
                )
                comp_normal = F.normalize(comp_normal, dim=-1) # (-1, 1)
                batched_normal = comp_normal.view(batch_size, -1,3)
                w2c = kwargs['c2w'][:, :3, :3].inverse()
                cam_normal = torch.matmul(w2c, batched_normal.transpose(1,2)).transpose(1,2)
                cam_normal = cam_normal * opacity.view(batch_size, -1, 1)
                cam_normal = cam_normal[0] # only visualize one image
                cam_normal[(cam_normal.sum(dim=1) == 0), 2] = 1
                cam_normal = torch.nn.functional.normalize(cam_normal,dim=1) # (-1,1)
                cam_normal_image = ((cam_normal+1.0)/2.0).clip(0, 1)

However, the normal map I get by this way looks different from ControlNet's:

https://github.com/threestudio-project/threestudio/assets/52091468/9fa89e05-5d43-43d0-baaf-683f05b66a23

Could you please help me? And I suggest adding this as a new feature. Thanks!

[buaacyw]

This is the result of ProlificDreamer in the Nerf phase, I choose 'analytic' as a normal map and the result looks bad. I have also tried other normal types but can't get a nice result. So is it possible to get a better normal map in the NeRF phase?

[yankeesong]

If you are using NormalBae to detect the normal map for ControlNet, one thing I noticed is the x-axis for NormalBae is flipped. Try to revert the x-axis for your comp_normal (before normalizing to [0,1]) and see if this solves the problem.