| from typing import Optional | |
| import torch | |
| import pytorch3d | |
| from pytorch3d.io import load_objs_as_meshes, load_obj, save_obj | |
| from pytorch3d.ops import interpolate_face_attributes | |
| from pytorch3d.structures import Meshes | |
| 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.lighting import AmbientLights | |
| from pytorch3d.renderer.materials import Materials | |
| from pytorch3d.renderer.mesh.shader import ShaderBase | |
| from pytorch3d.renderer.mesh.shading import _apply_lighting, flat_shading | |
| from pytorch3d.renderer.mesh.rasterizer import Fragments | |
| """ | |
| Customized the original pytorch3d hard flat shader to support N channel flat shading | |
| """ | |
| class HardNChannelFlatShader(ShaderBase): | |
| """ | |
| Per face lighting - the lighting model is applied using the average face | |
| position and the face normal. The blending function hard assigns | |
| the color of the closest face for each pixel. | |
| To use the default values, simply initialize the shader with the desired | |
| device e.g. | |
| .. code-block:: | |
| shader = HardFlatShader(device=torch.device("cuda:0")) | |
| """ | |
| def __init__( | |
| self, | |
| device="cpu", | |
| cameras: Optional[TensorProperties] = None, | |
| lights: Optional[TensorProperties] = None, | |
| materials: Optional[Materials] = None, | |
| blend_params: Optional[BlendParams] = None, | |
| channels: int = 3, | |
| ): | |
| self.channels = channels | |
| ones = ((1.0,) * channels,) | |
| zeros = ((0.0,) * channels,) | |
| if ( | |
| not isinstance(lights, AmbientLights) | |
| or not lights.ambient_color.shape[-1] == channels | |
| ): | |
| lights = AmbientLights( | |
| ambient_color=ones, | |
| device=device, | |
| ) | |
| if not materials or not materials.ambient_color.shape[-1] == channels: | |
| materials = Materials( | |
| device=device, | |
| diffuse_color=zeros, | |
| ambient_color=ones, | |
| specular_color=zeros, | |
| shininess=0.0, | |
| ) | |
| blend_params_new = BlendParams(background_color=(1.0,) * channels) | |
| if not isinstance(blend_params, BlendParams): | |
| blend_params = blend_params_new | |
| else: | |
| background_color_ = blend_params.background_color | |
| if ( | |
| isinstance(background_color_, Sequence[float]) | |
| and not len(background_color_) == channels | |
| ): | |
| blend_params = blend_params_new | |
| if ( | |
| isinstance(background_color_, torch.Tensor) | |
| and not background_color_.shape[-1] == channels | |
| ): | |
| blend_params = blend_params_new | |
| super().__init__( | |
| device, | |
| cameras, | |
| lights, | |
| materials, | |
| blend_params, | |
| ) | |
| def forward(self, fragments: Fragments, meshes: Meshes, **kwargs) -> torch.Tensor: | |
| cameras = super()._get_cameras(**kwargs) | |
| texels = meshes.sample_textures(fragments) | |
| lights = kwargs.get("lights", self.lights) | |
| materials = kwargs.get("materials", self.materials) | |
| blend_params = kwargs.get("blend_params", self.blend_params) | |
| colors = flat_shading( | |
| meshes=meshes, | |
| fragments=fragments, | |
| texels=texels, | |
| lights=lights, | |
| cameras=cameras, | |
| materials=materials, | |
| ) | |
| images = hard_rgb_blend(colors, fragments, blend_params) | |
| return images | |