Minor fixes to PNGCompression (#817)

* Add small value to avoid numerical issues

* Using predefined sort keys

- When 3DGS is extended with additional properties, e.g. appearance embeddings, the sorting can be erroneous. In this case, we sort only predefined attributes for 3DGS (means, quats, scales, opacities, and sh0 if appearance embedding is not used in the model).

---------

Co-authored-by: Chenyu <[email protected]>

Author: AIBluefisher

gsplat/compression/png_compression.py

@@ -325,6 +325,7 @@ def _compress_kmeans(
     params: Tensor,
     n_clusters: int = 65536,
     quantization: int = 6,
+    eps: float = 1e-6,
     verbose: bool = True,
     **kwargs,
 ) -> Dict[str, Any]:
@@ -339,6 +340,7 @@ def _compress_kmeans(
         params (Tensor): parameters to compress
         n_clusters (int): number of K-means clusters
         quantization (int): number of bits in quantization
+        eps (float, optional): small value to avoid numerical issues. Default to 1e-6.
         verbose (bool, optional): Whether to print verbose information. Default to True.
 
     Returns:
@@ -364,7 +366,7 @@ def _compress_kmeans(
     labels = labels.detach().cpu().numpy()
     centroids = kmeans.centroids.permute(1, 0)
 
-    mins = torch.min(centroids)
+    mins = torch.min(centroids) + eps
     maxs = torch.max(centroids)
     centroids_norm = (centroids - mins) / (maxs - mins)
     centroids_norm = centroids_norm.detach().cpu().numpy()

gsplat/compression/sort.py

@@ -29,7 +29,10 @@ def sort_splats(splats: Dict[str, Tensor], verbose: bool = True) -> Dict[str, Te
     n_sidelen = int(n_gs**0.5)
     assert n_sidelen**2 == n_gs, "Must be a perfect square"
 
-    sort_keys = [k for k in splats if k != "shN"]
+    sort_keys = ["means", "quats", "scales", "opacities"]
+    if "sh0" in splats:
+        sort_keys.append("sh0")
+    
     params_to_sort = torch.cat([splats[k].reshape(n_gs, -1) for k in sort_keys], dim=-1)
     shuffled_indices = torch.randperm(
         params_to_sort.shape[0], device=params_to_sort.device