souging/TRELLIS_TextTo3D

1import numpy as np2import torch3import torch.nn.functional as F4import math5import cv26from scipy.stats import qmc7from easydict import EasyDict as edict8from ..representations.octree import DfsOctree9 10 11def intrinsics_to_projection(12        intrinsics: torch.Tensor,13        near: float,14        far: float,15    ) -> torch.Tensor:16    """17    OpenCV intrinsics to OpenGL perspective matrix18 19    Args:20        intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix21        near (float): near plane to clip22        far (float): far plane to clip23    Returns:24        (torch.Tensor): [4, 4] OpenGL perspective matrix25    """26    fx, fy = intrinsics[0, 0], intrinsics[1, 1]27    cx, cy = intrinsics[0, 2], intrinsics[1, 2]28    ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)29    ret[0, 0] = 2 * fx30    ret[1, 1] = 2 * fy31    ret[0, 2] = 2 * cx - 132    ret[1, 2] = - 2 * cy + 133    ret[2, 2] = far / (far - near)34    ret[2, 3] = near * far / (near - far)35    ret[3, 2] = 1.36    return ret37 38 39def render(viewpoint_camera, octree : DfsOctree, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, used_rank = None, colors_overwrite = None, aux=None, halton_sampler=None):40    """41    Render the scene. 42    43    Background tensor (bg_color) must be on GPU!44    """45    # lazy import46    if 'OctreeTrivecRasterizer' not in globals():47        from diffoctreerast import OctreeVoxelRasterizer, OctreeGaussianRasterizer, OctreeTrivecRasterizer, OctreeDecoupolyRasterizer48    49    # Set up rasterization configuration50    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)51    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)52 53    raster_settings = edict(54        image_height=int(viewpoint_camera.image_height),55        image_width=int(viewpoint_camera.image_width),56        tanfovx=tanfovx,57        tanfovy=tanfovy,58        bg=bg_color,59        scale_modifier=scaling_modifier,60        viewmatrix=viewpoint_camera.world_view_transform,61        projmatrix=viewpoint_camera.full_proj_transform,62        sh_degree=octree.active_sh_degree,63        campos=viewpoint_camera.camera_center,64        with_distloss=pipe.with_distloss,65        jitter=pipe.jitter,66        debug=pipe.debug,67    )68 69    positions = octree.get_xyz70    if octree.primitive == "voxel":71        densities = octree.get_density72    elif octree.primitive == "gaussian":73        opacities = octree.get_opacity74    elif octree.primitive == "trivec":75        trivecs = octree.get_trivec76        densities = octree.get_density77        raster_settings.density_shift = octree.density_shift78    elif octree.primitive == "decoupoly":79        decoupolys_V, decoupolys_g = octree.get_decoupoly80        densities = octree.get_density81        raster_settings.density_shift = octree.density_shift82    else:83        raise ValueError(f"Unknown primitive {octree.primitive}")84    depths = octree.get_depth85 86    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors87    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.88    colors_precomp = None89    shs = octree.get_features90    if octree.primitive in ["voxel", "gaussian"] and colors_overwrite is not None:91        colors_precomp = colors_overwrite92        shs = None93 94    ret = edict()95 96    if octree.primitive == "voxel":97        renderer = OctreeVoxelRasterizer(raster_settings=raster_settings)98        rgb, depth, alpha, distloss = renderer(99            positions = positions,100            densities = densities,101            shs = shs,102            colors_precomp = colors_precomp,103            depths = depths,104            aabb = octree.aabb,105            aux = aux,106        )107        ret['rgb'] = rgb108        ret['depth'] = depth109        ret['alpha'] = alpha110        ret['distloss'] = distloss111    elif octree.primitive == "gaussian":112        renderer = OctreeGaussianRasterizer(raster_settings=raster_settings)113        rgb, depth, alpha = renderer(114            positions = positions,115            opacities = opacities,116            shs = shs,117            colors_precomp = colors_precomp,118            depths = depths,119            aabb = octree.aabb,120            aux = aux,121        )122        ret['rgb'] = rgb123        ret['depth'] = depth124        ret['alpha'] = alpha125    elif octree.primitive == "trivec":126        raster_settings.used_rank = used_rank if used_rank is not None else trivecs.shape[1]127        renderer = OctreeTrivecRasterizer(raster_settings=raster_settings)128        rgb, depth, alpha, percent_depth = renderer(129            positions = positions,130            trivecs = trivecs,131            densities = densities,132            shs = shs,133            colors_precomp = colors_precomp,134            colors_overwrite = colors_overwrite,135            depths = depths,136            aabb = octree.aabb,137            aux = aux,138            halton_sampler = halton_sampler,139        )140        ret['percent_depth'] = percent_depth141        ret['rgb'] = rgb142        ret['depth'] = depth143        ret['alpha'] = alpha144    elif octree.primitive == "decoupoly":145        raster_settings.used_rank = used_rank if used_rank is not None else decoupolys_V.shape[1]146        renderer = OctreeDecoupolyRasterizer(raster_settings=raster_settings)147        rgb, depth, alpha = renderer(148            positions = positions,149            decoupolys_V = decoupolys_V,150            decoupolys_g = decoupolys_g,151            densities = densities,152            shs = shs,153            colors_precomp = colors_precomp,154            depths = depths,155            aabb = octree.aabb,156            aux = aux,157        )158        ret['rgb'] = rgb159        ret['depth'] = depth160        ret['alpha'] = alpha161    162    return ret163 164 165class OctreeRenderer:166    """167    Renderer for the Voxel representation.168 169    Args:170        rendering_options (dict): Rendering options.171    """172 173    def __init__(self, rendering_options={}) -> None:174        try:175            import diffoctreerast176        except ImportError:177            print("\033[93m[WARNING] diffoctreerast is not installed. The renderer will be disabled.\033[0m")178            self.unsupported = True179        else:180            self.unsupported = False181        182        self.pipe = edict({183            "with_distloss": False,184            "with_aux": False,185            "scale_modifier": 1.0,186            "used_rank": None,187            "jitter": False,188            "debug": False,189        })190        self.rendering_options = edict({191            "resolution": None,192            "near": None,193            "far": None,194            "ssaa": 1,195            "bg_color": 'random',196        })197        self.halton_sampler = qmc.Halton(2, scramble=False)198        self.rendering_options.update(rendering_options)199        self.bg_color = None200    201    def render(202            self,203            octree: DfsOctree,204            extrinsics: torch.Tensor,205            intrinsics: torch.Tensor,206            colors_overwrite: torch.Tensor = None,207        ) -> edict:208        """209        Render the octree.210 211        Args:212            octree (Octree): octree213            extrinsics (torch.Tensor): (4, 4) camera extrinsics214            intrinsics (torch.Tensor): (3, 3) camera intrinsics215            colors_overwrite (torch.Tensor): (N, 3) override color216 217        Returns:218            edict containing:219                color (torch.Tensor): (3, H, W) rendered color220                depth (torch.Tensor): (H, W) rendered depth221                alpha (torch.Tensor): (H, W) rendered alpha222                distloss (Optional[torch.Tensor]): (H, W) rendered distance loss223                percent_depth (Optional[torch.Tensor]): (H, W) rendered percent depth224                aux (Optional[edict]): auxiliary tensors225        """226        resolution = self.rendering_options["resolution"]227        near = self.rendering_options["near"]228        far = self.rendering_options["far"]229        ssaa = self.rendering_options["ssaa"]230        231        if self.unsupported:232            image = np.zeros((512, 512, 3), dtype=np.uint8)233            text_bbox = cv2.getTextSize("Unsupported", cv2.FONT_HERSHEY_SIMPLEX, 2, 3)[0]234            origin = (512 - text_bbox[0]) // 2, (512 - text_bbox[1]) // 2235            image = cv2.putText(image, "Unsupported", origin, cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3, cv2.LINE_AA)236            return {237                'color': torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255,238            }239        240        if self.rendering_options["bg_color"] == 'random':241            self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda")242            if np.random.rand() < 0.5:243                self.bg_color += 1244        else:245            self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda")246 247        if self.pipe["with_aux"]:248            aux = {249                'grad_color2': torch.zeros((octree.num_leaf_nodes, 3), dtype=torch.float32, requires_grad=True, device="cuda") + 0,250                'contributions': torch.zeros((octree.num_leaf_nodes, 1), dtype=torch.float32, requires_grad=True, device="cuda") + 0,251            }252            for k in aux.keys():253                aux[k].requires_grad_()254                aux[k].retain_grad()255        else:256            aux = None257 258        view = extrinsics259        perspective = intrinsics_to_projection(intrinsics, near, far)260        camera = torch.inverse(view)[:3, 3]261        focalx = intrinsics[0, 0]262        focaly = intrinsics[1, 1]263        fovx = 2 * torch.atan(0.5 / focalx)264        fovy = 2 * torch.atan(0.5 / focaly)265            266        camera_dict = edict({267            "image_height": resolution * ssaa,268            "image_width": resolution * ssaa,269            "FoVx": fovx,270            "FoVy": fovy,271            "znear": near,272            "zfar": far,273            "world_view_transform": view.T.contiguous(),274            "projection_matrix": perspective.T.contiguous(),275            "full_proj_transform": (perspective @ view).T.contiguous(),276            "camera_center": camera277        })278 279        # Render280        render_ret = render(camera_dict, octree, self.pipe, self.bg_color, aux=aux, colors_overwrite=colors_overwrite, scaling_modifier=self.pipe.scale_modifier, used_rank=self.pipe.used_rank, halton_sampler=self.halton_sampler)281 282        if ssaa > 1:283            render_ret.rgb = F.interpolate(render_ret.rgb[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()284            render_ret.depth = F.interpolate(render_ret.depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()285            render_ret.alpha = F.interpolate(render_ret.alpha[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()286            if hasattr(render_ret, 'percent_depth'):287                render_ret.percent_depth = F.interpolate(render_ret.percent_depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()288 289        ret = edict({290            'color': render_ret.rgb,291            'depth': render_ret.depth,292            'alpha': render_ret.alpha,293        })294        if self.pipe["with_distloss"] and 'distloss' in render_ret:295            ret['distloss'] = render_ret.distloss296        if self.pipe["with_aux"]:297            ret['aux'] = aux298        if hasattr(render_ret, 'percent_depth'):299            ret['percent_depth'] = render_ret.percent_depth300        return ret301