Method encode

gsplat/compression/post_training/codec.py:106–223

(self, frame_dir: Path, splats: TensorDict, quant_ctx: QuantizationContext)

        self.config = config

    def encode(self, frame_dir: Path, splats: TensorDict, quant_ctx: QuantizationContext) -> CodecArtifacts:
        frame_dir.mkdir(parents=True, exist_ok=True)
        files: Dict[str, Dict[str, Any]] = {}

        for name, tensor in splats.items():
            stats = quant_ctx.field_stats.get(name)
            if stats is None:
                raise KeyError(f"Quantization stats missing for field '{name}'")

            method = getattr(stats, "method", "scalar")
            codebook_size = stats.codebook_shape[0] if getattr(stats, "codebook_shape", None) else 0
            logger.info(
                "PNG codec encoding field '%s'; method=%s bitwidth=%d channels=%d codebook_size=%d",
                name,
                method,
                stats.bitwidth,
                stats.channels,
                codebook_size,
            )
            int_tensor = quant_ctx.int_values.get(name)
            if int_tensor is None:
                raise ValueError(
                    f"Field '{name}' requires integer cache; set store_as_int=True in QuantFieldConfig."
                )

            n_points = int_tensor.shape[0]
            sidelen = int(math.isqrt(n_points))
            if sidelen * sidelen != n_points:
                raise ValueError(f"PNG codec requires perfect-square length after mapping, got {n_points} for '{name}'.")

            flattened = int_tensor.reshape(n_points, -1)
            channels = flattened.shape[1]

            if stats.bitwidth > 16:
                raise ValueError(f"PNG codec only supports up to 16-bit data, got {stats.bitwidth} for '{name}'.")

            # Skip fields with 0 channels (empty fields)
            if channels == 0:
                logger.info(f"Skipping field '{name}' with 0 channels")
                continue

            levels = (1 << stats.bitwidth) - 1
            clamped = torch.clamp(flattened, 0, levels)
            array = clamped.detach().cpu().numpy()

            data_shape = list(stats.tensor_shape or tensor.shape)

            if name == "means" and method == "scalar":
                array16 = array.astype(np.uint16)
                image16 = array16.reshape(sidelen, sidelen, channels)
                low = (image16 & 0xFF).astype(np.uint8)
                high = (image16 >> 8).astype(np.uint8)
                low_path = frame_dir / f"{name}_l.png"
                high_path = frame_dir / f"{name}_h.png"
                imageio.imwrite(low_path, low)
                imageio.imwrite(high_path, high)
                files[name] = {
                    "type": "png_split16",