Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import pytorch3d | |
| import torch.nn.functional as F | |
| from pytorch3d.ops import interpolate_face_attributes | |
| from pytorch3d.renderer import ( | |
| look_at_view_transform, | |
| FoVPerspectiveCameras, | |
| AmbientLights, | |
| PointLights, | |
| DirectionalLights, | |
| Materials, | |
| RasterizationSettings, | |
| MeshRenderer, | |
| MeshRasterizer, | |
| SoftPhongShader, | |
| SoftSilhouetteShader, | |
| HardPhongShader, | |
| TexturesVertex, | |
| TexturesUV, | |
| Materials, | |
| ) | |
| from pytorch3d.renderer.blending import BlendParams, hard_rgb_blend | |
| from pytorch3d.renderer.utils import convert_to_tensors_and_broadcast, TensorProperties | |
| from pytorch3d.renderer.mesh.shader import ShaderBase | |
| def get_cos_angle(points, normals, camera_position): | |
| """ | |
| calculate cosine similarity between view->surface and surface normal. | |
| """ | |
| if points.shape != normals.shape: | |
| msg = "Expected points and normals to have the same shape: got %r, %r" | |
| raise ValueError(msg % (points.shape, normals.shape)) | |
| # Ensure all inputs have same batch dimension as points | |
| matched_tensors = convert_to_tensors_and_broadcast( | |
| points, camera_position, device=points.device | |
| ) | |
| _, camera_position = matched_tensors | |
| # Reshape direction and color so they have all the arbitrary intermediate | |
| # dimensions as points. Assume first dim = batch dim and last dim = 3. | |
| points_dims = points.shape[1:-1] | |
| expand_dims = (-1,) + (1,) * len(points_dims) | |
| if camera_position.shape != normals.shape: | |
| camera_position = camera_position.view(expand_dims + (3,)) | |
| normals = F.normalize(normals, p=2, dim=-1, eps=1e-6) | |
| # Calculate the cosine value. | |
| view_direction = camera_position - points | |
| view_direction = F.normalize(view_direction, p=2, dim=-1, eps=1e-6) | |
| cos_angle = torch.sum(view_direction * normals, dim=-1, keepdim=True) | |
| cos_angle = cos_angle.clamp(0, 1) | |
| # Cosine of the angle between the reflected light ray and the viewer | |
| return cos_angle | |
| def _geometry_shading_with_pixels( | |
| meshes, fragments, lights, cameras, materials, texels | |
| ): | |
| """ | |
| Render pixel space vertex position, normal(world), depth, and cos angle | |
| Args: | |
| meshes: Batch of meshes | |
| fragments: Fragments named tuple with the outputs of rasterization | |
| lights: Lights class containing a batch of lights | |
| cameras: Cameras class containing a batch of cameras | |
| materials: Materials class containing a batch of material properties | |
| texels: texture per pixel of shape (N, H, W, K, 3) | |
| Returns: | |
| colors: (N, H, W, K, 3) | |
| pixel_coords: (N, H, W, K, 3), camera coordinates of each intersection. | |
| """ | |
| verts = meshes.verts_packed() # (V, 3) | |
| faces = meshes.faces_packed() # (F, 3) | |
| vertex_normals = meshes.verts_normals_packed() # (V, 3) | |
| faces_verts = verts[faces] | |
| faces_normals = vertex_normals[faces] | |
| pixel_coords_in_camera = interpolate_face_attributes( | |
| fragments.pix_to_face, fragments.bary_coords, faces_verts | |
| ) | |
| pixel_normals = interpolate_face_attributes( | |
| fragments.pix_to_face, fragments.bary_coords, faces_normals | |
| ) | |
| cos_angles = get_cos_angle( | |
| pixel_coords_in_camera, pixel_normals, cameras.get_camera_center() | |
| ) | |
| return pixel_coords_in_camera, pixel_normals, fragments.zbuf[..., None], cos_angles | |
| class HardGeometryShader(ShaderBase): | |
| """ | |
| renders common geometric informations. | |
| """ | |
| def forward(self, fragments, meshes, **kwargs): | |
| cameras = super()._get_cameras(**kwargs) | |
| texels = self.texel_from_uv(fragments, meshes) | |
| lights = kwargs.get("lights", self.lights) | |
| materials = kwargs.get("materials", self.materials) | |
| blend_params = kwargs.get("blend_params", self.blend_params) | |
| verts, normals, depths, cos_angles = _geometry_shading_with_pixels( | |
| meshes=meshes, | |
| fragments=fragments, | |
| texels=texels, | |
| lights=lights, | |
| cameras=cameras, | |
| materials=materials, | |
| ) | |
| texels = meshes.sample_textures(fragments) | |
| verts = hard_rgb_blend(verts, fragments, blend_params) | |
| normals = hard_rgb_blend(normals, fragments, blend_params) | |
| depths = hard_rgb_blend(depths, fragments, blend_params) | |
| cos_angles = hard_rgb_blend(cos_angles, fragments, blend_params) | |
| from IPython import embed | |
| embed() | |
| texels = hard_rgb_blend(texels, fragments, blend_params) | |
| return verts, normals, depths, cos_angles, texels, fragments | |
| def texel_from_uv(self, fragments, meshes): | |
| texture_tmp = meshes.textures | |
| maps_tmp = texture_tmp.maps_padded() | |
| uv_color = [[[1, 0], [1, 1]], [[0, 0], [0, 1]]] | |
| uv_color = ( | |
| torch.FloatTensor(uv_color).to(maps_tmp[0].device).type(maps_tmp[0].dtype) | |
| ) | |
| uv_texture = TexturesUV( | |
| [uv_color.clone() for t in maps_tmp], | |
| texture_tmp.faces_uvs_padded(), | |
| texture_tmp.verts_uvs_padded(), | |
| sampling_mode="bilinear", | |
| ) | |
| meshes.textures = uv_texture | |
| texels = meshes.sample_textures(fragments) | |
| meshes.textures = texture_tmp | |
| texels = torch.cat((texels, texels[..., -1:] * 0), dim=-1) | |
| return texels | |