Spaces:
Running on Zero
Running on Zero
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from einops import einsum, rearrange, repeat | |
| from jaxtyping import Float | |
| from plyfile import PlyData, PlyElement | |
| from scipy.spatial.transform import Rotation as R | |
| from torch import Tensor | |
| from ..misc.sh_rotation import rotate_sh | |
| def construct_list_of_attributes(num_rest: int) -> list[str]: | |
| attributes = ["x", "y", "z", "nx", "ny", "nz"] | |
| for i in range(3): | |
| attributes.append(f"f_dc_{i}") | |
| for i in range(num_rest): | |
| attributes.append(f"f_rest_{i}") | |
| attributes.append("opacity") | |
| for i in range(3): | |
| attributes.append(f"scale_{i}") | |
| for i in range(4): | |
| attributes.append(f"rot_{i}") | |
| return attributes | |
| def export_ply( | |
| extrinsics: Float[Tensor, "4 4"], | |
| means: Float[Tensor, "gaussian 3"], | |
| scales: Float[Tensor, "gaussian 3"], | |
| rotations: Float[Tensor, "gaussian 4"], | |
| harmonics: Float[Tensor, "gaussian 3 d_sh"], | |
| opacities: Float[Tensor, " gaussian"], | |
| path: Path, | |
| align_to_view: bool = False, | |
| ): | |
| path = Path(path) | |
| with torch.autocast(device_type=means.device.type, enabled=False): | |
| means = means.float() | |
| scales = scales.float() | |
| rotations = rotations.float() | |
| harmonics = harmonics.float() | |
| opacities = opacities.float() | |
| extrinsics = extrinsics.to(means) | |
| rotations = rotations / rotations.norm(dim=-1, keepdim=True).clamp_min(1e-8) | |
| if align_to_view: | |
| view_rotation = extrinsics[:3, :3].transpose(-1, -2) | |
| view_origin = extrinsics[:3, 3] | |
| means = einsum( | |
| view_rotation, | |
| means - view_origin, | |
| "i j, ... j -> ... i", | |
| ) | |
| harmonics = rotate_sh(harmonics, view_rotation) | |
| rotations = R.from_quat(rotations.detach().cpu().numpy()).as_matrix() | |
| rotations = view_rotation.detach().cpu().numpy() @ rotations | |
| rotations = R.from_matrix(rotations).as_quat() | |
| else: | |
| rotations = rotations.detach().cpu().numpy() | |
| x, y, z, w = rearrange(rotations, "g xyzw -> xyzw g") | |
| rotations = np.stack((w, x, y, z), axis=-1) | |
| means = means.detach().cpu().numpy() | |
| normals = np.zeros_like(means) | |
| harmonics_dc = harmonics[..., 0].detach().cpu().contiguous().numpy() | |
| harmonics_rest = ( | |
| harmonics[..., 1:].flatten(start_dim=1).detach().cpu().contiguous().numpy() | |
| ) | |
| logit_opacities = ( | |
| torch.logit(opacities[..., None].clamp(1e-6, 1 - 1e-6)) | |
| .detach() | |
| .cpu() | |
| .numpy() | |
| ) | |
| log_scales = scales.clamp_min(1e-12).log().detach().cpu().numpy() | |
| num_rest = 3 * (harmonics.shape[-1] - 1) | |
| dtype_full = [ | |
| (attribute, "f4") for attribute in construct_list_of_attributes(num_rest) | |
| ] | |
| elements = np.empty(means.shape[0], dtype=dtype_full) | |
| attributes = ( | |
| means, | |
| normals, | |
| harmonics_dc, | |
| harmonics_rest, | |
| logit_opacities, | |
| log_scales, | |
| rotations, | |
| ) | |
| attributes = np.concatenate(attributes, axis=1) | |
| elements[:] = list(map(tuple, attributes)) | |
| path.parent.mkdir(exist_ok=True, parents=True) | |
| PlyData([PlyElement.describe(elements, "vertex")]).write(path) | |
| def infer_gaussian_grid( | |
| num_gaussians: int, | |
| num_views: int, | |
| image_height: int, | |
| image_width: int, | |
| ) -> tuple[int, int, int]: | |
| if num_gaussians % num_views != 0: | |
| raise ValueError( | |
| f"Cannot split {num_gaussians} Gaussians over {num_views} views." | |
| ) | |
| gaussians_per_view = num_gaussians // num_views | |
| for downsample in (1, 2, 4, 8, 16, 32): | |
| if image_height % downsample != 0 or image_width % downsample != 0: | |
| continue | |
| height = image_height // downsample | |
| width = image_width // downsample | |
| pixels = height * width | |
| if pixels > 0 and gaussians_per_view % pixels == 0: | |
| return height, width, gaussians_per_view // pixels | |
| raise ValueError( | |
| "Could not infer Gaussian grid layout from " | |
| f"num_gaussians={num_gaussians}, num_views={num_views}, " | |
| f"image_shape=({image_height}, {image_width})." | |
| ) | |
| def select_gaussians( | |
| tensor: Tensor, | |
| mask: Tensor, | |
| num_views: int, | |
| height: int, | |
| width: int, | |
| points_per_pixel: int, | |
| ) -> Tensor: | |
| tensor = rearrange( | |
| tensor, | |
| "() (v h w spp) ... -> h w spp v ...", | |
| v=num_views, | |
| h=height, | |
| w=width, | |
| spp=points_per_pixel, | |
| ) | |
| return tensor[mask] | |
| def local_to_world_rotations( | |
| rotations: Tensor, | |
| context_extrinsics: Tensor, | |
| num_views: int, | |
| height: int, | |
| width: int, | |
| points_per_pixel: int, | |
| mask: Tensor | None, | |
| ) -> Tensor: | |
| if mask is None: | |
| local_rotations = rotations[0] | |
| c2w_rotations = repeat( | |
| context_extrinsics[0, :, :3, :3], | |
| "v i j -> (v h w spp) i j", | |
| h=height, | |
| w=width, | |
| spp=points_per_pixel, | |
| ) | |
| else: | |
| local_rotations = select_gaussians( | |
| rotations, | |
| mask, | |
| num_views, | |
| height, | |
| width, | |
| points_per_pixel, | |
| ) | |
| c2w_rotations = repeat( | |
| context_extrinsics[0, :, :3, :3], | |
| "v i j -> h w spp v i j", | |
| h=height, | |
| w=width, | |
| spp=points_per_pixel, | |
| )[mask] | |
| local_rotation_mats = R.from_quat( | |
| local_rotations.detach().cpu().numpy() | |
| ).as_matrix() | |
| world_rotation_mats = c2w_rotations.detach().cpu().numpy() @ local_rotation_mats | |
| world_rotations = R.from_matrix(world_rotation_mats).as_quat() | |
| return torch.from_numpy(world_rotations).to( | |
| device=rotations.device, | |
| dtype=rotations.dtype, | |
| ) | |
| def save_gaussian_ply( | |
| gaussians, | |
| visualization_dump, | |
| example, | |
| save_path, | |
| save_all_gaussians: bool = True, | |
| align_to_view: bool = False, | |
| ): | |
| del visualization_dump | |
| v, _, image_height, image_width = example["context"]["image"].shape[1:] | |
| h, w, points_per_pixel = infer_gaussian_grid( | |
| gaussians.means.shape[1], | |
| v, | |
| image_height, | |
| image_width, | |
| ) | |
| if gaussians.scales is None or gaussians.rotations is None: | |
| raise ValueError("PLY export requires Gaussians.scales and Gaussians.rotations.") | |
| if save_all_gaussians: | |
| mask = None | |
| selected_means = gaussians.means[0] | |
| selected_scales = gaussians.scales[0] | |
| selected_harmonics = gaussians.harmonics[0] | |
| selected_opacities = gaussians.opacities[0] | |
| else: | |
| mask = torch.zeros( | |
| (h, w, points_per_pixel, v), | |
| dtype=torch.bool, | |
| device=gaussians.means.device, | |
| ) | |
| gaussian_trim = 8 | |
| mask[gaussian_trim:-gaussian_trim, gaussian_trim:-gaussian_trim, :, :] = True | |
| selected_means = select_gaussians( | |
| gaussians.means, mask, v, h, w, points_per_pixel | |
| ) | |
| selected_scales = select_gaussians( | |
| gaussians.scales, mask, v, h, w, points_per_pixel | |
| ) | |
| selected_harmonics = select_gaussians( | |
| gaussians.harmonics, mask, v, h, w, points_per_pixel | |
| ) | |
| selected_opacities = select_gaussians( | |
| gaussians.opacities, mask, v, h, w, points_per_pixel | |
| ) | |
| world_rotations = local_to_world_rotations( | |
| gaussians.rotations, | |
| example["context"]["extrinsics"], | |
| v, | |
| h, | |
| w, | |
| points_per_pixel, | |
| mask, | |
| ) | |
| export_ply( | |
| example["context"]["extrinsics"][0, v // 2], | |
| selected_means, | |
| selected_scales, | |
| world_rotations, | |
| selected_harmonics, | |
| selected_opacities, | |
| save_path, | |
| align_to_view=align_to_view, | |
| ) | |