diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 01983d30f..331a45df8 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -66,6 +66,9 @@ def RGB2SH(rgb): C0 = 0.28209479177387814 return (rgb - 0.5) / C0 +def SH2RGB(sh): + C0 = 0.28209479177387814 + return sh * C0 + 0.5 def projection_matrix(znear, zfar, fovx, fovy, device="cpu"): t = znear * math.tan(0.5 * fovy) @@ -196,8 +199,12 @@ class GaussianSplattingModel(Model): ) @property - def get_colors(self): - return self.colors_all + def colors(self): + return SH2RGB(self.colors_all[:, 0, :]) + + @property + def shs_rest(self): + return self.colors_all[:, 1:, :] def load_state_dict(self, dict, **kwargs): # resize the parameters to match the new number of points @@ -546,13 +553,13 @@ class GaussianSplattingModel(Model): if crop_ids is not None: opacities_crop = self.opacities[crop_ids] means_crop = self.means[crop_ids] - colors_crop = self.get_colors[crop_ids] + colors_crop = self.colors_all[crop_ids] scales_crop = self.scales[crop_ids] quats_crop = self.quats[crop_ids] else: opacities_crop = self.opacities means_crop = self.means - colors_crop = self.get_colors + colors_crop = self.colors_all scales_crop = self.scales quats_crop = self.quats self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = ProjectGaussians.apply( diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 821dd952e..028cddd31 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -502,11 +502,12 @@ class ExportGaussianSplat(Exporter): map_to_tensors["normals"] = o3d.core.Tensor(np.zeros_like(positions), o3d.core.float32) colors = model.colors.data.cpu().numpy() + map_to_tensors["colors"] = (colors * 255).astype(np.uint8) for i in range(colors.shape[1]): - map_to_tensors[f"f_dc_{i}"] = colors[:, i] + map_to_tensors[f"f_dc_{i}"] = colors[:, i:i+1] shs = model.shs_rest.data.cpu().numpy() - if shs.size > 0: + if model.config.sh_degree > 0: shs = shs.reshape((colors.shape[0], -1, 1)) for i in range(shs.shape[-1]): map_to_tensors[f"f_rest_{i}"] = shs[:, i]