About Compression Algorithms

[LaFeuilleMorte]

Hi, I've read this blog :https://aras-p.info/blog/2023/09/27/Making-Gaussian-Splats-more-smaller/ about the compression paradigm of gaussian compression. Do you have any benchmark on the compression ratio? As far as I know, the compression algorithms in current version of supersplats have discarded the SH coefficients (which accounts for 75% of the bytes that a single splat has). And the compression ratio seemed to be only 1:4. After a throughly research on compression methods. I found this paper may have some clues:
https://fraunhoferhhi.github.io/Self-Organizing-Gaussians/

And there's a simpler one to implement this method:
https://github.com/nerfstudio-project/gsplat/blob/dd66cbd597f376f103e90de8b6265087b482eac1/gsplat/compression/png_compression.py#L75

From my experiment: this method will keep the SH coefficients, and achieve an 1000MB to 46 MB (1:21) compression ratio.

But there might be some issue implementing this method to supersplat:

1. The compression method depends on Plas to sort gaussians and https://github.com/DeMoriarty/TorchPQ?tab=readme-ov-file#install to calculate KMeans to compress SH coefficients. And the compression took quite long on my GPU. So the compression algorithm maybe very hard to implement in TS environment.

2. However, luckily, the decompression method has got rid of these dependencies. We can use just arrays to implement these calculation.

The compressed results are multiple files includes PNG and NPZ

[LaFeuilleMorte]

I just cleaned some code and only keep the decompression code:

import os
import json
from torch import Tensor
from typing import Any, Callable, Dict
import torch
import numpy as np

def inverse_log_transform(y):
    return torch.sign(y) * (torch.expm1(torch.abs(y)))

def _decompress_png_16bit(
    compress_dir: str, param_name: str, meta: Dict[str, Any]
) -> Tensor:
    """Decompress parameters from PNG files.

    Args:
        compress_dir (str): compression directory
        param_name (str): parameter field name
        meta (Dict[str, Any]): metadata

    Returns:
        Tensor: parameters
    """
    import imageio.v2 as imageio

    if not np.all(meta["shape"]):
        params = torch.zeros(meta["shape"], dtype=getattr(torch, meta["dtype"]))
        return meta

    img_l = imageio.imread(os.path.join(compress_dir, f"{param_name}_l.png"))
    img_u = imageio.imread(os.path.join(compress_dir, f"{param_name}_u.png"))
    img_u = img_u.astype(np.uint16)
    img = (img_u << 8) + img_l

    img_norm = img / (2**16 - 1)
    grid_norm = torch.tensor(img_norm)
    mins = torch.tensor(meta["mins"])
    maxs = torch.tensor(meta["maxs"])
    grid = grid_norm * (maxs - mins) + mins

    params = grid.reshape(meta["shape"])
    params = params.to(dtype=getattr(torch, meta["dtype"]))
    return params


def _decompress_npz(compress_dir: str, param_name: str, meta: Dict[str, Any]) -> Tensor:
    """Decompress parameters with numpy's NPZ compression."""
    arr = np.load(os.path.join(compress_dir, f"{param_name}.npz"))["arr"]
    params = torch.tensor(arr)
    params = params.reshape(meta["shape"])
    params = params.to(dtype=getattr(torch, meta["dtype"]))
    return params


def _decompress_png(compress_dir: str, param_name: str, meta: Dict[str, Any]) -> Tensor:
    """Decompress parameters from PNG file.

    Args:
        compress_dir (str): compression directory
        param_name (str): parameter field name
        meta (Dict[str, Any]): metadata

    Returns:
        Tensor: parameters
    """
    import imageio.v2 as imageio

    if not np.all(meta["shape"]):
        params = torch.zeros(meta["shape"], dtype=getattr(torch, meta["dtype"]))
        return meta

    img = imageio.imread(os.path.join(compress_dir, f"{param_name}.png"))
    img_norm = img / (2**8 - 1)

    grid_norm = torch.tensor(img_norm)
    mins = torch.tensor(meta["mins"])
    maxs = torch.tensor(meta["maxs"])
    grid = grid_norm * (maxs - mins) + mins

    params = grid.reshape(meta["shape"])
    params = params.to(dtype=getattr(torch, meta["dtype"]))
    return params



def _decompress_kmeans(
    compress_dir: str, param_name: str, meta: Dict[str, Any], **kwargs
) -> Tensor:
    """Decompress parameters from K-means compression.

    Args:
        compress_dir (str): compression directory
        param_name (str): parameter field name
        meta (Dict[str, Any]): metadata

    Returns:
        Tensor: parameters
    """
    if not np.all(meta["shape"]):
        params = torch.zeros(meta["shape"], dtype=getattr(torch, meta["dtype"]))
        return meta

    npz_dict = np.load(os.path.join(compress_dir, f"{param_name}.npz"))
    centroids_quant = npz_dict["centroids"]
    labels = npz_dict["labels"]

    centroids_norm = centroids_quant / (2 ** meta["quantization"] - 1)
    centroids_norm = torch.tensor(centroids_norm)
    mins = torch.tensor(meta["mins"])
    maxs = torch.tensor(meta["maxs"])
    centroids = centroids_norm * (maxs - mins) + mins

    params = centroids[labels]
    params = params.reshape(meta["shape"])
    params = params.to(dtype=getattr(torch, meta["dtype"]))
    return params

def _get_decompress_fn(param_name: str) -> Callable:
        decompress_fn_map = {
            "means": _decompress_png_16bit,
            "scales": _decompress_png,
            "quats": _decompress_png,
            "opacities": _decompress_png,
            "sh0": _decompress_png,
            "shN": _decompress_kmeans,
        }
        if param_name in decompress_fn_map:
            return decompress_fn_map[param_name]
        else:
            return _decompress_npz

def decompress(compress_dir: str) -> Dict[str, Tensor]:
    """Run decompression

    Args:
        compress_dir (str): directory that contains compressed files

    Returns:
        Dict[str, Tensor]: decompressed Gaussian splats
    """
    with open(os.path.join(compress_dir, "meta.json"), "r") as f:
        meta = json.load(f)

    splats = {}
    for param_name, param_meta in meta.items():
        decompress_fn = _get_decompress_fn(param_name)
        splats[param_name] = decompress_fn(compress_dir, param_name, param_meta)

    # Param-specific postprocessing
    splats["means"] = inverse_log_transform(splats["means"])
    return splats

if __name__ == "__main__":
    compress_dir = ""
    splats = decompress(compress_dir)

There're some array calculation using numpy and torch. Is there any tensor calculation framework in typescript?

[slimbuck]

Hi @LaFeuilleMorte,

I actually submitted a PR to the PlayCanvas engine which adds support for compressed spherical harmonics see here. (The PR for adding compression to SS is coming soon).

And the compression ratio seemed to be only 1:4

No our compression is better than that. You can read more details here. With the new spherical harmonic compression the 1.5GB bicycle scene comes down to 198MB (99MB without SH).

I haven't looked at the paper you link yet, but from the summary it looks like a train-time compression. The results look awesome, but what happens when you load that scene and edit it? Or combine it with another scene - presumably you must retrain?

There is currently no cuda implementation in the browser (so existing codebases won't easily run in the browser), but actually as we transition to WebGPU (which includes compute shaders) these techniques are going to become much more applicable. (Though in practise I expect we'd have to reimplement everything we want from scratch).

[LaFeuilleMorte]

Hi @LaFeuilleMorte,

I actually submitted a PR to the PlayCanvas engine which adds support for compressed spherical harmonics see here. (The PR for adding compression to SS is coming soon).

And the compression ratio seemed to be only 1:4

No our compression is better than that. You can read more details here. With the new spherical harmonic compression the 1.5GB bicycle scene comes down to 198MB (99MB without SH).

I haven't looked at the paper you link yet, but from the summary it looks like a train-time compression. The results look awesome, but what happens when you load that scene and edit it? Or combine it with another scene - presumably you must retrain?

There is currently no cuda implementation in the browser (so existing codebases won't easily run in the browser), but actually as we transition to WebGPU (which includes compute shaders) these techniques are going to become much more applicable. (Though in practise I expect we'd have to reimplement everything we want from scratch).

Yeah, this compression will train a 2D grid to represent the gaussian properties. So It cannot be directly used in cases like merging or editing scenes. And I'm so exited to try your newest compression algorithm!

[slimbuck]

Thanks for this additional context @LaFeuilleMorte . Let's create a ticket when we're closer to implementing such a beast :)