[DRAFT] Mip-Splatting implementation #274
Conversation
|
Awesome @niujinshuchong ! I'm thinking if there is a way to elegantly message some of implementation into the gsplat library, e.g., into It doesn't have to be living in the |
|
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 |
| self.splats["filters"] = torch.nn.Parameter(filter_3D) | ||
|
|
||
| @torch.no_grad() | ||
| def compute_3D_smoothing_filter_torch(self): |
There was a problem hiding this comment.
Would you mind move this into gsplat/cuda/_torch_impl.py and put up a test case under test/test_basic.py to protect the CUDA implementation compute_3D_smoothing_filter()?
|
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. |
| @@ -572,6 +572,34 @@ def rasterize_to_indices_in_range( | |||
| return out_gauss_ids, out_pixel_ids, out_camera_ids | |||
|
|
|||
|
|
|||
| def compute_3D_smoothing_filter( | |||
There was a problem hiding this comment.
nit: maybe use 3d instead of 3D in the naming of these functions to align with the naming convention in gsplat?
examples/benchmark_mipnerf360.py
Outdated
There was a problem hiding this comment.
Nice scripts!
- Does
GPUtilneed to be installed via pip? If so we should put it intoexamples/requirements.txt. - If you are going to merge the
simple_trainer_mip_splatting.pywithsimple_trainer.pythen we can just replace the original benchmark script with this one, and enable an input flag to enable mip splatting. But merging the benchmark script means that we probably want the default benchmarking set to be the original 3DGS.
| excluded_gpus = set([]) | ||
|
|
||
| # classic | ||
| result_dir = "results/benchmark_stmt" |
There was a problem hiding this comment.
what does stmt stands for? lol
There was a problem hiding this comment.
I think Single-scale-Training-and-Multi-scale-Testing, the idea being you train with a fixed resolution (here a factor of 1/8 of the original image res) and test on various other resolutions (1x full res, 1/2, 1/4, 1/8).
There was a problem hiding this comment.
how hard would it be to merge this one with the other script? Ideally an list argument could be passed in to specify the rendering factors that the evaluation happens on
There was a problem hiding this comment.
I think we could adapt the current benchmark.sh script into the .py style similar to the one proposed here, allowing for easier batch jobs on e.g. clusters and also streamlining the evaluation of the various new features of gsplat... I can look into it...
examples/show_mipnerf360.py
Outdated
There was a problem hiding this comment.
Nice script! Would love to have it! Maybe the name of this script could be something like "compare_results_mipnerf360" to be more readable? Actually maybe we could also get rid of the _mipnerf360 suffix until when another dataset is supported.
| @@ -0,0 +1,45 @@ | |||
| import json | |||
There was a problem hiding this comment.
actually, another idea of managing these scripts is to have a mip_splatting folder under examples, in which you could have multiple scripts in there, without worrying about merging or conflicting with other scripts. I think this is a scalable way to manage the scripts when more things are supported in the future.
And you could also have a customized simple_trainer.py living in it, basically whatever you need to be customized. But this is less preferred because a customized simple_trainer.py means the mip-splatting feature would be disjoint with other future features for user to play with. So ideally we merge the trainer but could put the scripts into an isolated folder.
|
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. |
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: autonomousvision/mip-splatting#48 |
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:
Here is the benchmark results of training with factor 8 and render with factor 8, 4, 2, 1.
metric is PSNR, SSIM and LPIPS.
Clean up later.