gsplat icon indicating copy to clipboard operation
gsplat copied to clipboard

[DRAFT] Mip-Splatting implementation

Open niujinshuchong opened this issue 1 year ago • 9 comments

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. Screenshot from 2024-07-09 17-12-26

Screenshot from 2024-07-09 17-13-27 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.

niujinshuchong avatar Jul 09 '24 15:07 niujinshuchong

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

liruilong940607 avatar Jul 09 '24 18:07 liruilong940607

I'll come back to it later with a closer look. Great job!

liruilong940607 avatar Jul 09 '24 18:07 liruilong940607

@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.

niujinshuchong avatar Jul 12 '24 13:07 niujinshuchong

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.

liruilong940607 avatar Jul 12 '24 18:07 liruilong940607

What's the plan on merging this PR? I am particularly interested in this feature (esp. the 3D filter part).

jb-ye avatar Aug 15 '24 16:08 jb-ye

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.

liruilong940607 avatar Aug 15 '24 17:08 liruilong940607

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.

jb-ye avatar Aug 15 '24 19:08 jb-ye

@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 avatar Aug 27 '24 07:08 niujinshuchong

@niujinshuchong We found one issue with the mip-splat implementation. See discussions: https://github.com/autonomousvision/mip-splatting/issues/48

jb-ye avatar Sep 05 '24 15:09 jb-ye