bugfix: apply new config to ply_loader_renderer.py

Author: JasonLSC

examples/ply_loader_renderer.py

@@ -6,7 +6,7 @@
 from contextlib import nullcontext
 from dataclasses import dataclass, field
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Union, ContextManager, TypedDict
+from typing import Dict, List, Optional, Tuple, Union, ContextManager, TypedDict, Any
 
 import imageio
 import nerfview
@@ -83,15 +83,59 @@ def update_schedule(self, **kwargs):
                 setattr(self, key, value)
         self.schedule = self._create_schedule()
 
-class CompressionConfig(TypedDict, total=False):
+@dataclass
+class CodecConfig:
+    encode: str
+    decode: str
+
+@dataclass
+class AttributeCodecs:
+    means: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_png_16bit", "_decompress_png_16bit"))
+    scales: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_factorized_ans", "_decompress_factorized_ans"))
+    quats: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_factorized_ans", "_decompress_factorized_ans"))
+    opacities: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_png", "_decompress_png"))
+    sh0: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_png", "_decompress_png"))
+    shN: CodecConfig = field(default_factory=lambda: CodecConfig("_compress_masked_kmeans", "_decompress_masked_kmeans"))
+    
+    def to_dict(self) -> Dict[str, Dict[str, str]]:
+        return {
+            attr: {"encode": getattr(self, attr).encode, "decode": getattr(self, attr).decode}
+            for attr in ["means", "scales", "quats", "opacities", "sh0", "shN"]
+        }
+
+@dataclass
+class CompressionConfig:
     # Use PLAS sort in compression or not
-    use_sort: Optional[bool] = field(default=None)
+    use_sort: bool = True
     # Verbose or not
-    verbose: Optional[bool] = field(default=None) 
+    verbose: bool = True
     # QP value for video coding
-    qp: int = 4
+    qp: Optional[int] = field(default=None)
     # Number of cluster of VQ for shN compression
     n_clusters: int = 32768
+    # Maps attribute names to their codec functions
+    attribute_codec_registry: Optional[AttributeCodecs] = field(default_factory=lambda: AttributeCodecs())
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the CompressionConfig instance to a dictionary.
+        If attribute_codec_registry is not None, it will be converted to a dictionary using its to_dict method.
+        Fields with None values (use_sort, verbose) will be excluded from the resulting dictionary.
+        """
+        result = {
+            "use_sort": self.use_sort,
+            "verbose": self.verbose,
+            "n_clusters": self.n_clusters,
+        }
+            
+        if self.qp is not None:
+            result["qp"] = self.qp
+        
+        # handle attribute_codec_registry
+        if self.attribute_codec_registry is not None:
+            result["attribute_codec_registry"] = self.attribute_codec_registry.to_dict()
+        
+        return result
 
 @dataclass
 class Config:
@@ -565,9 +609,11 @@ def __init__(
             if cfg.compression == "png":
                 self.compression_method = PngCompression()
             elif  cfg.compression == "entropy_coding":
-                self.compression_method = EntropyCodingCompression()
+                compression_cfg = cfg.compression_cfg.to_dict()
+                self.compression_method = EntropyCodingCompression(**compression_cfg)
             elif cfg.compression == "hevc":
-                self.compression_method = HevcCompression(**cfg.compression_cfg) # compression_cfg=cfg.compression_cfg
+                compression_cfg = cfg.compression_cfg.to_dict()
+                self.compression_method = HevcCompression(**compression_cfg) 
             else:
                 raise ValueError(f"Unknown compression strategy: {cfg.compression}")
 

gsplat/compression/hevc_compression.py

@@ -1,8 +1,8 @@
 import json
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field, InitVar
 import shutil
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional
 
 import numpy as np
 import torch
@@ -46,36 +46,55 @@ class HevcCompression:
 
     use_sort: bool = True
     verbose: bool = True
-    qp: int = 4
     n_clusters: int = 32768
+    qp: int = 4
+
+    attribute_codec_registry: InitVar[Optional[Dict[str, str]]] = None
+
+    compress_fn_map: Dict[str, Callable] = field(default_factory=lambda: {
+        "means": _compress_png_16bit,
+        "scales": _compress_hevc_kbit,
+        "quats": _compress_quats_hevc_kbit,
+        "opacities": _compress_hevc_kbit,
+        "sh0": _compress_hevc_kbit,
+        "shN": _compress_masked_kmeans,
+    })
+    decompress_fn_map: Dict[str, Callable] = field(default_factory=lambda: {
+        "means": _decompress_png_16bit,
+        "scales": _decompress_hevc_kbit,
+        "quats": _decompress_quats_hevc_kbit,
+        "opacities": _decompress_hevc_kbit,
+        "sh0": _decompress_hevc_kbit,
+        "shN": _decompress_masked_kmeans,
+    })
 
     def _get_compress_fn(self, param_name: str) -> Callable:
-        compress_fn_map = {
-            "means": _compress_png_16bit,
-            "scales": _compress_hevc_kbit,
-            "quats": _compress_quats_hevc_kbit,
-            "opacities": _compress_hevc_kbit,
-            "sh0": _compress_hevc_kbit,
-            # "shN": _compress_kmeans,
-            "shN": _compress_masked_kmeans,
-        }
-        if param_name in compress_fn_map:
-            return compress_fn_map[param_name]
+        # compress_fn_map = {
+        #     "means": _compress_png_16bit,
+        #     "scales": _compress_hevc_kbit,
+        #     "quats": _compress_quats_hevc_kbit,
+        #     "opacities": _compress_hevc_kbit,
+        #     "sh0": _compress_hevc_kbit,
+        #     # "shN": _compress_kmeans,
+        #     "shN": _compress_masked_kmeans,
+        # }
+        if param_name in self.compress_fn_map:
+            return self.compress_fn_map[param_name]
         else:
             return _compress_npz
 
     def _get_decompress_fn(self, param_name: str) -> Callable:
-        decompress_fn_map = {
-            "means": _decompress_png_16bit,
-            "scales": _decompress_hevc_kbit,
-            "quats": _decompress_quats_hevc_kbit,
-            "opacities": _decompress_hevc_kbit,
-            "sh0": _decompress_hevc_kbit,
-            # "shN": _decompress_kmeans,
-            "shN": _decompress_masked_kmeans,
-        }
-        if param_name in decompress_fn_map:
-            return decompress_fn_map[param_name]
+        # decompress_fn_map = {
+        #     "means": _decompress_png_16bit,
+        #     "scales": _decompress_hevc_kbit,
+        #     "quats": _decompress_quats_hevc_kbit,
+        #     "opacities": _decompress_hevc_kbit,
+        #     "sh0": _decompress_hevc_kbit,
+        #     # "shN": _decompress_kmeans,
+        #     "shN": _decompress_masked_kmeans,
+        # }
+        if param_name in self.decompress_fn_map:
+            return self.decompress_fn_map[param_name]
         else:
             return _decompress_npz