Render Surface Normal Map?

[DaoyiG]

Hi there!

Thanks for this amazing project! I'd like to know whether it's possible to render the normal map of an object at a given pose?

Thanks!

[gkioxari]

Yes it's possible. The output from the rasterization operation is a data structure called Fragments

pytorch3d/pytorch3d/renderer/mesh/rasterizer.py

Lines 16 to 21 in 4281df1

	# Class to store the outputs of mesh rasterization
	class Fragments(NamedTuple):
	pix_to_face: torch.Tensor
	zbuf: torch.Tensor
	bary_coords: torch.Tensor
	dists: torch.Tensor

Fragments contains a variable pix_to_face that maps each (x,y) image pixel to the K := faces_per_pixel faces of the input mesh that are rasterized to that pixel. When rendering RGB images, that information is used to get the color of the point of intersection for each face, this is what shading and blending does. Normal rendering is even simpler. For each image pixel (x,y) you have the corresponding K faces which carry normals. So all you have to do is blend these K normal directions (if K=1 no blending is necessary and only the nearest face is returned).

The need to render many shape attributes (such as texture or normals) is the motivation behind breaking the computation of rendering into rasterization - shading - blending in PyTorch3D. The rasterization step, the hardest step in the pipeline, remains the same. But the shading and blending, which are in pure PyTorch, should be customized. In your case, you want to customize these operators to perform normal rendering instead of RGB rendering. Follow the code to see how RGB rendering is performed, and adjust your code accordingly! Note that this is the intended use of PyTorch3D.

[DaoyiG]

Thanks for your fast reply!

If I understand you correctly, I can get the normal map by customizing a shader and the phong_shading function that returns the normals instead of the color of the K faces by something like the following:

def phong_normal_shading(self, meshes, fragments) -> torch.Tensor:
        faces = meshes.faces_packed()  # (F, 3)
        vertex_normals = meshes.verts_normals_packed()  # (V, 3)
        faces_normals = vertex_normals[faces]
      
        pixel_normals = interpolate_face_attributes(
            fragments.pix_to_face, fragments.bary_coords, faces_normals
        )
       
        return pixel_normals

The pixel_normals returned by the phong_normal_shading function will have the shape of (batch size * height * width * K * 3), then it only needs to be sent into softmax_rgb_blend function to get the blended normals (just as the blended RGB).

Is my above description correct? Please correct me if anything wrong : )

Thanks!

[gkioxari]

Very close!! Note that in phong_shading the interpolate_face_attributes will multiply the face_attributes (in your case the normals) with the barycentric coordinates.

pytorch3d/pytorch3d/ops/interp_face_attrs.py

Line 99 in 9585a58

pixel_vals = (barycentric_coords[..., None] * pixel_face_vals).sum(dim=-2)

For RGB, that makes sense. But for normals, not so much. So your version of interpolation should just be a simper face interpolation which just averages the normals across the K faces (equivalent to passing ones instead of fragments.bary_coords.

[DaoyiG]

Thanks for your reply!

Just want to be super sure here, if I customize the HardPhongShader so only the normal of the closest face should be returned, then no further blending is needed, and the following code should return a tensor of shape batchsize * height * width * 1 * 3, which we could just squeeze the penultimate dimension, and the resulting tensor of shape batchsize * height * width * 3 should already be the correct normal map?

    def phong_normal_shading(self, meshes, fragments) -> torch.Tensor:
        faces = meshes.faces_packed()  # (F, 3)
        vertex_normals = meshes.verts_normals_packed()  # (V, 3)
        faces_normals = vertex_normals[faces]
        ones = torch.ones_like(fragments.bary_coords)
        pixel_normals = interpolate_face_attributes(
            fragments.pix_to_face, ones, faces_normals
        )
        return pixel_normals

And if I customize the SoftPhongShader which 'K > 1', then normals should be blended. So after the above phong_normal_shading function, shall I just call the softmax_rgb_blend() function to blend like the following?

normal_map = softmax_rgb_blend(
            pixel_normals, fragments, blend_params, znear=znear, zfar=zfar
        )

And I also wonder what is the advantage of the 'soft' shader over the 'hard' shader, will the 'soft' shader brings more stable gradients?

Thanks!

[gkioxari]

That's right I think! If you set K=1 then no blending is needed. For K>1, then you need to blend and here is where you want to make some decisions based on your needs. You might want to blend by taking the average across K, or you might want to blend by taking a weighted average across K (for example, weighing disproportionate to how far the face is from the camera -- this info is in fragments.dist). The softmax_rgb_blend does something along those lines for the latter case. If you just want to average, then you can directly call mean on pixel_normals as returned from the interpolate_face_attributes.

The advantage of soft blending is exactly what you mention. Just taking the nearest face often leads to noisy Fragments especially when you learn and shapes evolve over time. So you want to propagate information to the K nearest faces instead of just the single nearest one.

[DaoyiG]

Thanks a lot! Now my problem is solved and I'm closing this issue : )

[WenjiaWang0312]

Dear @DaoyiG and @gkioxari , thank you for your sharing!
However, I found in shading that the code pass fragments.bary_coords to compute the pixel normals, so why do you suggest that "passing ones instead of fragments.bary_coords."?

[yiyuzhuang]

Dear @DaoyiG and @gkioxari , thank you for your sharing! However, I found in shading that the code pass fragments.bary_coords to compute the pixel normals, so why do you suggest that "passing ones instead of fragments.bary_coords."?

Hi dear @JasonBoy1 ,
I think the reason is the normal of all the points on a face should be the same, which is different with the utilization of color.

[zc-alexfan]

Custom shader for normal rendering in camera space:

class NormalShader(ShaderBase):
    def __init__(self, cameras):
        super().__init__()
        self.cameras = cameras

    def forward(self, fragments, meshes, **kwargs):
        cameras = self.cameras

        faces = meshes.faces_packed()
        vertex_normals = meshes.verts_normals_packed()
        face_normals = vertex_normals[faces]

        pix_to_face = fragments.pix_to_face
        bary_coords = fragments.bary_coords
        face_normals_at_pixels = face_normals[pix_to_face]  # (N, H, W, K, 3, 3)

        interpolated_normals = (bary_coords.unsqueeze(-1) * face_normals_at_pixels).sum(dim=-2)
        interpolated_normals = torch.nn.functional.normalize(interpolated_normals, dim=-1)

        # === Transform normals to camera space
        R = cameras.R
        R = R[:, None, None, None, :, :]
        normals_cam = torch.matmul(R, interpolated_normals.unsqueeze(-1)).squeeze(-1)

        # === Fixed rotation to make normals face +Z
        theta = torch.tensor(3.1415926, device=normals_cam.device)
        cos_t = torch.cos(theta)
        sin_t = torch.sin(theta)
        R_align = torch.tensor([
            [cos_t, 0, sin_t],
            [0,     1, 0],
            [-sin_t, 0, cos_t]
        ], device=normals_cam.device).T
        normals_cam = torch.matmul(normals_cam, R_align)

        # RGB
        normal_rgb = (normals_cam + 1.0) / 2.0  # (N, H, W, K, 3)

        # Alpha: 1.0 if face index >= 0, else 0.0
        alpha = (pix_to_face >= 0).float().unsqueeze(-1)  # (N, H, W, K, 1)

        rgba = torch.cat([normal_rgb, alpha], dim=-1)  # (N, H, W, K, 4)
        return rgba

Also created a gist for example usage here. Result on my example script:

Enjoy