| import torch | |
| from pytorch3d.renderer.mesh.shader import ShaderBase | |
| from pytorch3d.renderer import ( | |
| SoftPhongShader, | |
| ) | |
| from pytorch3d.renderer import BlendParams | |
| class MultiOutputShader(ShaderBase): | |
| def __init__(self, device, cameras, lights, materials, ccm_scale=1.0, choices=None): | |
| super().__init__() | |
| self.device = device | |
| self.cameras = cameras | |
| self.lights = lights | |
| self.materials = materials | |
| self.ccm_scale = ccm_scale | |
| if choices is None: | |
| self.choices = ["rgb", "mask", "depth", "normal", "albedo", "ccm"] | |
| else: | |
| self.choices = choices | |
| blend_params = BlendParams(sigma=1e-4, gamma=1e-4) | |
| self.phong_shader = SoftPhongShader( | |
| device=self.device, | |
| cameras=self.cameras, | |
| lights=self.lights, | |
| materials=self.materials, | |
| blend_params=blend_params | |
| ) | |
| def forward(self, fragments, meshes, **kwargs): | |
| batch_size, H, W, _ = fragments.zbuf.shape | |
| output = {} | |
| if "rgb" in self.choices: | |
| rgb_images = self.phong_shader(fragments, meshes, **kwargs) | |
| rgb = rgb_images[..., :3] | |
| output["rgb"] = rgb | |
| if "mask" in self.choices: | |
| alpha = rgb_images[..., 3:4] | |
| mask = (alpha > 0).float() | |
| output["mask"] = mask | |
| if "albedo" in self.choices: | |
| albedo = meshes.sample_textures(fragments) | |
| output["albedo"] = albedo[..., 0, :] | |
| if "depth" in self.choices: | |
| depth = fragments.zbuf | |
| output["depth"] = depth | |
| if "normal" in self.choices: | |
| pix_to_face = fragments.pix_to_face[..., 0] | |
| bary_coords = fragments.bary_coords[..., 0, :] | |
| valid_mask = pix_to_face >= 0 | |
| face_indices = pix_to_face[valid_mask] | |
| faces_packed = meshes.faces_packed() | |
| normals_packed = meshes.verts_normals_packed() | |
| face_vertex_normals = normals_packed[faces_packed[face_indices]] | |
| bary = bary_coords.view(-1, 3)[valid_mask.view(-1)] | |
| interpolated_normals = ( | |
| bary[..., 0:1] * face_vertex_normals[:, 0, :] + | |
| bary[..., 1:2] * face_vertex_normals[:, 1, :] + | |
| bary[..., 2:3] * face_vertex_normals[:, 2, :] | |
| ) | |
| interpolated_normals = interpolated_normals / interpolated_normals.norm(dim=-1, keepdim=True) | |
| normal = torch.zeros(batch_size, H, W, 3, device=self.device) | |
| normal[valid_mask] = interpolated_normals | |
| output["normal"] = normal | |
| if "ccm" in self.choices: | |
| face_vertices = meshes.verts_packed()[meshes.faces_packed()] | |
| faces_at_pixels = face_vertices[fragments.pix_to_face] | |
| ccm = torch.sum(fragments.bary_coords.unsqueeze(-1) * faces_at_pixels, dim=-2) | |
| ccm = (ccm[..., 0, :] * self.ccm_scale + 1) / 2 | |
| output["ccm"] = ccm | |
| return output |