Spaces:
Running on Zero
Running on Zero
| # | |
| # 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 george.drettakis@inria.fr | |
| # | |
| 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, pipe, bg_color: torch.Tensor, scaling_modifier=1.0, override_color=None): | |
| # lazy import | |
| if "rasterization" not in globals(): | |
| from gsplat import rasterization | |
| tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) | |
| tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) | |
| focal_length_x = viewpoint_camera.image_width / (2 * tanfovx) | |
| focal_length_y = viewpoint_camera.image_height / (2 * tanfovy) | |
| K = torch.tensor( | |
| [ | |
| [focal_length_x, 0, viewpoint_camera.image_width / 2.0], | |
| [0, focal_length_y, viewpoint_camera.image_height / 2.0], | |
| [0, 0, 1], | |
| ], | |
| device=pc.get_xyz.device, | |
| dtype=torch.float32, | |
| ) | |
| means3D = pc.get_xyz | |
| opacity = pc.get_opacity | |
| scales = pc.get_scaling * scaling_modifier | |
| rotations = pc.get_rotation | |
| if override_color is not None: | |
| colors = override_color # [N, 3] | |
| sh_degree = None | |
| else: | |
| colors = pc.get_features # [N, K, 3] | |
| sh_degree = pc.active_sh_degree | |
| viewmat = viewpoint_camera.world_view_transform.transpose(0, 1) | |
| render_colors, render_alphas, info = rasterization( | |
| means=means3D, # [N, 3] | |
| quats=rotations, # [N, 4] | |
| scales=scales, # [N, 3] | |
| opacities=opacity.squeeze(-1), # [N] | |
| colors=colors, | |
| viewmats=viewmat[None], # [1, 4, 4] | |
| Ks=K[None], # [1, 3, 3] | |
| backgrounds=bg_color[None], | |
| width=int(viewpoint_camera.image_width), | |
| height=int(viewpoint_camera.image_height), | |
| packed=False, | |
| sh_degree=sh_degree, | |
| rasterize_mode='antialiased' | |
| ) | |
| rendered_image = render_colors[0].permute(2, 0, 1) | |
| radii = info["radii"].squeeze(0) | |
| try: | |
| info["means2d"].retain_grad() | |
| except Exception: | |
| pass | |
| return edict({ | |
| "render": rendered_image, | |
| "viewspace_points": info["means2d"], | |
| "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 | |