| 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 |
|
|
