Compressing gaussian splats

[jefequien]

Thanks largely to MCMC's improved densification strategy, a simple post-training compression scheme is enough to top the 3D Gaussian Splatting Compression Methods leaderboard.
This PR implements quantization, sorting, and spherical harmonic coefficients K-means clustering for compression.

Relevant links:
https://github.com/fraunhoferhhi/Self-Organizing-Gaussians
https://aras-p.info/blog/2023/09/27/Making-Gaussian-Splats-more-smaller/

Results

Evaluated on MipNeRF360 dataset.

Method	PSNR	SSIM	LPIPS	Size (MB)
mcmc (0.5M) w/ compression	28.24	0.852	0.175	8.89
mcmc (1M) w/ compression	28.65	0.864	0.154	16.5
mcmc (2M) w/ compression	29.02	0.872	0.138	31.2

Method	PSNR	SSIM	LPIPS	Size (MB)
mcmc (1M) w/ quant and kmeans	28.65	0.864	0.154	16.5
mcmc (1M) w/ quant only	29.02	0.870	0.146	31.6
mcmc (1M) w/o compression	29.18	0.874	0.142	236

Note: Using K-means centroids and labels to initialize a codebook sets the high score for a given file size, at the cost of another training run. This is not implemented in the PR.

Data format

This PR quantizes means to 16 bits, opacities, quats, scales, sh0 to 8 bits, and shN to 6 bit. Sorting these fields with PLAS and saving them as PNGs saves a few MBs and generates nice visuals.
means are quantized in log space.
shN coefficients are clustered with K-means into 2**16 clusters and their centroids and labels are saved to a compressed npz file.

Bonsai scene's quantized and sorted sh0 png.

[brentyi]

I think all dict needs to be Dict imported from typing.

As an FYI from __future__ import annotations at the top of files will also let you use dict[K, V], tuple[T], etc!

[liruilong940607]

Did some cleanup and put up a doc page for it (sphinx-build docs/source _build). I think I'm happy to see it merged at this stage. The only thing left is to put some evaluation numbers into the evaluation page but I think its an optional. Let me know if all these modifications that I made looks reasonable to you @jefequien

[jefequien]

Typo. shN?

[liruilong940607]

fixed!

[Florian-Barthel]

Great to see Self-Organizing-Gaussians beeing used!

[SharkWipf]

FYI: I was messing around with the code with Claude, and it seems to inadvertently have caught and fixed a bug related to the compression in simple_trainer.py:

It seems correct to me at a glance, but as I am not read in into the code I don't want to open an issue/PR for an accidental Claude change that I haven't looked into, I figured I'd just put it here instead so you can judge it yourself.

[ichsan2895]

Hello, sorry for late party. I have used this method in gsplat==1.3.0. But I get the compression format is PNG instead of PLY.

I must decompress it, so I can convert it to PLY. But, unfortunatelly, the size of PLY is remain same.

The question is, any tutorial for getting compressed PLY with this method?

[ichsan2895]

FYI: I was messing around with the code with Claude, and it seems to inadvertently have caught and fixed a bug related to the compression in simple_trainer.py:

It seems correct to me at a glance, but as I am not read in into the code I don't want to open an issue/PR for an accidental Claude change that I haven't looked into, I figured I'd just put it here instead so you can judge it yourself.

Yeah, I got this error too. It does not error somehow in my setup with 1 GPU, but if it has more >1 GPU, it caused error

It correct code should be

compress_dir = f"{self.cfg.result_dir}/compression/rank{world_rank}"

[Ben-Mack]

Hi @jefequien, are there anyway to load Checkpoint from Nerfstudio and run just the compression step, instead of training from ground up in gsplat?

[jefequien]

@Ben-Mack Assuming Nerfstudio uses the same .pt file as gsplat, it should be simple. Just use simple_trainer.py with --ckpt

gsplat/examples/benchmarks/compression/mcmc.sh

Line 39 in cc2fb4c

CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --disable_viewer --data_factor $DATA_FACTOR \