Spaces:
Running
on
Zero
Running
on
Zero
| import gsplat as gs | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from easydict import EasyDict as edict | |
| class GSplatRenderer: | |
| 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, | |
| "use_mip_gaussian": True | |
| }) | |
| 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, | |
| gaussian, | |
| extrinsics: torch.Tensor, | |
| intrinsics: torch.Tensor, | |
| colors_overwrite: torch.Tensor = None | |
| ) -> edict: | |
| resolution = self.rendering_options["resolution"] | |
| 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" | |
| ) | |
| height = resolution * ssaa | |
| width = resolution * ssaa | |
| # Set up background color | |
| 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" | |
| ) | |
| Ks_scaled = intrinsics.clone() | |
| Ks_scaled[0, 0] *= width | |
| Ks_scaled[1, 1] *= height | |
| Ks_scaled[0, 2] *= width | |
| Ks_scaled[1, 2] *= height | |
| Ks_scaled = Ks_scaled.unsqueeze(0) | |
| near_plane = 0.01 | |
| far_plane = 1000.0 | |
| # Rasterize with gsplat | |
| render_colors, render_alphas, meta = gs.rasterization( | |
| means=gaussian.get_xyz, | |
| quats=F.normalize(gaussian.get_rotation, dim=-1), | |
| scales=gaussian.get_scaling / intrinsics[0, 0], | |
| opacities=gaussian.get_opacity.squeeze(-1), | |
| colors=colors_overwrite.unsqueeze(0) if colors_overwrite is not None else torch.sigmoid( | |
| gaussian.get_features.squeeze(1)).unsqueeze(0), | |
| viewmats=extrinsics.unsqueeze(0), | |
| Ks=Ks_scaled, | |
| width=width, | |
| height=height, | |
| near_plane=near_plane, | |
| far_plane=far_plane, | |
| radius_clip=3.0, | |
| eps2d=0.3, | |
| render_mode="RGB", | |
| backgrounds=self.bg_color.unsqueeze(0), | |
| camera_model="pinhole" | |
| ) | |
| rendered_image = render_colors[0, ..., 0:3].permute(2, 0, 1) | |
| # Apply supersampling if needed | |
| if ssaa > 1: | |
| rendered_image = F.interpolate( | |
| rendered_image[None], | |
| size=(resolution, resolution), | |
| mode='bilinear', | |
| align_corners=False, | |
| antialias=True | |
| ).squeeze() | |
| return edict({'color': rendered_image}) |