[DRAFT] Mip-Splatting implementation
Implementation of Mip-Splatting. Tested on bicycle scenes and works OK.
Here is a comparison when trained with factor 4 and render with factor 1.
Left is from gsplat's default setting and right is from this PR.
Benchmark results on 7 scenes of Mip-NeRF 360 dataset:
| PSNR | SSIM | LPIPS | num_GS (M) | |
|---|---|---|---|---|
| classic | 28.989 | 0.871 | 0.136 | 3.23 |
| antialiased | 29.043 | 0.871 | 0.139 | 3.37 |
| mip-splatting (this PR) | 29.051 | 0.872 | 0.138 | 3.16 |
| mip-splatting (this PR) + cuda implementation of compute_3D_smoothing_filter | 29.017 | 0.873 | 0.139 | 3.14 |
Here is the benchmark results of training with factor 8 and render with factor 8, 4, 2, 1.
metric is PSNR, SSIM and LPIPS.
| 1x | 2x | 4x | 8x | avg. | 1x | 2x | 4x | 8x | avg. | 1x | 2x | 4x | 8x | avg. | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| classic | 30.08 | 23.14 | 20.34 | 19.18 | 23.18 | 0.912 | 0.722 | 0.597 | 0.605 | 0.709 | 0.049 | 0.169 | 0.327 | 0.470 | 0.254 |
| antialiased | 29.77 | 24.69 | 22.39 | 21.65 | 24.63 | 0.904 | 0.736 | 0.605 | 0.588 | 0.708 | 0.056 | 0.174 | 0.304 | 0.438 | 0.243 |
| mip-splatting (this PR) | 30.29 | 26.84 | 26.07 | 26.14 | 27.33 | 0.916 | 0.799 | 0.748 | 0.776 | 0.810 | 0.050 | 0.192 | 0.309 | 0.428 | 0.245 |
Clean up later.
Awesome @niujinshuchong !
I'm thinking if there is a way to elegantly message some of implementation into the gsplat library, e.g., into rasterization() function so that the downstream users (e.g., nerfstudio) can just import and use.
It doesn't have to be living in the rasterization() function if it does not fit there. We could also put some implementation in to like gsplat/mip.py, so that ppl can use it with:
from gsplat.mip import ...
# and do their thing
I'll come back to it later with a closer look. Great job!
@liruilong940607 Thanks for the suggestions. I added a cuda implementation of compute_3D_smoothing_filter and it can be imported from gsplat. The benchmark results (see top) are reasonable. For now, I create a new script simple_trainer_mip_splatting.py and it almost the same as simple_trainer.py, do you think we should merge them or make it separate? I also added a training script to use multiple gpus.
Thanks for the benchmarking! Really nice! Yeah from a quick going through of the training script, it seems to be very compatible with the original one so I would suggest we merge the two! Just need a flag in the config, e.g.(--mip_splatting?) and a couple if else in the code.
What's the plan on merging this PR? I am particularly interested in this feature (esp. the 3D filter part).
I think this PR is just about cleaning up, add doc strings and update webpage. It would be nice to finish it up.
Not sure if @niujinshuchong still has time for it. But I might be able to find some cycles in the next two weeks to help with finishing up this PR.
I think this PR is just about cleaning up, add doc strings and update webpage. It would be nice to finish it up.
Not sure if @niujinshuchong still has time for it. But I might be able to find some cycles in the next two weeks to help with finishing up this PR.
Based on my experience with antialiased mode rasterization, I think the 3D filter should be always enabled for antialiased rasterization. Otherwise, a significant percentage of gaussians could become very thin and flat "2D" gaussians, which is often undesirable for modeling non-lambertian effect.
Also with 3D filter, one usually get less number of gaussians at the same quality post training.
@liruilong940607 Yes, it just need to be clean up. I will find some time to continue. However, it would be great if you can help with it. Sorry for the late reply.
@niujinshuchong We found one issue with the mip-splat implementation. See discussions: https://github.com/autonomousvision/mip-splatting/issues/48