| | from typing import * |
| | import torch |
| | from easydict import EasyDict as edict |
| | from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode |
| | import torch.nn.functional as F |
| |
|
| |
|
| | def intrinsics_to_projection( |
| | intrinsics: torch.Tensor, |
| | near: float, |
| | far: float, |
| | ) -> torch.Tensor: |
| | """ |
| | OpenCV intrinsics to OpenGL perspective matrix |
| | |
| | Args: |
| | intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix |
| | near (float): near plane to clip |
| | far (float): far plane to clip |
| | Returns: |
| | (torch.Tensor): [4, 4] OpenGL perspective matrix |
| | """ |
| | fx, fy = intrinsics[0, 0], intrinsics[1, 1] |
| | cx, cy = intrinsics[0, 2], intrinsics[1, 2] |
| | ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) |
| | ret[0, 0] = 2 * fx |
| | ret[1, 1] = 2 * fy |
| | ret[0, 2] = 2 * cx - 1 |
| | ret[1, 2] = - 2 * cy + 1 |
| | ret[2, 2] = (far + near) / (far - near) |
| | ret[2, 3] = 2 * near * far / (near - far) |
| | ret[3, 2] = 1. |
| | return ret |
| | |
| |
|
| | class MeshRenderer: |
| | """ |
| | Renderer for the Mesh representation. |
| | |
| | Args: |
| | rendering_options (dict): Rendering options. |
| | """ |
| | def __init__(self, rendering_options={}, device='cuda'): |
| | if 'dr' not in globals(): |
| | import nvdiffrast.torch as dr |
| | |
| | self.rendering_options = edict({ |
| | "resolution": None, |
| | "near": None, |
| | "far": None, |
| | "ssaa": 1, |
| | "chunk_size": None, |
| | "antialias": True, |
| | "clamp_barycentric_coords": False, |
| | }) |
| | self.rendering_options.update(rendering_options) |
| | self.glctx = dr.RasterizeCudaContext(device=device) |
| | self.device=device |
| | |
| | def render( |
| | self, |
| | mesh : Mesh, |
| | extrinsics: torch.Tensor, |
| | intrinsics: torch.Tensor, |
| | return_types = ["mask", "normal", "depth"], |
| | transformation : Optional[torch.Tensor] = None |
| | ) -> edict: |
| | """ |
| | Render the mesh. |
| | |
| | Args: |
| | mesh : meshmodel |
| | extrinsics (torch.Tensor): (4, 4) camera extrinsics |
| | intrinsics (torch.Tensor): (3, 3) camera intrinsics |
| | return_types (list): list of return types, can be "attr", "mask", "depth", "coord", "normal" |
| | |
| | Returns: |
| | edict based on return_types containing: |
| | attr (torch.Tensor): [C, H, W] rendered attr image |
| | depth (torch.Tensor): [H, W] rendered depth image |
| | normal (torch.Tensor): [3, H, W] rendered normal image |
| | mask (torch.Tensor): [H, W] rendered mask image |
| | """ |
| | if 'dr' not in globals(): |
| | import nvdiffrast.torch as dr |
| | |
| | resolution = self.rendering_options["resolution"] |
| | near = self.rendering_options["near"] |
| | far = self.rendering_options["far"] |
| | ssaa = self.rendering_options["ssaa"] |
| | chunk_size = self.rendering_options["chunk_size"] |
| | antialias = self.rendering_options["antialias"] |
| | clamp_barycentric_coords = self.rendering_options["clamp_barycentric_coords"] |
| | |
| | if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0: |
| | ret_dict = edict() |
| | for type in return_types: |
| | if type == "mask" : |
| | ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device) |
| | elif type == "depth": |
| | ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device) |
| | elif type == "normal": |
| | ret_dict[type] = torch.full((3, resolution, resolution), 0.5, dtype=torch.float32, device=self.device) |
| | elif type == "coord": |
| | ret_dict[type] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device) |
| | elif type == "attr": |
| | if isinstance(mesh, MeshWithVoxel): |
| | ret_dict[type] = torch.zeros((mesh.attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device) |
| | else: |
| | ret_dict[type] = torch.zeros((mesh.vertex_attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device) |
| | return ret_dict |
| | |
| | perspective = intrinsics_to_projection(intrinsics, near, far) |
| | |
| | full_proj = (perspective @ extrinsics).unsqueeze(0) |
| | extrinsics = extrinsics.unsqueeze(0) |
| | |
| | vertices = mesh.vertices.unsqueeze(0) |
| | vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1) |
| | if transformation is not None: |
| | vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2)) |
| | vertices = vertices_homo[..., :3].contiguous() |
| | vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2)) |
| | vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2)) |
| | faces = mesh.faces |
| | |
| | if 'normal' in return_types: |
| | v0 = vertices_camera[0, mesh.faces[:, 0], :3] |
| | v1 = vertices_camera[0, mesh.faces[:, 1], :3] |
| | v2 = vertices_camera[0, mesh.faces[:, 2], :3] |
| | e0 = v1 - v0 |
| | e1 = v2 - v0 |
| | face_normal = torch.cross(e0, e1, dim=1) |
| | face_normal = F.normalize(face_normal, dim=1) |
| | face_normal = torch.where(torch.sum(face_normal * v0, dim=1, keepdim=True) > 0, face_normal, -face_normal) |
| | |
| | out_dict = edict() |
| | if chunk_size is None: |
| | rast, rast_db = dr.rasterize( |
| | self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa) |
| | ) |
| | if clamp_barycentric_coords: |
| | rast[..., :2] = torch.clamp(rast[..., :2], 0, 1) |
| | rast[..., :2] /= torch.where(rast[..., :2].sum(dim=-1, keepdim=True) > 1, rast[..., :2].sum(dim=-1, keepdim=True), torch.ones_like(rast[..., :2])) |
| | for type in return_types: |
| | img = None |
| | if type == "mask" : |
| | img = (rast[..., -1:] > 0).float() |
| | if antialias: img = dr.antialias(img, rast, vertices_clip, faces) |
| | elif type == "depth": |
| | img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0] |
| | if antialias: img = dr.antialias(img, rast, vertices_clip, faces) |
| | elif type == "normal" : |
| | img = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0] |
| | if antialias: img = dr.antialias(img, rast, vertices_clip, faces) |
| | img = (img + 1) / 2 |
| | elif type == "coord": |
| | img = dr.interpolate(vertices, rast, faces)[0] |
| | if antialias: img = dr.antialias(img, rast, vertices_clip, faces) |
| | elif type == "attr": |
| | if isinstance(mesh, MeshWithVoxel): |
| | if 'grid_sample_3d' not in globals(): |
| | from flex_gemm.ops.grid_sample import grid_sample_3d |
| | mask = rast[..., -1:] > 0 |
| | xyz = dr.interpolate(vertices, rast, faces)[0] |
| | xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3) |
| | img = grid_sample_3d( |
| | mesh.attrs, |
| | torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1), |
| | mesh.voxel_shape, |
| | xyz, |
| | mode='trilinear' |
| | ) |
| | img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask |
| | elif isinstance(mesh, MeshWithPbrMaterial): |
| | tri_id = rast[0, :, :, -1:] |
| | mask = tri_id > 0 |
| | uv_coords = mesh.uv_coords.reshape(1, -1, 2) |
| | texc, texd = dr.interpolate( |
| | uv_coords, |
| | rast, |
| | torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3), |
| | rast_db=rast_db, |
| | diff_attrs='all' |
| | ) |
| | |
| | texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3) |
| | texc = torch.clamp(texc, min=-1e3, max=1e3) |
| | texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3) |
| | texd = torch.clamp(texd, min=-1e3, max=1e3) |
| | mid = mesh.material_ids[(tri_id - 1).long()] |
| | imgs = { |
| | 'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device), |
| | 'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device), |
| | 'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device), |
| | 'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) |
| | } |
| | for id, mat in enumerate(mesh.materials): |
| | mat_mask = (mid == id).float() * mask.float() |
| | mat_texc = texc * mat_mask |
| | mat_texd = texd * mat_mask |
| |
|
| | if mat.base_color_texture is not None: |
| | base_color = dr.texture( |
| | mat.base_color_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['base_color'] += base_color * mat.base_color_factor * mat_mask |
| | else: |
| | imgs['base_color'] += mat.base_color_factor * mat_mask |
| | |
| | if mat.metallic_texture is not None: |
| | metallic = dr.texture( |
| | mat.metallic_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['metallic'] += metallic * mat.metallic_factor * mat_mask |
| | else: |
| | imgs['metallic'] += mat.metallic_factor * mat_mask |
| |
|
| | if mat.roughness_texture is not None: |
| | roughness = dr.texture( |
| | mat.roughness_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['roughness'] += roughness * mat.roughness_factor * mat_mask |
| | else: |
| | imgs['roughness'] += mat.roughness_factor * mat_mask |
| |
|
| | if mat.alpha_mode == AlphaMode.OPAQUE: |
| | imgs['alpha'] += 1.0 * mat_mask |
| | else: |
| | if mat.alpha_texture is not None: |
| | alpha = dr.texture( |
| | mat.alpha_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | if mat.alpha_mode == AlphaMode.MASK: |
| | imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask |
| | elif mat.alpha_mode == AlphaMode.BLEND: |
| | imgs['alpha'] += alpha * mat.alpha_factor * mat_mask |
| | else: |
| | if mat.alpha_mode == AlphaMode.MASK: |
| | imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask |
| | elif mat.alpha_mode == AlphaMode.BLEND: |
| | imgs['alpha'] += mat.alpha_factor * mat_mask |
| | |
| | img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0) |
| | else: |
| | img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces)[0] |
| | if antialias: img = dr.antialias(img, rast, vertices_clip, faces) |
| | |
| | out_dict[type] = img |
| | else: |
| | z_buffer = torch.full((1, resolution * ssaa, resolution * ssaa), torch.inf, device=self.device, dtype=torch.float32) |
| | for i in range(0, faces.shape[0], chunk_size): |
| | faces_chunk = faces[i:i+chunk_size] |
| | rast, rast_db = dr.rasterize( |
| | self.glctx, vertices_clip, faces_chunk, (resolution * ssaa, resolution * ssaa) |
| | ) |
| | z_filter = torch.logical_and( |
| | rast[..., 3] != 0, |
| | rast[..., 2] < z_buffer |
| | ) |
| | z_buffer[z_filter] = rast[z_filter][..., 2] |
| | |
| | for type in return_types: |
| | img = None |
| | if type == "mask" : |
| | img = (rast[..., -1:] > 0).float() |
| | elif type == "depth": |
| | img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_chunk)[0] |
| | elif type == "normal" : |
| | face_normal_chunk = face_normal[i:i+chunk_size] |
| | img = dr.interpolate(face_normal_chunk.unsqueeze(0), rast, torch.arange(face_normal_chunk.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0] |
| | img = (img + 1) / 2 |
| | elif type == "coord": |
| | img = dr.interpolate(vertices, rast, faces_chunk)[0] |
| | elif type == "attr": |
| | if isinstance(mesh, MeshWithVoxel): |
| | if 'grid_sample_3d' not in globals(): |
| | from flex_gemm.ops.grid_sample import grid_sample_3d |
| | mask = rast[..., -1:] > 0 |
| | xyz = dr.interpolate(vertices, rast, faces_chunk)[0] |
| | xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3) |
| | img = grid_sample_3d( |
| | mesh.attrs, |
| | torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1), |
| | mesh.voxel_shape, |
| | xyz, |
| | mode='trilinear' |
| | ) |
| | img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask |
| | elif isinstance(mesh, MeshWithPbrMaterial): |
| | tri_id = rast[0, :, :, -1:] |
| | mask = tri_id > 0 |
| | uv_coords = mesh.uv_coords.reshape(1, -1, 2) |
| | texc, texd = dr.interpolate( |
| | uv_coords, |
| | rast, |
| | torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3), |
| | rast_db=rast_db, |
| | diff_attrs='all' |
| | ) |
| | |
| | texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3) |
| | texc = torch.clamp(texc, min=-1e3, max=1e3) |
| | texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3) |
| | texd = torch.clamp(texd, min=-1e3, max=1e3) |
| | mid = mesh.material_ids[(tri_id - 1).long()] |
| | imgs = { |
| | 'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device), |
| | 'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device), |
| | 'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device), |
| | 'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) |
| | } |
| | for id, mat in enumerate(mesh.materials): |
| | mat_mask = (mid == id).float() * mask.float() |
| | mat_texc = texc * mat_mask |
| | mat_texd = texd * mat_mask |
| |
|
| | if mat.base_color_texture is not None: |
| | base_color = dr.texture( |
| | mat.base_color_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['base_color'] += base_color * mat.base_color_factor * mat_mask |
| | else: |
| | imgs['base_color'] += mat.base_color_factor * mat_mask |
| | |
| | if mat.metallic_texture is not None: |
| | metallic = dr.texture( |
| | mat.metallic_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['metallic'] += metallic * mat.metallic_factor * mat_mask |
| | else: |
| | imgs['metallic'] += mat.metallic_factor * mat_mask |
| |
|
| | if mat.roughness_texture is not None: |
| | roughness = dr.texture( |
| | mat.roughness_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | imgs['roughness'] += roughness * mat.roughness_factor * mat_mask |
| | else: |
| | imgs['roughness'] += mat.roughness_factor * mat_mask |
| |
|
| | if mat.alpha_mode == AlphaMode.OPAQUE: |
| | imgs['alpha'] += 1.0 * mat_mask |
| | else: |
| | if mat.alpha_texture is not None: |
| | alpha = dr.texture( |
| | mat.alpha_texture.image.unsqueeze(0), |
| | mat_texc, |
| | mat_texd, |
| | filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', |
| | boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' |
| | )[0] |
| | if mat.alpha_mode == AlphaMode.MASK: |
| | imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask |
| | elif mat.alpha_mode == AlphaMode.BLEND: |
| | imgs['alpha'] += alpha * mat.alpha_factor * mat_mask |
| | else: |
| | if mat.alpha_mode == AlphaMode.MASK: |
| | imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask |
| | elif mat.alpha_mode == AlphaMode.BLEND: |
| | imgs['alpha'] += mat.alpha_factor * mat_mask |
| | |
| | img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0) |
| | else: |
| | img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces_chunk)[0] |
| | |
| | if type not in out_dict: |
| | out_dict[type] = img |
| | else: |
| | out_dict[type][z_filter] = img[z_filter] |
| |
|
| | for type in return_types: |
| | img = out_dict[type] |
| | if ssaa > 1: |
| | img = F.interpolate(img.permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True) |
| | img = img.squeeze() |
| | else: |
| | img = img.permute(0, 3, 1, 2).squeeze() |
| | out_dict[type] = img |
| |
|
| | if isinstance(mesh, (MeshWithVoxel, MeshWithPbrMaterial)) and 'attr' in return_types: |
| | for k, s in mesh.layout.items(): |
| | out_dict[k] = out_dict['attr'][s] |
| | del out_dict['attr'] |
| | |
| | return out_dict |
| |
|
