souging/TRELLIS_TextTo3D

1# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.2#3# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.8import torch9import nvdiffrast.torch as dr10from easydict import EasyDict as edict11from ..representations.mesh import MeshExtractResult12import torch.nn.functional as F13 14 15def intrinsics_to_projection(16        intrinsics: torch.Tensor,17        near: float,18        far: float,19    ) -> torch.Tensor:20    """21    OpenCV intrinsics to OpenGL perspective matrix22 23    Args:24        intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix25        near (float): near plane to clip26        far (float): far plane to clip27    Returns:28        (torch.Tensor): [4, 4] OpenGL perspective matrix29    """30    fx, fy = intrinsics[0, 0], intrinsics[1, 1]31    cx, cy = intrinsics[0, 2], intrinsics[1, 2]32    ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)33    ret[0, 0] = 2 * fx34    ret[1, 1] = 2 * fy35    ret[0, 2] = 2 * cx - 136    ret[1, 2] = - 2 * cy + 137    ret[2, 2] = far / (far - near)38    ret[2, 3] = near * far / (near - far)39    ret[3, 2] = 1.40    return ret41 42 43class MeshRenderer:44    """45    Renderer for the Mesh representation.46 47    Args:48        rendering_options (dict): Rendering options.49        glctx (nvdiffrast.torch.RasterizeGLContext): RasterizeGLContext object for CUDA/OpenGL interop.50        """51    def __init__(self, rendering_options={}, device='cuda'):52        self.rendering_options = edict({53            "resolution": None,54            "near": None,55            "far": None,56            "ssaa": 157        })58        self.rendering_options.update(rendering_options)59        self.glctx = dr.RasterizeCudaContext(device=device)60        self.device=device61        62    def render(63            self,64            mesh : MeshExtractResult,65            extrinsics: torch.Tensor,66            intrinsics: torch.Tensor,67            return_types = ["mask", "normal", "depth"]68        ) -> edict:69        """70        Render the mesh.71 72        Args:73            mesh : meshmodel74            extrinsics (torch.Tensor): (4, 4) camera extrinsics75            intrinsics (torch.Tensor): (3, 3) camera intrinsics76            return_types (list): list of return types, can be "mask", "depth", "normal_map", "normal", "color"77 78        Returns:79            edict based on return_types containing:80                color (torch.Tensor): [3, H, W] rendered color image81                depth (torch.Tensor): [H, W] rendered depth image82                normal (torch.Tensor): [3, H, W] rendered normal image83                normal_map (torch.Tensor): [3, H, W] rendered normal map image84                mask (torch.Tensor): [H, W] rendered mask image85        """86        resolution = self.rendering_options["resolution"]87        near = self.rendering_options["near"]88        far = self.rendering_options["far"]89        ssaa = self.rendering_options["ssaa"]90        91        if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:92            default_img = torch.zeros((1, resolution, resolution, 3), dtype=torch.float32, device=self.device)93            ret_dict = {k : default_img if k in ['normal', 'normal_map', 'color'] else default_img[..., :1] for k in return_types}94            return ret_dict95        96        perspective = intrinsics_to_projection(intrinsics, near, far)97        98        RT = extrinsics.unsqueeze(0)99        full_proj = (perspective @ extrinsics).unsqueeze(0)100        101        vertices = mesh.vertices.unsqueeze(0)102 103        vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1)104        vertices_camera = torch.bmm(vertices_homo, RT.transpose(-1, -2))105        vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2))106        faces_int = mesh.faces.int()107        rast, _ = dr.rasterize(108            self.glctx, vertices_clip, faces_int, (resolution * ssaa, resolution * ssaa))109        110        out_dict = edict()111        for type in return_types:112            img = None113            if type == "mask" :114                img = dr.antialias((rast[..., -1:] > 0).float(), rast, vertices_clip, faces_int)115            elif type == "depth":116                img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_int)[0]117                img = dr.antialias(img, rast, vertices_clip, faces_int)118            elif type == "normal" :119                img = dr.interpolate(120                    mesh.face_normal.reshape(1, -1, 3), rast,121                    torch.arange(mesh.faces.shape[0] * 3, device=self.device, dtype=torch.int).reshape(-1, 3)122                )[0]123                img = dr.antialias(img, rast, vertices_clip, faces_int)124                # normalize norm pictures125                img = (img + 1) / 2126            elif type == "normal_map" :127                img = dr.interpolate(mesh.vertex_attrs[:, 3:].contiguous(), rast, faces_int)[0]128                img = dr.antialias(img, rast, vertices_clip, faces_int)129            elif type == "color" :130                img = dr.interpolate(mesh.vertex_attrs[:, :3].contiguous(), rast, faces_int)[0]131                img = dr.antialias(img, rast, vertices_clip, faces_int)132 133            if ssaa > 1:134                img = F.interpolate(img.permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True)135                img = img.squeeze()136            else:137                img = img.permute(0, 3, 1, 2).squeeze()138            out_dict[type] = img139 140        return out_dict141