Delete trellis/renderers

trellis/renderers/__init__.py

@@ -1,31 +0,0 @@
-import importlib
-
-__attributes = {
-    'OctreeRenderer': 'octree_renderer',
-    'GaussianRenderer': 'gaussian_render',
-    'MeshRenderer': 'mesh_renderer',
-}
-
-__submodules = []
-
-__all__ = list(__attributes.keys()) + __submodules
-
-def __getattr__(name):
-    if name not in globals():
-        if name in __attributes:
-            module_name = __attributes[name]
-            module = importlib.import_module(f".{module_name}", __name__)
-            globals()[name] = getattr(module, name)
-        elif name in __submodules:
-            module = importlib.import_module(f".{name}", __name__)
-            globals()[name] = module
-        else:
-            raise AttributeError(f"module {__name__} has no attribute {name}")
-    return globals()[name]
-
-
-# For Pylance
-if __name__ == '__main__':
-    from .octree_renderer import OctreeRenderer
-    from .gaussian_render import GaussianRenderer
-    from .mesh_renderer import MeshRenderer

trellis/renderers/gaussian_render.py

@@ -1,231 +0,0 @@
-#
-# Copyright (C) 2023, Inria
-# GRAPHDECO research group, https://team.inria.fr/graphdeco
-# All rights reserved.
-#
-# This software is free for non-commercial, research and evaluation use 
-# under the terms of the LICENSE.md file.
-#
-# For inquiries contact  [email protected]
-#
-
-import torch
-import math
-from easydict import EasyDict as edict
-import numpy as np
-from ..representations.gaussian import Gaussian
-from .sh_utils import eval_sh
-import torch.nn.functional as F
-from easydict import EasyDict as edict
-
-
-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 / (far - near)
-    ret[2, 3] = near * far / (near - far)
-    ret[3, 2] = 1.
-    return ret
-
-
-def render(viewpoint_camera, pc : Gaussian, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):
-    """
-    Render the scene. 
-    
-    Background tensor (bg_color) must be on GPU!
-    """
-    # lazy import
-    if 'GaussianRasterizer' not in globals():
-        from diff_gaussian_rasterization import GaussianRasterizer, GaussianRasterizationSettings
-    
-    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
-    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
-    try:
-        screenspace_points.retain_grad()
-    except:
-        pass
-    # Set up rasterization configuration
-    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
-    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
-    
-    kernel_size = pipe.kernel_size
-    subpixel_offset = torch.zeros((int(viewpoint_camera.image_height), int(viewpoint_camera.image_width), 2), dtype=torch.float32, device="cuda")
-
-    raster_settings = GaussianRasterizationSettings(
-        image_height=int(viewpoint_camera.image_height),
-        image_width=int(viewpoint_camera.image_width),
-        tanfovx=tanfovx,
-        tanfovy=tanfovy,
-        kernel_size=kernel_size,
-        subpixel_offset=subpixel_offset,
-        bg=bg_color,
-        scale_modifier=scaling_modifier,
-        viewmatrix=viewpoint_camera.world_view_transform,
-        projmatrix=viewpoint_camera.full_proj_transform,
-        sh_degree=pc.active_sh_degree,
-        campos=viewpoint_camera.camera_center,
-        prefiltered=False,
-        debug=pipe.debug
-    )
-    
-    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
-
-    means3D = pc.get_xyz
-    means2D = screenspace_points
-    opacity = pc.get_opacity
-
-    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
-    # scaling / rotation by the rasterizer.
-    scales = None
-    rotations = None
-    cov3D_precomp = None
-    if pipe.compute_cov3D_python:
-        cov3D_precomp = pc.get_covariance(scaling_modifier)
-    else:
-        scales = pc.get_scaling
-        rotations = pc.get_rotation
-
-    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
-    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
-    shs = None
-    colors_precomp = None
-    if override_color is None:
-        if pipe.convert_SHs_python:
-            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
-            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
-            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)
-            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
-            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
-        else:
-            shs = pc.get_features
-    else:
-        colors_precomp = override_color
-
-    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
-    rendered_image, radii = rasterizer(
-        means3D = means3D,
-        means2D = means2D,
-        shs = shs,
-        colors_precomp = colors_precomp,
-        opacities = opacity,
-        scales = scales,
-        rotations = rotations,
-        cov3D_precomp = cov3D_precomp
-    )
-
-    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
-    # They will be excluded from value updates used in the splitting criteria.
-    return edict({"render": rendered_image,
-            "viewspace_points": screenspace_points,
-            "visibility_filter" : radii > 0,
-            "radii": radii})
-
-
-class GaussianRenderer:
-    """
-    Renderer for the Voxel representation.
-
-    Args:
-        rendering_options (dict): Rendering options.
-    """
-
-    def __init__(self, rendering_options={}) -> None:
-        self.pipe = edict({
-            "kernel_size": 0.1,
-            "convert_SHs_python": False,
-            "compute_cov3D_python": False,
-            "scale_modifier": 1.0,
-            "debug": False
-        })
-        self.rendering_options = edict({
-            "resolution": None,
-            "near": None,
-            "far": None,
-            "ssaa": 1,
-            "bg_color": 'random',
-        })
-        self.rendering_options.update(rendering_options)
-        self.bg_color = None
-    
-    def render(
-            self,
-            gausssian: Gaussian,
-            extrinsics: torch.Tensor,
-            intrinsics: torch.Tensor,
-            colors_overwrite: torch.Tensor = None
-        ) -> edict:
-        """
-        Render the gausssian.
-
-        Args:
-            gaussian : gaussianmodule
-            extrinsics (torch.Tensor): (4, 4) camera extrinsics
-            intrinsics (torch.Tensor): (3, 3) camera intrinsics
-            colors_overwrite (torch.Tensor): (N, 3) override color
-
-        Returns:
-            edict containing:
-                color (torch.Tensor): (3, H, W) rendered color image
-        """
-        resolution = self.rendering_options["resolution"]
-        near = self.rendering_options["near"]
-        far = self.rendering_options["far"]
-        ssaa = self.rendering_options["ssaa"]
-        
-        if self.rendering_options["bg_color"] == 'random':
-            self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda")
-            if np.random.rand() < 0.5:
-                self.bg_color += 1
-        else:
-            self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda")
-
-        view = extrinsics
-        perspective = intrinsics_to_projection(intrinsics, near, far)
-        camera = torch.inverse(view)[:3, 3]
-        focalx = intrinsics[0, 0]
-        focaly = intrinsics[1, 1]
-        fovx = 2 * torch.atan(0.5 / focalx)
-        fovy = 2 * torch.atan(0.5 / focaly)
-            
-        camera_dict = edict({
-            "image_height": resolution * ssaa,
-            "image_width": resolution * ssaa,
-            "FoVx": fovx,
-            "FoVy": fovy,
-            "znear": near,
-            "zfar": far,
-            "world_view_transform": view.T.contiguous(),
-            "projection_matrix": perspective.T.contiguous(),
-            "full_proj_transform": (perspective @ view).T.contiguous(),
-            "camera_center": camera
-        })
-
-        # Render
-        render_ret = render(camera_dict, gausssian, self.pipe, self.bg_color, override_color=colors_overwrite, scaling_modifier=self.pipe.scale_modifier)
-
-        if ssaa > 1:
-            render_ret.render = F.interpolate(render_ret.render[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()
-            
-        ret = edict({
-            'color': render_ret['render']
-        })
-        return ret

trellis/renderers/mesh_renderer.py

@@ -1,140 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
-#
-# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
-import torch
-import nvdiffrast.torch as dr
-from easydict import EasyDict as edict
-from ..representations.mesh import MeshExtractResult
-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 / (far - near)
-    ret[2, 3] = near * far / (near - far)
-    ret[3, 2] = 1.
-    return ret
-
-
-class MeshRenderer:
-    """
-    Renderer for the Mesh representation.
-
-    Args:
-        rendering_options (dict): Rendering options.
-        glctx (nvdiffrast.torch.RasterizeGLContext): RasterizeGLContext object for CUDA/OpenGL interop.
-        """
-    def __init__(self, rendering_options={}, device='cuda'):
-        self.rendering_options = edict({
-            "resolution": None,
-            "near": None,
-            "far": None,
-            "ssaa": 1
-        })
-        self.rendering_options.update(rendering_options)
-        self.glctx = dr.RasterizeCudaContext(device=device)
-        self.device=device
-        
-    def render(
-            self,
-            mesh : MeshExtractResult,
-            extrinsics: torch.Tensor,
-            intrinsics: torch.Tensor,
-            return_types = ["mask", "normal", "depth"]
-        ) -> 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 "mask", "depth", "normal_map", "normal", "color"
-
-        Returns:
-            edict based on return_types containing:
-                color (torch.Tensor): [3, H, W] rendered color image
-                depth (torch.Tensor): [H, W] rendered depth image
-                normal (torch.Tensor): [3, H, W] rendered normal image
-                normal_map (torch.Tensor): [3, H, W] rendered normal map image
-                mask (torch.Tensor): [H, W] rendered mask image
-        """
-        resolution = self.rendering_options["resolution"]
-        near = self.rendering_options["near"]
-        far = self.rendering_options["far"]
-        ssaa = self.rendering_options["ssaa"]
-        
-        if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
-            default_img = torch.zeros((1, resolution, resolution, 3), dtype=torch.float32, device=self.device)
-            ret_dict = {k : default_img if k in ['normal', 'normal_map', 'color'] else default_img[..., :1] for k in return_types}
-            return ret_dict
-        
-        perspective = intrinsics_to_projection(intrinsics, near, far)
-        
-        RT = extrinsics.unsqueeze(0)
-        full_proj = (perspective @ extrinsics).unsqueeze(0)
-        
-        vertices = mesh.vertices.unsqueeze(0)
-
-        vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1)
-        vertices_camera = torch.bmm(vertices_homo, RT.transpose(-1, -2))
-        vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2))
-        faces_int = mesh.faces.int()
-        rast, _ = dr.rasterize(
-            self.glctx, vertices_clip, faces_int, (resolution * ssaa, resolution * ssaa))
-        
-        out_dict = edict()
-        for type in return_types:
-            img = None
-            if type == "mask" :
-                img = dr.antialias((rast[..., -1:] > 0).float(), rast, vertices_clip, faces_int)
-            elif type == "depth":
-                img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_int)[0]
-                img = dr.antialias(img, rast, vertices_clip, faces_int)
-            elif type == "normal" :
-                img = dr.interpolate(
-                    mesh.face_normal.reshape(1, -1, 3), rast,
-                    torch.arange(mesh.faces.shape[0] * 3, device=self.device, dtype=torch.int).reshape(-1, 3)
-                )[0]
-                img = dr.antialias(img, rast, vertices_clip, faces_int)
-                # normalize norm pictures
-                img = (img + 1) / 2
-            elif type == "normal_map" :
-                img = dr.interpolate(mesh.vertex_attrs[:, 3:].contiguous(), rast, faces_int)[0]
-                img = dr.antialias(img, rast, vertices_clip, faces_int)
-            elif type == "color" :
-                img = dr.interpolate(mesh.vertex_attrs[:, :3].contiguous(), rast, faces_int)[0]
-                img = dr.antialias(img, rast, vertices_clip, faces_int)
-
-            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
-
-        return out_dict

trellis/renderers/octree_renderer.py

@@ -1,300 +0,0 @@
-import numpy as np
-import torch
-import torch.nn.functional as F
-import math
-import cv2
-from scipy.stats import qmc
-from easydict import EasyDict as edict
-from ..representations.octree import DfsOctree
-
-
-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 / (far - near)
-    ret[2, 3] = near * far / (near - far)
-    ret[3, 2] = 1.
-    return ret
-
-
-def 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):
-    """
-    Render the scene. 
-    
-    Background tensor (bg_color) must be on GPU!
-    """
-    # lazy import
-    if 'OctreeTrivecRasterizer' not in globals():
-        from diffoctreerast import OctreeVoxelRasterizer, OctreeGaussianRasterizer, OctreeTrivecRasterizer, OctreeDecoupolyRasterizer
-    
-    # Set up rasterization configuration
-    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
-    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
-
-    raster_settings = edict(
-        image_height=int(viewpoint_camera.image_height),
-        image_width=int(viewpoint_camera.image_width),
-        tanfovx=tanfovx,
-        tanfovy=tanfovy,
-        bg=bg_color,
-        scale_modifier=scaling_modifier,
-        viewmatrix=viewpoint_camera.world_view_transform,
-        projmatrix=viewpoint_camera.full_proj_transform,
-        sh_degree=octree.active_sh_degree,
-        campos=viewpoint_camera.camera_center,
-        with_distloss=pipe.with_distloss,
-        jitter=pipe.jitter,
-        debug=pipe.debug,
-    )
-
-    positions = octree.get_xyz
-    if octree.primitive == "voxel":
-        densities = octree.get_density
-    elif octree.primitive == "gaussian":
-        opacities = octree.get_opacity
-    elif octree.primitive == "trivec":
-        trivecs = octree.get_trivec
-        densities = octree.get_density
-        raster_settings.density_shift = octree.density_shift
-    elif octree.primitive == "decoupoly":
-        decoupolys_V, decoupolys_g = octree.get_decoupoly
-        densities = octree.get_density
-        raster_settings.density_shift = octree.density_shift
-    else:
-        raise ValueError(f"Unknown primitive {octree.primitive}")
-    depths = octree.get_depth
-
-    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
-    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
-    colors_precomp = None
-    shs = octree.get_features
-    if octree.primitive in ["voxel", "gaussian"] and colors_overwrite is not None:
-        colors_precomp = colors_overwrite
-        shs = None
-
-    ret = edict()
-
-    if octree.primitive == "voxel":
-        renderer = OctreeVoxelRasterizer(raster_settings=raster_settings)
-        rgb, depth, alpha, distloss = renderer(
-            positions = positions,
-            densities = densities,
-            shs = shs,
-            colors_precomp = colors_precomp,
-            depths = depths,
-            aabb = octree.aabb,
-            aux = aux,
-        )
-        ret['rgb'] = rgb
-        ret['depth'] = depth
-        ret['alpha'] = alpha
-        ret['distloss'] = distloss
-    elif octree.primitive == "gaussian":
-        renderer = OctreeGaussianRasterizer(raster_settings=raster_settings)
-        rgb, depth, alpha = renderer(
-            positions = positions,
-            opacities = opacities,
-            shs = shs,
-            colors_precomp = colors_precomp,
-            depths = depths,
-            aabb = octree.aabb,
-            aux = aux,
-        )
-        ret['rgb'] = rgb
-        ret['depth'] = depth
-        ret['alpha'] = alpha
-    elif octree.primitive == "trivec":
-        raster_settings.used_rank = used_rank if used_rank is not None else trivecs.shape[1]
-        renderer = OctreeTrivecRasterizer(raster_settings=raster_settings)
-        rgb, depth, alpha, percent_depth = renderer(
-            positions = positions,
-            trivecs = trivecs,
-            densities = densities,
-            shs = shs,
-            colors_precomp = colors_precomp,
-            colors_overwrite = colors_overwrite,
-            depths = depths,
-            aabb = octree.aabb,
-            aux = aux,
-            halton_sampler = halton_sampler,
-        )
-        ret['percent_depth'] = percent_depth
-        ret['rgb'] = rgb
-        ret['depth'] = depth
-        ret['alpha'] = alpha
-    elif octree.primitive == "decoupoly":
-        raster_settings.used_rank = used_rank if used_rank is not None else decoupolys_V.shape[1]
-        renderer = OctreeDecoupolyRasterizer(raster_settings=raster_settings)
-        rgb, depth, alpha = renderer(
-            positions = positions,
-            decoupolys_V = decoupolys_V,
-            decoupolys_g = decoupolys_g,
-            densities = densities,
-            shs = shs,
-            colors_precomp = colors_precomp,
-            depths = depths,
-            aabb = octree.aabb,
-            aux = aux,
-        )
-        ret['rgb'] = rgb
-        ret['depth'] = depth
-        ret['alpha'] = alpha
-    
-    return ret
-
-
-class OctreeRenderer:
-    """
-    Renderer for the Voxel representation.
-
-    Args:
-        rendering_options (dict): Rendering options.
-    """
-
-    def __init__(self, rendering_options={}) -> None:
-        try:
-            import diffoctreerast
-        except ImportError:
-            print("\033[93m[WARNING] diffoctreerast is not installed. The renderer will be disabled.\033[0m")
-            self.unsupported = True
-        else:
-            self.unsupported = False
-        
-        self.pipe = edict({
-            "with_distloss": False,
-            "with_aux": False,
-            "scale_modifier": 1.0,
-            "used_rank": None,
-            "jitter": False,
-            "debug": False,
-        })
-        self.rendering_options = edict({
-            "resolution": None,
-            "near": None,
-            "far": None,
-            "ssaa": 1,
-            "bg_color": 'random',
-        })
-        self.halton_sampler = qmc.Halton(2, scramble=False)
-        self.rendering_options.update(rendering_options)
-        self.bg_color = None
-    
-    def render(
-            self,
-            octree: DfsOctree,
-            extrinsics: torch.Tensor,
-            intrinsics: torch.Tensor,
-            colors_overwrite: torch.Tensor = None,
-        ) -> edict:
-        """
-        Render the octree.
-
-        Args:
-            octree (Octree): octree
-            extrinsics (torch.Tensor): (4, 4) camera extrinsics
-            intrinsics (torch.Tensor): (3, 3) camera intrinsics
-            colors_overwrite (torch.Tensor): (N, 3) override color
-
-        Returns:
-            edict containing:
-                color (torch.Tensor): (3, H, W) rendered color
-                depth (torch.Tensor): (H, W) rendered depth
-                alpha (torch.Tensor): (H, W) rendered alpha
-                distloss (Optional[torch.Tensor]): (H, W) rendered distance loss
-                percent_depth (Optional[torch.Tensor]): (H, W) rendered percent depth
-                aux (Optional[edict]): auxiliary tensors
-        """
-        resolution = self.rendering_options["resolution"]
-        near = self.rendering_options["near"]
-        far = self.rendering_options["far"]
-        ssaa = self.rendering_options["ssaa"]
-        
-        if self.unsupported:
-            image = np.zeros((512, 512, 3), dtype=np.uint8)
-            text_bbox = cv2.getTextSize("Unsupported", cv2.FONT_HERSHEY_SIMPLEX, 2, 3)[0]
-            origin = (512 - text_bbox[0]) // 2, (512 - text_bbox[1]) // 2
-            image = cv2.putText(image, "Unsupported", origin, cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3, cv2.LINE_AA)
-            return {
-                'color': torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255,
-            }
-        
-        if self.rendering_options["bg_color"] == 'random':
-            self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda")
-            if np.random.rand() < 0.5:
-                self.bg_color += 1
-        else:
-            self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda")
-
-        if self.pipe["with_aux"]:
-            aux = {
-                'grad_color2': torch.zeros((octree.num_leaf_nodes, 3), dtype=torch.float32, requires_grad=True, device="cuda") + 0,
-                'contributions': torch.zeros((octree.num_leaf_nodes, 1), dtype=torch.float32, requires_grad=True, device="cuda") + 0,
-            }
-            for k in aux.keys():
-                aux[k].requires_grad_()
-                aux[k].retain_grad()
-        else:
-            aux = None
-
-        view = extrinsics
-        perspective = intrinsics_to_projection(intrinsics, near, far)
-        camera = torch.inverse(view)[:3, 3]
-        focalx = intrinsics[0, 0]
-        focaly = intrinsics[1, 1]
-        fovx = 2 * torch.atan(0.5 / focalx)
-        fovy = 2 * torch.atan(0.5 / focaly)
-            
-        camera_dict = edict({
-            "image_height": resolution * ssaa,
-            "image_width": resolution * ssaa,
-            "FoVx": fovx,
-            "FoVy": fovy,
-            "znear": near,
-            "zfar": far,
-            "world_view_transform": view.T.contiguous(),
-            "projection_matrix": perspective.T.contiguous(),
-            "full_proj_transform": (perspective @ view).T.contiguous(),
-            "camera_center": camera
-        })
-
-        # Render
-        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)
-
-        if ssaa > 1:
-            render_ret.rgb = F.interpolate(render_ret.rgb[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()
-            render_ret.depth = F.interpolate(render_ret.depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()
-            render_ret.alpha = F.interpolate(render_ret.alpha[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()
-            if hasattr(render_ret, 'percent_depth'):
-                render_ret.percent_depth = F.interpolate(render_ret.percent_depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()
-
-        ret = edict({
-            'color': render_ret.rgb,
-            'depth': render_ret.depth,
-            'alpha': render_ret.alpha,
-        })
-        if self.pipe["with_distloss"] and 'distloss' in render_ret:
-            ret['distloss'] = render_ret.distloss
-        if self.pipe["with_aux"]:
-            ret['aux'] = aux
-        if hasattr(render_ret, 'percent_depth'):
-            ret['percent_depth'] = render_ret.percent_depth
-        return ret

trellis/renderers/sh_utils.py

@@ -1,118 +0,0 @@
-#  Copyright 2021 The PlenOctree Authors.
-#  Redistribution and use in source and binary forms, with or without
-#  modification, are permitted provided that the following conditions are met:
-#
-#  1. Redistributions of source code must retain the above copyright notice,
-#  this list of conditions and the following disclaimer.
-#
-#  2. Redistributions in binary form must reproduce the above copyright notice,
-#  this list of conditions and the following disclaimer in the documentation
-#  and/or other materials provided with the distribution.
-#
-#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
-#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-#  POSSIBILITY OF SUCH DAMAGE.
-
-import torch
-
-C0 = 0.28209479177387814
-C1 = 0.4886025119029199
-C2 = [
-    1.0925484305920792,
-    -1.0925484305920792,
-    0.31539156525252005,
-    -1.0925484305920792,
-    0.5462742152960396
-]
-C3 = [
-    -0.5900435899266435,
-    2.890611442640554,
-    -0.4570457994644658,
-    0.3731763325901154,
-    -0.4570457994644658,
-    1.445305721320277,
-    -0.5900435899266435
-]
-C4 = [
-    2.5033429417967046,
-    -1.7701307697799304,
-    0.9461746957575601,
-    -0.6690465435572892,
-    0.10578554691520431,
-    -0.6690465435572892,
-    0.47308734787878004,
-    -1.7701307697799304,
-    0.6258357354491761,
-]   
-
-
-def eval_sh(deg, sh, dirs):
-    """
-    Evaluate spherical harmonics at unit directions
-    using hardcoded SH polynomials.
-    Works with torch/np/jnp.
-    ... Can be 0 or more batch dimensions.
-    Args:
-        deg: int SH deg. Currently, 0-3 supported
-        sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
-        dirs: jnp.ndarray unit directions [..., 3]
-    Returns:
-        [..., C]
-    """
-    assert deg <= 4 and deg >= 0
-    coeff = (deg + 1) ** 2
-    assert sh.shape[-1] >= coeff
-
-    result = C0 * sh[..., 0]
-    if deg > 0:
-        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
-        result = (result -
-                C1 * y * sh[..., 1] +
-                C1 * z * sh[..., 2] -
-                C1 * x * sh[..., 3])
-
-        if deg > 1:
-            xx, yy, zz = x * x, y * y, z * z
-            xy, yz, xz = x * y, y * z, x * z
-            result = (result +
-                    C2[0] * xy * sh[..., 4] +
-                    C2[1] * yz * sh[..., 5] +
-                    C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
-                    C2[3] * xz * sh[..., 7] +
-                    C2[4] * (xx - yy) * sh[..., 8])
-
-            if deg > 2:
-                result = (result +
-                C3[0] * y * (3 * xx - yy) * sh[..., 9] +
-                C3[1] * xy * z * sh[..., 10] +
-                C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
-                C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
-                C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
-                C3[5] * z * (xx - yy) * sh[..., 14] +
-                C3[6] * x * (xx - 3 * yy) * sh[..., 15])
-
-                if deg > 3:
-                    result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
-                            C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
-                            C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
-                            C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
-                            C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
-                            C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
-                            C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
-                            C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
-                            C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
-    return result
-
-def RGB2SH(rgb):
-    return (rgb - 0.5) / C0
-
-def SH2RGB(sh):
-    return sh * C0 + 0.5