├── README.md ├── custom_mc ├── _marching_cubes_lewiner.py ├── _marching_cubes_lewiner_cy.cpp ├── _marching_cubes_lewiner_cy.cpython-36m-x86_64-linux-gnu.so ├── _marching_cubes_lewiner_cy.pyx ├── _marching_cubes_lewiner_luts.py └── setup.py ├── evaluation └── clean_eval_dtu_mesh.py ├── extract_mesh.py ├── gaussian_splatting ├── LICENSE.md ├── arguments │ ├── __init__.py │ └── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ └── __init__.cpython-38.pyc ├── convert.py ├── environment.yml ├── full_eval.py ├── gaussian_renderer │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-38.pyc │ │ └── network_gui.cpython-38.pyc │ └── network_gui.py ├── lpipsPyTorch │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-38.pyc │ └── modules │ │ ├── __pycache__ │ │ ├── lpips.cpython-38.pyc │ │ ├── networks.cpython-38.pyc │ │ └── utils.cpython-38.pyc │ │ ├── lpips.py │ │ ├── networks.py │ │ └── utils.py ├── metrics.py ├── render.py ├── scene │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── cameras.cpython-310.pyc │ │ ├── cameras.cpython-38.pyc │ │ ├── colmap_loader.cpython-310.pyc │ │ ├── colmap_loader.cpython-38.pyc │ │ ├── dataset_readers.cpython-310.pyc │ │ ├── dataset_readers.cpython-38.pyc │ │ ├── gaussian_model.cpython-310.pyc │ │ └── gaussian_model.cpython-38.pyc │ ├── cameras.py │ ├── colmap_loader.py │ ├── dataset_readers.py │ └── gaussian_model.py ├── submodules │ ├── diff-gaussian-rasterization │ │ ├── .gitignore │ │ ├── .gitmodules │ │ ├── CMakeLists.txt │ │ ├── LICENSE.md │ │ ├── README.md │ │ ├── cuda_rasterizer │ │ │ ├── auxiliary.h │ │ │ ├── backward.cu │ │ │ ├── backward.h │ │ │ ├── config.h │ │ │ ├── forward.cu │ │ │ ├── forward.h │ │ │ ├── rasterizer.h │ │ │ ├── rasterizer_impl.cu │ │ │ └── rasterizer_impl.h │ │ ├── diff_gaussian_rasterization │ │ │ └── __init__.py │ │ ├── ext.cpp │ │ ├── rasterize_points.cu │ │ ├── rasterize_points.h │ │ ├── setup.py │ │ └── third_party │ │ │ └── stbi_image_write.h │ └── simple-knn │ │ ├── ext.cpp │ │ ├── setup.py │ │ ├── simple_knn.cu │ │ ├── simple_knn.h │ │ ├── simple_knn │ │ └── .gitkeep │ │ ├── spatial.cu │ │ └── spatial.h ├── train.py └── utils │ ├── __pycache__ │ ├── camera_utils.cpython-310.pyc │ ├── camera_utils.cpython-38.pyc │ ├── general_utils.cpython-310.pyc │ ├── general_utils.cpython-38.pyc │ ├── graphics_utils.cpython-310.pyc │ ├── graphics_utils.cpython-38.pyc │ ├── image_utils.cpython-38.pyc │ ├── loss_utils.cpython-38.pyc │ ├── sh_utils.cpython-310.pyc │ ├── sh_utils.cpython-38.pyc │ ├── system_utils.cpython-310.pyc │ └── system_utils.cpython-38.pyc │ ├── camera_utils.py │ ├── general_utils.py │ ├── graphics_utils.py │ ├── image_utils.py │ ├── loss_utils.py │ ├── sh_utils.py │ └── system_utils.py ├── lpipsPyTorch ├── __init__.py ├── __pycache__ │ └── __init__.cpython-38.pyc └── modules │ ├── __pycache__ │ ├── lpips.cpython-38.pyc │ ├── networks.cpython-38.pyc │ └── utils.cpython-38.pyc │ ├── lpips.py │ ├── networks.py │ └── utils.py ├── media └── overview.png ├── metrics.py ├── np_utils ├── dataset.py ├── extensions │ └── chamfer_dist │ │ ├── __init__.py │ │ ├── chamfer.cu │ │ ├── chamfer_cuda.cpp │ │ ├── setup.py │ │ └── test.py ├── extract_mesh_meshudf.py ├── owndata.conf ├── train.py └── utils.py ├── requirements.txt ├── sugar_scene ├── __pycache__ │ ├── cameras.cpython-310.pyc │ ├── gs_model.cpython-310.pyc │ ├── sugar_densifier.cpython-310.pyc │ ├── sugar_model.cpython-310.pyc │ └── sugar_optimizer.cpython-310.pyc ├── cameras.py ├── gs_model.py ├── sugar_compositor.py ├── sugar_densifier.py ├── sugar_model.py └── sugar_optimizer.py ├── sugar_trainers ├── __pycache__ │ ├── coarse_density.cpython-310.pyc │ ├── coarse_sdf.cpython-310.pyc │ └── refine.cpython-310.pyc └── coarse_sdf.py ├── sugar_utils ├── __pycache__ │ ├── general_utils.cpython-310.pyc │ ├── graphics_utils.cpython-310.pyc │ ├── loss_utils.cpython-310.pyc │ └── spherical_harmonics.cpython-310.pyc ├── general_utils.py ├── graphics_utils.py ├── loss_utils.py └── spherical_harmonics.py └── train.py /README.md: -------------------------------------------------------------------------------- 1 |

2 |

Neural Signed Distance Function Inference through Splatting 3D Gaussians Pulled on Zero-Level Set

3 | 4 |

5 | Wenyuan Zhang 6 | · 7 | Yu-Shen Liu 8 | · 9 | Zhizhong Han 10 |

11 |

NeurIPS 2024

12 |

Paper | Project Page

13 |
14 |
15 |

16 | 17 |

18 | 19 | In this paper, we propose to seamlessly combine 3D Gaussians with the learning of neural SDFs. Our method provides a novel perspective to jointly learn 3D Gaussians and neural SDFs by more effectively using multi-view consistency and imposing geometry constraints. 20 | 21 | # Preprocessed Datasets & Pretrained Meshes 22 | 23 | Our preprocessed datasets are provided in [This link](https://drive.google.com/drive/folders/1I3mSRrQ6oMV5nlNkaUtXmLVfCX_z-iRS?usp=sharing). 24 | 25 | Pretrained meshes are provided in [This link](https://drive.google.com/drive/folders/13Aao4i-j5bG7Ss-jf6cMBjiit_uFElP5?usp=sharing). 26 | 27 | # Setup 28 | 29 | ## Installation 30 | 31 | Clone the repository and create an anaconda environment called gspull using 32 | ```shell 33 | git clone git@github.com:wen-yuan-zhang/GS-Pull.git 34 | cd GS-Pull 35 | 36 | conda create -n gspull python=3.10 37 | conda activate gspull 38 | 39 | conda install pytorch=1.13.0 torchvision=0.14.0 cudatoolkit=11.7 -c pytorch 40 | conda install cudatoolkit-dev=11.7 -c conda-forge 41 | 42 | pip install -r requirements.txt 43 | ``` 44 | 45 | To install the differentiable splatting kernel, use 46 | ```shell 47 | cd gaussian_splatting/submodules 48 | pip install diff-gaussian-rasterization 49 | pip install simple-knn 50 | ``` 51 | 52 | To install the C++ extensions for NeuralPull, use 53 | ```shell 54 | cd np_utils/extensions/chamfer_dist 55 | python setup.py install 56 | ``` 57 | 58 | (Optional) To try training UDFs, install udf extraction extensions 59 | ```shell 60 | cd custom_mc 61 | python setup.py build_ext --inplace 62 | ``` 63 | 64 | # Training 65 | 66 | To train a scene, firstly run original Gaussian Splatting for 7000 iterations 67 | ```shell 68 | cd gaussian_splatting 69 | python train.py -s -m --iterations 7000 70 | ``` 71 | For example, to train scan24 of DTU dataset, use 72 | ```shell 73 | python train.py -s data/DTU/scan24 -m output/DTU/scan24 --iterations 7000 74 | ``` 75 | The default background color is black. To use white background, you need to add a '-w' argument. 76 | 77 | Then train GS-Pull using 78 | ```shell 79 | cd ../ 80 | python train.py -s -c --output 81 | ``` 82 | For example, to continue training scan24 of DTU dataset, use 83 | ```shell 84 | python train.py -s data/DTU/scan24 -c gaussian_splatting/output/DTU/scan24 --output output/DTU/scan24 85 | ``` 86 | Note that we will identify the scene name in training, so please ensure that the output directory ends with the exact scene name of the dataset. 87 | 88 | # Mesh Extraction 89 | 90 | To extract meshes from checkpoints, use 91 | ```shell 92 | python extract_mesh.py -s -g -o 93 | ``` 94 | For example, to extract mesh of scan24 of DTU dataset, use 95 | ```shell 96 | python extract_mesh.py -s data/DTU/scan24/ -g gaussian_splatting/output/DTU/scan24/ -o output/DTU/scan24 97 | ``` 98 | 99 | # Evaluation 100 | 101 | To evaluate DTU scenes, put the ground truth of DTU dataset under `data/`, and then use 102 | ```shell 103 | cd evaluation 104 | python clean_eval_dtu_mesh.py --datadir --expdir --scan 105 | ``` 106 | For example, to evaluate scan24, use 107 | ```shell 108 | python clean_eval_dtu_mesh.py --datadir ../data/DTU --expdir ../output/DTU/scan24 --scan 24 109 | ``` 110 | 111 | To evaluate Tanks and Temples scenes, follow the official evaluation scipts provided by [TNT dataset](https://github.com/isl-org/TanksAndTemples/tree/master/python_toolbox/evaluation). 112 | 113 | 114 | 115 | # Acknowledgements 116 | 117 | This project is built upon [3DGS](https://github.com/graphdeco-inria/gaussian-splatting), [SuGaR](https://github.com/Anttwo/SuGaR), [Neural-Pull](https://github.com/mabaorui/NeuralPull) and [CAP-UDF](https://github.com/junshengzhou/CAP-UDF). We thank all the authors for their great repos. 118 | 119 | 120 | # Citation 121 | 122 | If you find our code or paper useful, please consider citing 123 | ```bibtex 124 | @inproceedings{zhang2024gspull, 125 | title = {Neural Signed Distance Function Inference through Splatting 3D Gaussians Pulled on Zero-Level Set}, 126 | author = {Wenyuan Zhang and Yu-Shen Liu and Zhizhong Han}, 127 | booktitle = {Advances in Neural Information Processing Systems}, 128 | year = {2024}, 129 | } 130 | ``` 131 | 132 | -------------------------------------------------------------------------------- /custom_mc/_marching_cubes_lewiner.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import numpy as np 3 | import _marching_cubes_lewiner_luts as mcluts 4 | import _marching_cubes_lewiner_cy 5 | 6 | 7 | def marching_cubes_lewiner(volume, level=None, spacing=(1., 1., 1.), 8 | gradient_direction='descent', step_size=1, 9 | allow_degenerate=True, use_classic=False, mask=None): 10 | """Lewiner et al. algorithm for marching cubes. See 11 | marching_cubes_lewiner for documentation. 12 | """ 13 | 14 | # Check volume and ensure its in the format that the alg needs 15 | if not isinstance(volume, np.ndarray) or (volume.ndim != 3): 16 | raise ValueError('Input volume should be a 3D numpy array.') 17 | if volume.shape[0] < 2 or volume.shape[1] < 2 or volume.shape[2] < 2: 18 | raise ValueError("Input array must be at least 2x2x2.") 19 | volume = np.ascontiguousarray(volume, 20 | np.float32) # no copy if not necessary 21 | 22 | # Check/convert other inputs: 23 | # level 24 | if level is None: 25 | level = 0.5 * (volume.min() + volume.max()) 26 | else: 27 | level = float(level) 28 | if level < volume.min() or level > volume.max(): 29 | raise ValueError("Surface level must be within volume data range.") 30 | # spacing 31 | if len(spacing) != 3: 32 | raise ValueError("`spacing` must consist of three floats.") 33 | # step_size 34 | step_size = int(step_size) 35 | if step_size < 1: 36 | raise ValueError('step_size must be at least one.') 37 | # use_classic 38 | use_classic = bool(use_classic) 39 | 40 | # Get LutProvider class (reuse if possible) 41 | L = _get_mc_luts() 42 | 43 | # Check if a mask array is passed 44 | if mask is not None: 45 | if not mask.shape == volume.shape: 46 | raise ValueError('volume and mask must have the same shape.') 47 | 48 | # Apply algorithm 49 | func = _marching_cubes_lewiner_cy.marching_cubes 50 | vertices, faces, normals, values = func(volume, level, L, 51 | step_size, use_classic, mask) 52 | 53 | if not len(vertices): 54 | raise RuntimeError('No surface found at the given iso value.') 55 | 56 | # Output in z-y-x order, as is common in skimage 57 | vertices = np.fliplr(vertices) 58 | normals = np.fliplr(normals) 59 | 60 | # Finishing touches to output 61 | faces.shape = -1, 3 62 | if gradient_direction == 'descent': 63 | # MC implementation is right-handed, but gradient_direction is 64 | # left-handed 65 | faces = np.fliplr(faces) 66 | elif not gradient_direction == 'ascent': 67 | raise ValueError("Incorrect input %s in `gradient_direction`, see " 68 | "docstring." % (gradient_direction)) 69 | if not np.array_equal(spacing, (1, 1, 1)): 70 | vertices = vertices * np.r_[spacing] 71 | 72 | if allow_degenerate: 73 | return vertices, faces, normals, values 74 | else: 75 | fun = _marching_cubes_lewiner_cy.remove_degenerate_faces 76 | return fun(vertices.astype(np.float32), faces, normals, values) 77 | 78 | 79 | 80 | def udf_mc_lewiner(volume, grads, spacing=(1., 1., 1.), 81 | gradient_direction='descent', step_size=1, 82 | allow_degenerate=True, use_classic=False, mask=None): 83 | """Lewiner et al. algorithm for marching cubes. See 84 | marching_cubes_lewiner for documentation. 85 | """ 86 | 87 | # Check volume and ensure its in the format that the alg needs 88 | if not isinstance(volume, np.ndarray) or (volume.ndim != 3): 89 | raise ValueError('Input volume should be a 3D numpy array.') 90 | if volume.shape[0] < 2 or volume.shape[1] < 2 or volume.shape[2] < 2: 91 | raise ValueError("Input array must be at least 2x2x2.") 92 | volume = np.ascontiguousarray(volume, 93 | np.float32) # no copy if not necessary 94 | 95 | # spacing 96 | if len(spacing) != 3: 97 | raise ValueError("`spacing` must consist of three floats.") 98 | # step_size 99 | step_size = int(step_size) 100 | if step_size < 1: 101 | raise ValueError('step_size must be at least one.') 102 | # use_classic 103 | use_classic = bool(use_classic) 104 | 105 | # Get LutProvider class (reuse if possible) 106 | L = _get_mc_luts() 107 | 108 | # Check if a mask array is passed 109 | if mask is not None: 110 | if not mask.shape == volume.shape: 111 | raise ValueError('volume and mask must have the same shape.') 112 | 113 | # Apply algorithm 114 | func = _marching_cubes_lewiner_cy.marching_cubes_udf 115 | vertices, faces, normals, values = func(volume, grads, L, 116 | step_size, use_classic, mask) 117 | 118 | if not len(vertices): 119 | raise RuntimeError('No surface found at the given iso value.') 120 | 121 | # Output in z-y-x order, as is common in skimage 122 | vertices = np.fliplr(vertices) 123 | normals = np.fliplr(normals) 124 | 125 | # Finishing touches to output 126 | faces.shape = -1, 3 127 | if gradient_direction == 'descent': 128 | # MC implementation is right-handed, but gradient_direction is 129 | # left-handed 130 | faces = np.fliplr(faces) 131 | elif not gradient_direction == 'ascent': 132 | raise ValueError("Incorrect input %s in `gradient_direction`, see " 133 | "docstring." % (gradient_direction)) 134 | if not np.array_equal(spacing, (1, 1, 1)): 135 | vertices = vertices * np.r_[spacing] 136 | 137 | if allow_degenerate: 138 | return vertices, faces, normals, values 139 | else: 140 | fun = _marching_cubes_lewiner_cy.remove_degenerate_faces 141 | return fun(vertices.astype(np.float32), faces, normals, values) 142 | 143 | 144 | 145 | def _to_array(args): 146 | shape, text = args 147 | byts = base64.decodebytes(text.encode('utf-8')) 148 | ar = np.frombuffer(byts, dtype='int8') 149 | ar.shape = shape 150 | return ar 151 | 152 | 153 | # Map an edge-index to two relative pixel positions. The ege index 154 | # represents a point that lies somewhere in between these pixels. 155 | # Linear interpolation should be used to determine where it is exactly. 156 | # 0 157 | # 3 1 -> 0x 158 | # 2 xx 159 | EDGETORELATIVEPOSX = np.array([ [0,1],[1,1],[1,0],[0,0], [0,1],[1,1],[1,0],[0,0], [0,0],[1,1],[1,1],[0,0] ], 'int8') 160 | EDGETORELATIVEPOSY = np.array([ [0,0],[0,1],[1,1],[1,0], [0,0],[0,1],[1,1],[1,0], [0,0],[0,0],[1,1],[1,1] ], 'int8') 161 | EDGETORELATIVEPOSZ = np.array([ [0,0],[0,0],[0,0],[0,0], [1,1],[1,1],[1,1],[1,1], [0,1],[0,1],[0,1],[0,1] ], 'int8') 162 | 163 | 164 | def _get_mc_luts(): 165 | """ Kind of lazy obtaining of the luts. 166 | """ 167 | if not hasattr(mcluts, 'THE_LUTS'): 168 | 169 | mcluts.THE_LUTS = _marching_cubes_lewiner_cy.LutProvider( 170 | EDGETORELATIVEPOSX, EDGETORELATIVEPOSY, EDGETORELATIVEPOSZ, 171 | 172 | _to_array(mcluts.CASESCLASSIC), _to_array(mcluts.CASES), 173 | 174 | _to_array(mcluts.TILING1), _to_array(mcluts.TILING2), _to_array(mcluts.TILING3_1), _to_array(mcluts.TILING3_2), 175 | _to_array(mcluts.TILING4_1), _to_array(mcluts.TILING4_2), _to_array(mcluts.TILING5), _to_array(mcluts.TILING6_1_1), 176 | _to_array(mcluts.TILING6_1_2), _to_array(mcluts.TILING6_2), _to_array(mcluts.TILING7_1), 177 | _to_array(mcluts.TILING7_2), _to_array(mcluts.TILING7_3), _to_array(mcluts.TILING7_4_1), 178 | _to_array(mcluts.TILING7_4_2), _to_array(mcluts.TILING8), _to_array(mcluts.TILING9), 179 | _to_array(mcluts.TILING10_1_1), _to_array(mcluts.TILING10_1_1_), _to_array(mcluts.TILING10_1_2), 180 | _to_array(mcluts.TILING10_2), _to_array(mcluts.TILING10_2_), _to_array(mcluts.TILING11), 181 | _to_array(mcluts.TILING12_1_1), _to_array(mcluts.TILING12_1_1_), _to_array(mcluts.TILING12_1_2), 182 | _to_array(mcluts.TILING12_2), _to_array(mcluts.TILING12_2_), _to_array(mcluts.TILING13_1), 183 | _to_array(mcluts.TILING13_1_), _to_array(mcluts.TILING13_2), _to_array(mcluts.TILING13_2_), 184 | _to_array(mcluts.TILING13_3), _to_array(mcluts.TILING13_3_), _to_array(mcluts.TILING13_4), 185 | _to_array(mcluts.TILING13_5_1), _to_array(mcluts.TILING13_5_2), _to_array(mcluts.TILING14), 186 | 187 | _to_array(mcluts.TEST3), _to_array(mcluts.TEST4), _to_array(mcluts.TEST6), 188 | _to_array(mcluts.TEST7), _to_array(mcluts.TEST10), _to_array(mcluts.TEST12), 189 | _to_array(mcluts.TEST13), _to_array(mcluts.SUBCONFIG13), 190 | ) 191 | 192 | return mcluts.THE_LUTS 193 | -------------------------------------------------------------------------------- /custom_mc/_marching_cubes_lewiner_cy.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/custom_mc/_marching_cubes_lewiner_cy.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /custom_mc/setup.py: -------------------------------------------------------------------------------- 1 | # python setup.py build_ext --inplace 2 | from setuptools import setup 3 | from Cython.Build import cythonize 4 | import numpy as np 5 | import os 6 | 7 | includes_numpy = '-I ' + np.get_include() + ' ' 8 | os.environ['CFLAGS'] = includes_numpy + (os.environ['CFLAGS'] if 'CFLAGS' in os.environ else '') 9 | 10 | setup( 11 | name="My MC", 12 | ext_modules=cythonize("_marching_cubes_lewiner_cy.pyx", include_path=[np.get_include()], language="c++"), 13 | ) 14 | -------------------------------------------------------------------------------- /gaussian_splatting/LICENSE.md: -------------------------------------------------------------------------------- 1 | Gaussian-Splatting License 2 | =========================== 3 | 4 | **Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**. 5 | The *Software* is in the process of being registered with the Agence pour la Protection des 6 | Programmes (APP). 7 | 8 | The *Software* is still being developed by the *Licensor*. 9 | 10 | *Licensor*'s goal is to allow the research community to use, test and evaluate 11 | the *Software*. 12 | 13 | ## 1. Definitions 14 | 15 | *Licensee* means any person or entity that uses the *Software* and distributes 16 | its *Work*. 17 | 18 | *Licensor* means the owners of the *Software*, i.e Inria and MPII 19 | 20 | *Software* means the original work of authorship made available under this 21 | License ie gaussian-splatting. 22 | 23 | *Work* means the *Software* and any additions to or derivative works of the 24 | *Software* that are made available under this License. 25 | 26 | 27 | ## 2. Purpose 28 | This license is intended to define the rights granted to the *Licensee* by 29 | Licensors under the *Software*. 30 | 31 | ## 3. Rights granted 32 | 33 | For the above reasons Licensors have decided to distribute the *Software*. 34 | Licensors grant non-exclusive rights to use the *Software* for research purposes 35 | to research users (both academic and industrial), free of charge, without right 36 | to sublicense.. The *Software* may be used "non-commercially", i.e., for research 37 | and/or evaluation purposes only. 38 | 39 | Subject to the terms and conditions of this License, you are granted a 40 | non-exclusive, royalty-free, license to reproduce, prepare derivative works of, 41 | publicly display, publicly perform and distribute its *Work* and any resulting 42 | derivative works in any form. 43 | 44 | ## 4. Limitations 45 | 46 | **4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do 47 | so under this License, (b) you include a complete copy of this License with 48 | your distribution, and (c) you retain without modification any copyright, 49 | patent, trademark, or attribution notices that are present in the *Work*. 50 | 51 | **4.2 Derivative Works.** You may specify that additional or different terms apply 52 | to the use, reproduction, and distribution of your derivative works of the *Work* 53 | ("Your Terms") only if (a) Your Terms provide that the use limitation in 54 | Section 2 applies to your derivative works, and (b) you identify the specific 55 | derivative works that are subject to Your Terms. Notwithstanding Your Terms, 56 | this License (including the redistribution requirements in Section 3.1) will 57 | continue to apply to the *Work* itself. 58 | 59 | **4.3** Any other use without of prior consent of Licensors is prohibited. Research 60 | users explicitly acknowledge having received from Licensors all information 61 | allowing to appreciate the adequacy between of the *Software* and their needs and 62 | to undertake all necessary precautions for its execution and use. 63 | 64 | **4.4** The *Software* is provided both as a compiled library file and as source 65 | code. In case of using the *Software* for a publication or other results obtained 66 | through the use of the *Software*, users are strongly encouraged to cite the 67 | corresponding publications as explained in the documentation of the *Software*. 68 | 69 | ## 5. Disclaimer 70 | 71 | THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES 72 | WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY 73 | UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL 74 | CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES 75 | OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL 76 | USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR 77 | ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE 78 | AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 79 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE 80 | GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) 81 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 82 | LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR 83 | IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*. 84 | -------------------------------------------------------------------------------- /gaussian_splatting/arguments/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from argparse import ArgumentParser, Namespace 13 | import sys 14 | import os 15 | 16 | class GroupParams: 17 | pass 18 | 19 | class ParamGroup: 20 | def __init__(self, parser: ArgumentParser, name : str, fill_none = False): 21 | group = parser.add_argument_group(name) 22 | for key, value in vars(self).items(): 23 | shorthand = False 24 | if key.startswith("_"): 25 | shorthand = True 26 | key = key[1:] 27 | t = type(value) 28 | value = value if not fill_none else None 29 | if shorthand: 30 | if t == bool: 31 | group.add_argument("--" + key, ("-" + key[0:1]), default=value, action="store_true") 32 | else: 33 | group.add_argument("--" + key, ("-" + key[0:1]), default=value, type=t) 34 | else: 35 | if t == bool: 36 | group.add_argument("--" + key, default=value, action="store_true") 37 | else: 38 | group.add_argument("--" + key, default=value, type=t) 39 | 40 | def extract(self, args): 41 | group = GroupParams() 42 | for arg in vars(args).items(): 43 | if arg[0] in vars(self) or ("_" + arg[0]) in vars(self): 44 | setattr(group, arg[0], arg[1]) 45 | return group 46 | 47 | class ModelParams(ParamGroup): 48 | def __init__(self, parser, sentinel=False): 49 | self.sh_degree = 3 50 | self._source_path = "" 51 | self._model_path = "" 52 | self._images = "images" 53 | self._resolution = -1 54 | self._white_background = False 55 | self.data_device = "cuda" 56 | self.eval = False 57 | super().__init__(parser, "Loading Parameters", sentinel) 58 | 59 | def extract(self, args): 60 | g = super().extract(args) 61 | g.source_path = os.path.abspath(g.source_path) 62 | return g 63 | 64 | class PipelineParams(ParamGroup): 65 | def __init__(self, parser): 66 | self.convert_SHs_python = False 67 | self.compute_cov3D_python = False 68 | self.debug = False 69 | super().__init__(parser, "Pipeline Parameters") 70 | 71 | class OptimizationParams(ParamGroup): 72 | def __init__(self, parser): 73 | self.iterations = 30_000 74 | self.position_lr_init = 0.00016 75 | self.position_lr_final = 0.0000016 76 | self.position_lr_delay_mult = 0.01 77 | self.position_lr_max_steps = 30_000 78 | self.feature_lr = 0.0025 79 | self.opacity_lr = 0.05 80 | self.scaling_lr = 0.005 81 | self.rotation_lr = 0.001 82 | self.percent_dense = 0.01 83 | self.lambda_dssim = 0.2 84 | self.densification_interval = 500 85 | self.opacity_reset_interval = 3000 86 | self.densify_from_iter = 3000 87 | self.densify_until_iter = 15_000 88 | self.densify_grad_threshold = 0.0002 89 | self.random_background = False 90 | super().__init__(parser, "Optimization Parameters") 91 | 92 | def get_combined_args(parser : ArgumentParser): 93 | cmdlne_string = sys.argv[1:] 94 | cfgfile_string = "Namespace()" 95 | args_cmdline = parser.parse_args(cmdlne_string) 96 | 97 | try: 98 | cfgfilepath = os.path.join(args_cmdline.model_path, "cfg_args") 99 | print("Looking for config file in", cfgfilepath) 100 | with open(cfgfilepath) as cfg_file: 101 | print("Config file found: {}".format(cfgfilepath)) 102 | cfgfile_string = cfg_file.read() 103 | except TypeError: 104 | print("Config file not found at") 105 | pass 106 | args_cfgfile = eval(cfgfile_string) 107 | 108 | merged_dict = vars(args_cfgfile).copy() 109 | for k,v in vars(args_cmdline).items(): 110 | if v != None: 111 | merged_dict[k] = v 112 | return Namespace(**merged_dict) 113 | -------------------------------------------------------------------------------- /gaussian_splatting/arguments/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/arguments/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/arguments/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/arguments/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/convert.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import os 13 | import logging 14 | from argparse import ArgumentParser 15 | import shutil 16 | 17 | # This Python script is based on the shell converter script provided in the MipNerF 360 repository. 18 | parser = ArgumentParser("Colmap converter") 19 | parser.add_argument("--no_gpu", action='store_true') 20 | parser.add_argument("--skip_matching", action='store_true') 21 | parser.add_argument("--source_path", "-s", required=True, type=str) 22 | parser.add_argument("--camera", default="OPENCV", type=str) 23 | parser.add_argument("--colmap_executable", default="", type=str) 24 | parser.add_argument("--resize", action="store_true") 25 | parser.add_argument("--magick_executable", default="", type=str) 26 | args = parser.parse_args() 27 | colmap_command = '"{}"'.format(args.colmap_executable) if len(args.colmap_executable) > 0 else "colmap" 28 | magick_command = '"{}"'.format(args.magick_executable) if len(args.magick_executable) > 0 else "magick" 29 | use_gpu = 1 if not args.no_gpu else 0 30 | 31 | if not args.skip_matching: 32 | os.makedirs(args.source_path + "/distorted/sparse", exist_ok=True) 33 | 34 | ## Feature extraction 35 | feat_extracton_cmd = colmap_command + " feature_extractor "\ 36 | "--database_path " + args.source_path + "/distorted/database.db \ 37 | --image_path " + args.source_path + "/input \ 38 | --ImageReader.single_camera 1 \ 39 | --ImageReader.camera_model " + args.camera + " \ 40 | --SiftExtraction.use_gpu " + str(use_gpu) 41 | exit_code = os.system(feat_extracton_cmd) 42 | if exit_code != 0: 43 | logging.error(f"Feature extraction failed with code {exit_code}. Exiting.") 44 | exit(exit_code) 45 | 46 | ## Feature matching 47 | feat_matching_cmd = colmap_command + " exhaustive_matcher \ 48 | --database_path " + args.source_path + "/distorted/database.db \ 49 | --SiftMatching.use_gpu " + str(use_gpu) 50 | exit_code = os.system(feat_matching_cmd) 51 | if exit_code != 0: 52 | logging.error(f"Feature matching failed with code {exit_code}. Exiting.") 53 | exit(exit_code) 54 | 55 | ### Bundle adjustment 56 | # The default Mapper tolerance is unnecessarily large, 57 | # decreasing it speeds up bundle adjustment steps. 58 | mapper_cmd = (colmap_command + " mapper \ 59 | --database_path " + args.source_path + "/distorted/database.db \ 60 | --image_path " + args.source_path + "/input \ 61 | --output_path " + args.source_path + "/distorted/sparse \ 62 | --Mapper.ba_global_function_tolerance=0.000001") 63 | exit_code = os.system(mapper_cmd) 64 | if exit_code != 0: 65 | logging.error(f"Mapper failed with code {exit_code}. Exiting.") 66 | exit(exit_code) 67 | 68 | ### Image undistortion 69 | ## We need to undistort our images into ideal pinhole intrinsics. 70 | img_undist_cmd = (colmap_command + " image_undistorter \ 71 | --image_path " + args.source_path + "/input \ 72 | --input_path " + args.source_path + "/distorted/sparse/0 \ 73 | --output_path " + args.source_path + "\ 74 | --output_type COLMAP") 75 | exit_code = os.system(img_undist_cmd) 76 | if exit_code != 0: 77 | logging.error(f"Mapper failed with code {exit_code}. Exiting.") 78 | exit(exit_code) 79 | 80 | files = os.listdir(args.source_path + "/sparse") 81 | os.makedirs(args.source_path + "/sparse/0", exist_ok=True) 82 | # Copy each file from the source directory to the destination directory 83 | for file in files: 84 | if file == '0': 85 | continue 86 | source_file = os.path.join(args.source_path, "sparse", file) 87 | destination_file = os.path.join(args.source_path, "sparse", "0", file) 88 | shutil.move(source_file, destination_file) 89 | 90 | if(args.resize): 91 | print("Copying and resizing...") 92 | 93 | # Resize images. 94 | os.makedirs(args.source_path + "/images_2", exist_ok=True) 95 | os.makedirs(args.source_path + "/images_4", exist_ok=True) 96 | os.makedirs(args.source_path + "/images_8", exist_ok=True) 97 | # Get the list of files in the source directory 98 | files = os.listdir(args.source_path + "/images") 99 | # Copy each file from the source directory to the destination directory 100 | for file in files: 101 | source_file = os.path.join(args.source_path, "images", file) 102 | 103 | destination_file = os.path.join(args.source_path, "images_2", file) 104 | shutil.copy2(source_file, destination_file) 105 | exit_code = os.system(magick_command + " mogrify -resize 50% " + destination_file) 106 | if exit_code != 0: 107 | logging.error(f"50% resize failed with code {exit_code}. Exiting.") 108 | exit(exit_code) 109 | 110 | destination_file = os.path.join(args.source_path, "images_4", file) 111 | shutil.copy2(source_file, destination_file) 112 | exit_code = os.system(magick_command + " mogrify -resize 25% " + destination_file) 113 | if exit_code != 0: 114 | logging.error(f"25% resize failed with code {exit_code}. Exiting.") 115 | exit(exit_code) 116 | 117 | destination_file = os.path.join(args.source_path, "images_8", file) 118 | shutil.copy2(source_file, destination_file) 119 | exit_code = os.system(magick_command + " mogrify -resize 12.5% " + destination_file) 120 | if exit_code != 0: 121 | logging.error(f"12.5% resize failed with code {exit_code}. Exiting.") 122 | exit(exit_code) 123 | 124 | print("Done.") 125 | -------------------------------------------------------------------------------- /gaussian_splatting/environment.yml: -------------------------------------------------------------------------------- 1 | name: gaussian_splatting 2 | channels: 3 | - pytorch 4 | - conda-forge 5 | - defaults 6 | dependencies: 7 | - cudatoolkit=11.6 8 | - plyfile=0.8.1 9 | - python=3.7.13 10 | - pip=22.3.1 11 | - pytorch=1.12.1 12 | - torchaudio=0.12.1 13 | - torchvision=0.13.1 14 | - tqdm 15 | - pip: 16 | - submodules/diff-gaussian-rasterization 17 | - submodules/simple-knn -------------------------------------------------------------------------------- /gaussian_splatting/full_eval.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import os 13 | from argparse import ArgumentParser 14 | 15 | mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"] 16 | mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"] 17 | tanks_and_temples_scenes = ["truck", "train"] 18 | deep_blending_scenes = ["drjohnson", "playroom"] 19 | 20 | parser = ArgumentParser(description="Full evaluation script parameters") 21 | parser.add_argument("--skip_training", action="store_true") 22 | parser.add_argument("--skip_rendering", action="store_true") 23 | parser.add_argument("--skip_metrics", action="store_true") 24 | parser.add_argument("--output_path", default="./eval") 25 | args, _ = parser.parse_known_args() 26 | 27 | all_scenes = [] 28 | all_scenes.extend(mipnerf360_outdoor_scenes) 29 | all_scenes.extend(mipnerf360_indoor_scenes) 30 | all_scenes.extend(tanks_and_temples_scenes) 31 | all_scenes.extend(deep_blending_scenes) 32 | 33 | if not args.skip_training or not args.skip_rendering: 34 | parser.add_argument('--mipnerf360', "-m360", required=True, type=str) 35 | parser.add_argument("--tanksandtemples", "-tat", required=True, type=str) 36 | parser.add_argument("--deepblending", "-db", required=True, type=str) 37 | args = parser.parse_args() 38 | 39 | if not args.skip_training: 40 | common_args = " --quiet --eval --test_iterations -1 " 41 | for scene in mipnerf360_outdoor_scenes: 42 | source = args.mipnerf360 + "/" + scene 43 | os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args) 44 | for scene in mipnerf360_indoor_scenes: 45 | source = args.mipnerf360 + "/" + scene 46 | os.system("python train.py -s " + source + " -i images_2 -m " + args.output_path + "/" + scene + common_args) 47 | for scene in tanks_and_temples_scenes: 48 | source = args.tanksandtemples + "/" + scene 49 | os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 50 | for scene in deep_blending_scenes: 51 | source = args.deepblending + "/" + scene 52 | os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 53 | 54 | if not args.skip_rendering: 55 | all_sources = [] 56 | for scene in mipnerf360_outdoor_scenes: 57 | all_sources.append(args.mipnerf360 + "/" + scene) 58 | for scene in mipnerf360_indoor_scenes: 59 | all_sources.append(args.mipnerf360 + "/" + scene) 60 | for scene in tanks_and_temples_scenes: 61 | all_sources.append(args.tanksandtemples + "/" + scene) 62 | for scene in deep_blending_scenes: 63 | all_sources.append(args.deepblending + "/" + scene) 64 | 65 | common_args = " --quiet --eval --skip_train" 66 | for scene, source in zip(all_scenes, all_sources): 67 | os.system("python render.py --iteration 7000 -s " + source + " -m " + args.output_path + "/" + scene + common_args) 68 | os.system("python render.py --iteration 30000 -s " + source + " -m " + args.output_path + "/" + scene + common_args) 69 | 70 | if not args.skip_metrics: 71 | scenes_string = "" 72 | for scene in all_scenes: 73 | scenes_string += "\"" + args.output_path + "/" + scene + "\" " 74 | 75 | os.system("python metrics.py -m " + scenes_string) -------------------------------------------------------------------------------- /gaussian_splatting/gaussian_renderer/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import math 14 | from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer 15 | from scene.gaussian_model import GaussianModel 16 | from utils.sh_utils import eval_sh 17 | 18 | def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None): 19 | """ 20 | Render the scene. 21 | 22 | Background tensor (bg_color) must be on GPU! 23 | """ 24 | 25 | # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means 26 | screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 27 | try: 28 | screenspace_points.retain_grad() 29 | except: 30 | pass 31 | 32 | # Set up rasterization configuration 33 | tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) 34 | tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) 35 | 36 | raster_settings = GaussianRasterizationSettings( 37 | image_height=int(viewpoint_camera.image_height), 38 | image_width=int(viewpoint_camera.image_width), 39 | tanfovx=tanfovx, 40 | tanfovy=tanfovy, 41 | bg=bg_color, 42 | scale_modifier=scaling_modifier, 43 | viewmatrix=viewpoint_camera.world_view_transform, 44 | projmatrix=viewpoint_camera.full_proj_transform, 45 | sh_degree=pc.active_sh_degree, 46 | campos=viewpoint_camera.camera_center, 47 | prefiltered=False, 48 | debug=pipe.debug 49 | ) 50 | 51 | rasterizer = GaussianRasterizer(raster_settings=raster_settings) 52 | 53 | means3D = pc.get_xyz 54 | means2D = screenspace_points 55 | opacity = pc.get_opacity 56 | 57 | # If precomputed 3d covariance is provided, use it. If not, then it will be computed from 58 | # scaling / rotation by the rasterizer. 59 | scales = None 60 | rotations = None 61 | cov3D_precomp = None 62 | if pipe.compute_cov3D_python: 63 | cov3D_precomp = pc.get_covariance(scaling_modifier) 64 | else: 65 | scales = pc.get_scaling 66 | rotations = pc.get_rotation 67 | 68 | # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors 69 | # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. 70 | shs = None 71 | colors_precomp = None 72 | if override_color is None: 73 | if pipe.convert_SHs_python: 74 | shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) 75 | dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) 76 | dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) 77 | sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) 78 | colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) 79 | else: 80 | shs = pc.get_features 81 | else: 82 | colors_precomp = override_color 83 | 84 | # Rasterize visible Gaussians to image, obtain their radii (on screen). 85 | rendered_image, radii = rasterizer( 86 | means3D = means3D, 87 | means2D = means2D, 88 | shs = shs, 89 | colors_precomp = colors_precomp, 90 | opacities = opacity, 91 | scales = scales, 92 | rotations = rotations, 93 | cov3D_precomp = cov3D_precomp) 94 | 95 | # Those Gaussians that were frustum culled or had a radius of 0 were not visible. 96 | # They will be excluded from value updates used in the splitting criteria. 97 | return {"render": rendered_image, 98 | "viewspace_points": screenspace_points, 99 | "visibility_filter" : radii > 0, 100 | "radii": radii} 101 | -------------------------------------------------------------------------------- /gaussian_splatting/gaussian_renderer/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/gaussian_renderer/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/gaussian_renderer/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/gaussian_renderer/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/gaussian_renderer/__pycache__/network_gui.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/gaussian_renderer/__pycache__/network_gui.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/gaussian_renderer/network_gui.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import traceback 14 | import socket 15 | import json 16 | from scene.cameras import MiniCam 17 | 18 | host = "127.0.0.1" 19 | port = 6009 20 | 21 | conn = None 22 | addr = None 23 | 24 | listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 25 | 26 | def init(wish_host, wish_port): 27 | global host, port, listener 28 | host = wish_host 29 | port = wish_port 30 | listener.bind((host, port)) 31 | listener.listen() 32 | listener.settimeout(0) 33 | 34 | def try_connect(): 35 | global conn, addr, listener 36 | try: 37 | conn, addr = listener.accept() 38 | print(f"\nConnected by {addr}") 39 | conn.settimeout(None) 40 | except Exception as inst: 41 | pass 42 | 43 | def read(): 44 | global conn 45 | messageLength = conn.recv(4) 46 | messageLength = int.from_bytes(messageLength, 'little') 47 | message = conn.recv(messageLength) 48 | return json.loads(message.decode("utf-8")) 49 | 50 | def send(message_bytes, verify): 51 | global conn 52 | if message_bytes != None: 53 | conn.sendall(message_bytes) 54 | conn.sendall(len(verify).to_bytes(4, 'little')) 55 | conn.sendall(bytes(verify, 'ascii')) 56 | 57 | def receive(): 58 | message = read() 59 | 60 | width = message["resolution_x"] 61 | height = message["resolution_y"] 62 | 63 | if width != 0 and height != 0: 64 | try: 65 | do_training = bool(message["train"]) 66 | fovy = message["fov_y"] 67 | fovx = message["fov_x"] 68 | znear = message["z_near"] 69 | zfar = message["z_far"] 70 | do_shs_python = bool(message["shs_python"]) 71 | do_rot_scale_python = bool(message["rot_scale_python"]) 72 | keep_alive = bool(message["keep_alive"]) 73 | scaling_modifier = message["scaling_modifier"] 74 | world_view_transform = torch.reshape(torch.tensor(message["view_matrix"]), (4, 4)).cuda() 75 | world_view_transform[:,1] = -world_view_transform[:,1] 76 | world_view_transform[:,2] = -world_view_transform[:,2] 77 | full_proj_transform = torch.reshape(torch.tensor(message["view_projection_matrix"]), (4, 4)).cuda() 78 | full_proj_transform[:,1] = -full_proj_transform[:,1] 79 | custom_cam = MiniCam(width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform) 80 | except Exception as e: 81 | print("") 82 | traceback.print_exc() 83 | raise e 84 | return custom_cam, do_training, do_shs_python, do_rot_scale_python, keep_alive, scaling_modifier 85 | else: 86 | return None, None, None, None, None, None -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .modules.lpips import LPIPS 4 | 5 | 6 | def lpips(x: torch.Tensor, 7 | y: torch.Tensor, 8 | net_type: str = 'alex', 9 | version: str = '0.1'): 10 | r"""Function that measures 11 | Learned Perceptual Image Patch Similarity (LPIPS). 12 | 13 | Arguments: 14 | x, y (torch.Tensor): the input tensors to compare. 15 | net_type (str): the network type to compare the features: 16 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 17 | version (str): the version of LPIPS. Default: 0.1. 18 | """ 19 | device = x.device 20 | criterion = LPIPS(net_type, version).to(device) 21 | return criterion(x, y) 22 | -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/lpipsPyTorch/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/__pycache__/lpips.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/lpipsPyTorch/modules/__pycache__/lpips.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/__pycache__/networks.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/lpipsPyTorch/modules/__pycache__/networks.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/lpipsPyTorch/modules/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/lpips.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .networks import get_network, LinLayers 5 | from .utils import get_state_dict 6 | 7 | 8 | class LPIPS(nn.Module): 9 | r"""Creates a criterion that measures 10 | Learned Perceptual Image Patch Similarity (LPIPS). 11 | 12 | Arguments: 13 | net_type (str): the network type to compare the features: 14 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 15 | version (str): the version of LPIPS. Default: 0.1. 16 | """ 17 | def __init__(self, net_type: str = 'alex', version: str = '0.1'): 18 | 19 | assert version in ['0.1'], 'v0.1 is only supported now' 20 | 21 | super(LPIPS, self).__init__() 22 | 23 | # pretrained network 24 | self.net = get_network(net_type) 25 | 26 | # linear layers 27 | self.lin = LinLayers(self.net.n_channels_list) 28 | self.lin.load_state_dict(get_state_dict(net_type, version)) 29 | 30 | def forward(self, x: torch.Tensor, y: torch.Tensor): 31 | feat_x, feat_y = self.net(x), self.net(y) 32 | 33 | diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] 34 | res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] 35 | 36 | return torch.sum(torch.cat(res, 0), 0, True) 37 | -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/networks.py: -------------------------------------------------------------------------------- 1 | from typing import Sequence 2 | 3 | from itertools import chain 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torchvision import models 8 | 9 | from .utils import normalize_activation 10 | 11 | 12 | def get_network(net_type: str): 13 | if net_type == 'alex': 14 | return AlexNet() 15 | elif net_type == 'squeeze': 16 | return SqueezeNet() 17 | elif net_type == 'vgg': 18 | return VGG16() 19 | else: 20 | raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') 21 | 22 | 23 | class LinLayers(nn.ModuleList): 24 | def __init__(self, n_channels_list: Sequence[int]): 25 | super(LinLayers, self).__init__([ 26 | nn.Sequential( 27 | nn.Identity(), 28 | nn.Conv2d(nc, 1, 1, 1, 0, bias=False) 29 | ) for nc in n_channels_list 30 | ]) 31 | 32 | for param in self.parameters(): 33 | param.requires_grad = False 34 | 35 | 36 | class BaseNet(nn.Module): 37 | def __init__(self): 38 | super(BaseNet, self).__init__() 39 | 40 | # register buffer 41 | self.register_buffer( 42 | 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) 43 | self.register_buffer( 44 | 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) 45 | 46 | def set_requires_grad(self, state: bool): 47 | for param in chain(self.parameters(), self.buffers()): 48 | param.requires_grad = state 49 | 50 | def z_score(self, x: torch.Tensor): 51 | return (x - self.mean) / self.std 52 | 53 | def forward(self, x: torch.Tensor): 54 | x = self.z_score(x) 55 | 56 | output = [] 57 | for i, (_, layer) in enumerate(self.layers._modules.items(), 1): 58 | x = layer(x) 59 | if i in self.target_layers: 60 | output.append(normalize_activation(x)) 61 | if len(output) == len(self.target_layers): 62 | break 63 | return output 64 | 65 | 66 | class SqueezeNet(BaseNet): 67 | def __init__(self): 68 | super(SqueezeNet, self).__init__() 69 | 70 | self.layers = models.squeezenet1_1(True).features 71 | self.target_layers = [2, 5, 8, 10, 11, 12, 13] 72 | self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] 73 | 74 | self.set_requires_grad(False) 75 | 76 | 77 | class AlexNet(BaseNet): 78 | def __init__(self): 79 | super(AlexNet, self).__init__() 80 | 81 | self.layers = models.alexnet(True).features 82 | self.target_layers = [2, 5, 8, 10, 12] 83 | self.n_channels_list = [64, 192, 384, 256, 256] 84 | 85 | self.set_requires_grad(False) 86 | 87 | 88 | class VGG16(BaseNet): 89 | def __init__(self): 90 | super(VGG16, self).__init__() 91 | 92 | self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features 93 | self.target_layers = [4, 9, 16, 23, 30] 94 | self.n_channels_list = [64, 128, 256, 512, 512] 95 | 96 | self.set_requires_grad(False) 97 | -------------------------------------------------------------------------------- /gaussian_splatting/lpipsPyTorch/modules/utils.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | import torch 4 | 5 | 6 | def normalize_activation(x, eps=1e-10): 7 | norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) 8 | return x / (norm_factor + eps) 9 | 10 | 11 | def get_state_dict(net_type: str = 'alex', version: str = '0.1'): 12 | # build url 13 | url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ 14 | + f'master/lpips/weights/v{version}/{net_type}.pth' 15 | 16 | # download 17 | old_state_dict = torch.hub.load_state_dict_from_url( 18 | url, progress=True, 19 | map_location=None if torch.cuda.is_available() else torch.device('cpu') 20 | ) 21 | 22 | # rename keys 23 | new_state_dict = OrderedDict() 24 | for key, val in old_state_dict.items(): 25 | new_key = key 26 | new_key = new_key.replace('lin', '') 27 | new_key = new_key.replace('model.', '') 28 | new_state_dict[new_key] = val 29 | 30 | return new_state_dict 31 | -------------------------------------------------------------------------------- /gaussian_splatting/metrics.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from pathlib import Path 13 | import os 14 | from PIL import Image 15 | import torch 16 | import torchvision.transforms.functional as tf 17 | from utils.loss_utils import ssim 18 | from lpipsPyTorch import lpips 19 | import json 20 | from tqdm import tqdm 21 | from utils.image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | def readImages(renders_dir, gt_dir): 25 | renders = [] 26 | gts = [] 27 | image_names = [] 28 | for fname in os.listdir(renders_dir): 29 | render = Image.open(renders_dir / fname) 30 | gt = Image.open(gt_dir / fname) 31 | renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda()) 32 | gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda()) 33 | image_names.append(fname) 34 | return renders, gts, image_names 35 | 36 | def evaluate(model_paths): 37 | 38 | full_dict = {} 39 | per_view_dict = {} 40 | full_dict_polytopeonly = {} 41 | per_view_dict_polytopeonly = {} 42 | print("") 43 | 44 | for scene_dir in model_paths: 45 | # try: 46 | print("Scene:", scene_dir) 47 | full_dict[scene_dir] = {} 48 | per_view_dict[scene_dir] = {} 49 | full_dict_polytopeonly[scene_dir] = {} 50 | per_view_dict_polytopeonly[scene_dir] = {} 51 | 52 | test_dir = Path(scene_dir) / "test" 53 | 54 | for method in os.listdir(test_dir): 55 | print("Method:", method) 56 | 57 | full_dict[scene_dir][method] = {} 58 | per_view_dict[scene_dir][method] = {} 59 | full_dict_polytopeonly[scene_dir][method] = {} 60 | per_view_dict_polytopeonly[scene_dir][method] = {} 61 | 62 | method_dir = test_dir / method 63 | gt_dir = method_dir/ "gt" 64 | renders_dir = method_dir / "renders" 65 | renders, gts, image_names = readImages(renders_dir, gt_dir) 66 | 67 | ssims = [] 68 | psnrs = [] 69 | lpipss = [] 70 | 71 | for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"): 72 | ssims.append(ssim(renders[idx], gts[idx])) 73 | psnrs.append(psnr(renders[idx], gts[idx])) 74 | lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg')) 75 | 76 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 77 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 78 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 79 | print("") 80 | 81 | full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(), 82 | "PSNR": torch.tensor(psnrs).mean().item(), 83 | "LPIPS": torch.tensor(lpipss).mean().item()}) 84 | per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 85 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 86 | "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}}) 87 | 88 | with open(scene_dir + "/results.json", 'w') as fp: 89 | json.dump(full_dict[scene_dir], fp, indent=True) 90 | with open(scene_dir + "/per_view.json", 'w') as fp: 91 | json.dump(per_view_dict[scene_dir], fp, indent=True) 92 | # except: 93 | # print("Unable to compute metrics for model", scene_dir) 94 | 95 | if __name__ == "__main__": 96 | device = torch.device("cuda:0") 97 | torch.cuda.set_device(device) 98 | 99 | # Set up command line argument parser 100 | parser = ArgumentParser(description="Training script parameters") 101 | parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[]) 102 | args = parser.parse_args() 103 | evaluate(args.model_paths) 104 | -------------------------------------------------------------------------------- /gaussian_splatting/render.py: -------------------------------------------------------------------------------- 1 | # import pydevd_pycharm 2 | # pydevd_pycharm.settrace('166.111.81.124', port=44112, stdoutToServer=True, stderrToServer=True) 3 | # 4 | # Copyright (C) 2023, Inria 5 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 6 | # All rights reserved. 7 | # 8 | # This software is free for non-commercial, research and evaluation use 9 | # under the terms of the LICENSE.md file. 10 | # 11 | # For inquiries contact george.drettakis@inria.fr 12 | # 13 | 14 | import torch 15 | from scene import Scene 16 | import os 17 | from tqdm import tqdm 18 | from os import makedirs 19 | from gaussian_renderer import render 20 | import torchvision 21 | from utils.general_utils import safe_state 22 | from argparse import ArgumentParser 23 | from arguments import ModelParams, PipelineParams, get_combined_args 24 | from gaussian_renderer import GaussianModel 25 | 26 | def render_set(model_path, name, iteration, views, gaussians, pipeline, background): 27 | render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders") 28 | gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt") 29 | 30 | makedirs(render_path, exist_ok=True) 31 | makedirs(gts_path, exist_ok=True) 32 | 33 | for idx, view in enumerate(tqdm(views, desc="Rendering progress")): 34 | rendering = render(view, gaussians, pipeline, background)["render"] 35 | gt = view.original_image[0:3, :, :] 36 | torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png")) 37 | torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png")) 38 | 39 | def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool): 40 | if not skip_test: 41 | dataset.eval = True 42 | with torch.no_grad(): 43 | gaussians = GaussianModel(dataset.sh_degree) 44 | scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False) 45 | 46 | bg_color = [1,1,1] if dataset.white_background else [0, 0, 0] 47 | background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") 48 | 49 | if not skip_train: 50 | render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background) 51 | 52 | if not skip_test: 53 | render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background) 54 | 55 | if __name__ == "__main__": 56 | # Set up command line argument parser 57 | parser = ArgumentParser(description="Testing script parameters") 58 | model = ModelParams(parser, sentinel=True) 59 | pipeline = PipelineParams(parser) 60 | parser.add_argument("--iteration", default=-1, type=int) 61 | parser.add_argument("--skip_train", action="store_true") 62 | parser.add_argument("--skip_test", action="store_true") 63 | parser.add_argument("--quiet", action="store_true") 64 | args = get_combined_args(parser) 65 | print("Rendering " + args.model_path) 66 | 67 | # Initialize system state (RNG) 68 | safe_state(args.quiet) 69 | 70 | render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test) -------------------------------------------------------------------------------- /gaussian_splatting/scene/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import os 13 | import random 14 | import json 15 | from utils.system_utils import searchForMaxIteration 16 | from scene.dataset_readers import sceneLoadTypeCallbacks 17 | from scene.gaussian_model import GaussianModel 18 | from arguments import ModelParams 19 | from utils.camera_utils import cameraList_from_camInfos, camera_to_JSON 20 | 21 | class Scene: 22 | 23 | gaussians : GaussianModel 24 | 25 | def __init__(self, args : ModelParams, gaussians : GaussianModel, load_iteration=None, shuffle=True, resolution_scales=[1.0]): 26 | """b 27 | :param path: Path to colmap scene main folder. 28 | """ 29 | self.model_path = args.model_path 30 | self.loaded_iter = None 31 | self.gaussians = gaussians 32 | 33 | if load_iteration: 34 | if load_iteration == -1: 35 | self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud")) 36 | else: 37 | self.loaded_iter = load_iteration 38 | print("Loading trained model at iteration {}".format(self.loaded_iter)) 39 | 40 | self.train_cameras = {} 41 | self.test_cameras = {} 42 | 43 | if os.path.exists(os.path.join(args.source_path, "sparse")): 44 | scene_info = sceneLoadTypeCallbacks["Colmap"](args.source_path, args.images, args.eval) 45 | elif os.path.exists(os.path.join(args.source_path, "transforms_train.json")): 46 | print("Found transforms_train.json file, assuming Blender data set!") 47 | scene_info = sceneLoadTypeCallbacks["Blender"](args.source_path, args.white_background, args.eval) 48 | elif os.path.exists(os.path.join(args.source_path, "inputs/sfm_scene.json")): 49 | print("Found sfm_scene.json file, assuming NeILF data set!") 50 | scene_info = sceneLoadTypeCallbacks["NeILF"](args.source_path, args.white_background, args.eval) 51 | else: 52 | assert False, "Could not recognize scene type!" 53 | 54 | if not self.loaded_iter: 55 | with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file: 56 | dest_file.write(src_file.read()) 57 | json_cams = [] 58 | camlist = [] 59 | if scene_info.test_cameras: 60 | camlist.extend(scene_info.test_cameras) 61 | if scene_info.train_cameras: 62 | camlist.extend(scene_info.train_cameras) 63 | for id, cam in enumerate(camlist): 64 | json_cams.append(camera_to_JSON(id, cam)) 65 | with open(os.path.join(self.model_path, "cameras.json"), 'w') as file: 66 | json.dump(json_cams, file) 67 | 68 | if shuffle: 69 | random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling 70 | random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling 71 | 72 | self.cameras_extent = scene_info.nerf_normalization["radius"] 73 | 74 | for resolution_scale in resolution_scales: 75 | print("Loading Training Cameras") 76 | self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args) 77 | print("Loading Test Cameras") 78 | self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args) 79 | 80 | if self.loaded_iter: 81 | self.gaussians.load_ply(os.path.join(self.model_path, 82 | "point_cloud", 83 | "iteration_" + str(self.loaded_iter), 84 | "point_cloud.ply")) 85 | else: 86 | self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) 87 | 88 | def save(self, iteration): 89 | point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) 90 | self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) 91 | 92 | def getTrainCameras(self, scale=1.0): 93 | return self.train_cameras[scale] 94 | 95 | def getTestCameras(self, scale=1.0): 96 | return self.test_cameras[scale] -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/cameras.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/cameras.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/cameras.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/cameras.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/colmap_loader.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/colmap_loader.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/colmap_loader.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/colmap_loader.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/dataset_readers.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/dataset_readers.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/dataset_readers.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/dataset_readers.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/gaussian_model.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/gaussian_model.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/__pycache__/gaussian_model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/scene/__pycache__/gaussian_model.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/scene/cameras.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | from torch import nn 14 | import numpy as np 15 | from utils.graphics_utils import getWorld2View2, getProjectionMatrix 16 | 17 | class Camera(nn.Module): 18 | def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, 19 | image_name, uid, 20 | trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda" 21 | ): 22 | super(Camera, self).__init__() 23 | 24 | self.uid = uid 25 | self.colmap_id = colmap_id 26 | self.R = R 27 | self.T = T 28 | self.FoVx = FoVx 29 | self.FoVy = FoVy 30 | self.image_name = image_name 31 | 32 | try: 33 | self.data_device = torch.device(data_device) 34 | except Exception as e: 35 | print(e) 36 | print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" ) 37 | self.data_device = torch.device("cuda") 38 | 39 | self.original_image = image.clamp(0.0, 1.0).to(self.data_device) 40 | self.image_width = self.original_image.shape[2] 41 | self.image_height = self.original_image.shape[1] 42 | 43 | if gt_alpha_mask is not None: 44 | self.original_image *= gt_alpha_mask.to(self.data_device) 45 | else: 46 | self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device) 47 | 48 | self.zfar = 100.0 49 | self.znear = 0.01 50 | 51 | self.trans = trans 52 | self.scale = scale 53 | 54 | self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda() 55 | self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda() 56 | self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) 57 | self.camera_center = self.world_view_transform.inverse()[3, :3] 58 | 59 | class MiniCam: 60 | def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): 61 | self.image_width = width 62 | self.image_height = height 63 | self.FoVy = fovy 64 | self.FoVx = fovx 65 | self.znear = znear 66 | self.zfar = zfar 67 | self.world_view_transform = world_view_transform 68 | self.full_proj_transform = full_proj_transform 69 | view_inv = torch.inverse(self.world_view_transform) 70 | self.camera_center = view_inv[3][:3] 71 | 72 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | diff_gaussian_rasterization.egg-info/ 3 | dist/ 4 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "third_party/glm"] 2 | path = third_party/glm 3 | url = https://github.com/g-truc/glm.git 4 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | cmake_minimum_required(VERSION 3.20) 13 | 14 | project(DiffRast LANGUAGES CUDA CXX) 15 | 16 | set(CMAKE_CXX_STANDARD 17) 17 | set(CMAKE_CXX_EXTENSIONS OFF) 18 | set(CMAKE_CUDA_STANDARD 17) 19 | 20 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}") 21 | 22 | add_library(CudaRasterizer 23 | cuda_rasterizer/backward.h 24 | cuda_rasterizer/backward.cu 25 | cuda_rasterizer/forward.h 26 | cuda_rasterizer/forward.cu 27 | cuda_rasterizer/auxiliary.h 28 | cuda_rasterizer/rasterizer_impl.cu 29 | cuda_rasterizer/rasterizer_impl.h 30 | cuda_rasterizer/rasterizer.h 31 | ) 32 | 33 | set_target_properties(CudaRasterizer PROPERTIES CUDA_ARCHITECTURES "70;75;86") 34 | 35 | target_include_directories(CudaRasterizer PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/cuda_rasterizer) 36 | target_include_directories(CudaRasterizer PRIVATE third_party/glm ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) 37 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/LICENSE.md: -------------------------------------------------------------------------------- 1 | Gaussian-Splatting License 2 | =========================== 3 | 4 | **Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**. 5 | The *Software* is in the process of being registered with the Agence pour la Protection des 6 | Programmes (APP). 7 | 8 | The *Software* is still being developed by the *Licensor*. 9 | 10 | *Licensor*'s goal is to allow the research community to use, test and evaluate 11 | the *Software*. 12 | 13 | ## 1. Definitions 14 | 15 | *Licensee* means any person or entity that uses the *Software* and distributes 16 | its *Work*. 17 | 18 | *Licensor* means the owners of the *Software*, i.e Inria and MPII 19 | 20 | *Software* means the original work of authorship made available under this 21 | License ie gaussian-splatting. 22 | 23 | *Work* means the *Software* and any additions to or derivative works of the 24 | *Software* that are made available under this License. 25 | 26 | 27 | ## 2. Purpose 28 | This license is intended to define the rights granted to the *Licensee* by 29 | Licensors under the *Software*. 30 | 31 | ## 3. Rights granted 32 | 33 | For the above reasons Licensors have decided to distribute the *Software*. 34 | Licensors grant non-exclusive rights to use the *Software* for research purposes 35 | to research users (both academic and industrial), free of charge, without right 36 | to sublicense.. The *Software* may be used "non-commercially", i.e., for research 37 | and/or evaluation purposes only. 38 | 39 | Subject to the terms and conditions of this License, you are granted a 40 | non-exclusive, royalty-free, license to reproduce, prepare derivative works of, 41 | publicly display, publicly perform and distribute its *Work* and any resulting 42 | derivative works in any form. 43 | 44 | ## 4. Limitations 45 | 46 | **4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do 47 | so under this License, (b) you include a complete copy of this License with 48 | your distribution, and (c) you retain without modification any copyright, 49 | patent, trademark, or attribution notices that are present in the *Work*. 50 | 51 | **4.2 Derivative Works.** You may specify that additional or different terms apply 52 | to the use, reproduction, and distribution of your derivative works of the *Work* 53 | ("Your Terms") only if (a) Your Terms provide that the use limitation in 54 | Section 2 applies to your derivative works, and (b) you identify the specific 55 | derivative works that are subject to Your Terms. Notwithstanding Your Terms, 56 | this License (including the redistribution requirements in Section 3.1) will 57 | continue to apply to the *Work* itself. 58 | 59 | **4.3** Any other use without of prior consent of Licensors is prohibited. Research 60 | users explicitly acknowledge having received from Licensors all information 61 | allowing to appreciate the adequacy between of the *Software* and their needs and 62 | to undertake all necessary precautions for its execution and use. 63 | 64 | **4.4** The *Software* is provided both as a compiled library file and as source 65 | code. In case of using the *Software* for a publication or other results obtained 66 | through the use of the *Software*, users are strongly encouraged to cite the 67 | corresponding publications as explained in the documentation of the *Software*. 68 | 69 | ## 5. Disclaimer 70 | 71 | THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES 72 | WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY 73 | UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL 74 | CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES 75 | OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL 76 | USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR 77 | ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE 78 | AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 79 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE 80 | GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) 81 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 82 | LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR 83 | IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*. 84 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/README.md: -------------------------------------------------------------------------------- 1 | # Differential Gaussian Rasterization 2 | 3 | Used as the rasterization engine for the paper "3D Gaussian Splatting for Real-Time Rendering of Radiance Fields". If you can make use of it in your own research, please be so kind to cite us. 4 | 5 |
6 |
7 |

BibTeX

8 | 
@Article{kerbl3Dgaussians,
 9 |       author       = {Kerbl, Bernhard and Kopanas, Georgios and Leimk{\"u}hler, Thomas and Drettakis, George},
10 |       title        = {3D Gaussian Splatting for Real-Time Radiance Field Rendering},
11 |       journal      = {ACM Transactions on Graphics},
12 |       number       = {4},
13 |       volume       = {42},
14 |       month        = {July},
15 |       year         = {2023},
16 |       url          = {https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/}
17 | }
18 |
19 |
-------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/auxiliary.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef CUDA_RASTERIZER_AUXILIARY_H_INCLUDED 13 | #define CUDA_RASTERIZER_AUXILIARY_H_INCLUDED 14 | 15 | #include "config.h" 16 | #include "stdio.h" 17 | 18 | #define BLOCK_SIZE (BLOCK_X * BLOCK_Y) 19 | #define NUM_WARPS (BLOCK_SIZE/32) 20 | // Spherical harmonics coefficients 21 | __device__ const float SH_C0 = 0.28209479177387814f; 22 | __device__ const float SH_C1 = 0.4886025119029199f; 23 | __device__ const float SH_C2[] = { 24 | 1.0925484305920792f, 25 | -1.0925484305920792f, 26 | 0.31539156525252005f, 27 | -1.0925484305920792f, 28 | 0.5462742152960396f 29 | }; 30 | __device__ const float SH_C3[] = { 31 | -0.5900435899266435f, 32 | 2.890611442640554f, 33 | -0.4570457994644658f, 34 | 0.3731763325901154f, 35 | -0.4570457994644658f, 36 | 1.445305721320277f, 37 | -0.5900435899266435f 38 | }; 39 | 40 | __forceinline__ __device__ float ndc2Pix(float v, int S) 41 | { 42 | return ((v + 1.0) * S - 1.0) * 0.5; 43 | } 44 | 45 | __forceinline__ __device__ void getRect(const float2 p, int max_radius, uint2& rect_min, uint2& rect_max, dim3 grid) 46 | { 47 | rect_min = { 48 | min(grid.x, max((int)0, (int)((p.x - max_radius) / BLOCK_X))), 49 | min(grid.y, max((int)0, (int)((p.y - max_radius) / BLOCK_Y))) 50 | }; 51 | rect_max = { 52 | min(grid.x, max((int)0, (int)((p.x + max_radius + BLOCK_X - 1) / BLOCK_X))), 53 | min(grid.y, max((int)0, (int)((p.y + max_radius + BLOCK_Y - 1) / BLOCK_Y))) 54 | }; 55 | } 56 | 57 | __forceinline__ __device__ void getRect(const float2 p, int2 ext_rect, uint2& rect_min, uint2& rect_max, dim3 grid) 58 | { 59 | rect_min = { 60 | min(grid.x, max((int)0, (int)((p.x - ext_rect.x) / BLOCK_X))), 61 | min(grid.y, max((int)0, (int)((p.y - ext_rect.y) / BLOCK_Y))) 62 | }; 63 | rect_max = { 64 | min(grid.x, max((int)0, (int)((p.x + ext_rect.x + BLOCK_X - 1) / BLOCK_X))), 65 | min(grid.y, max((int)0, (int)((p.y + ext_rect.y + BLOCK_Y - 1) / BLOCK_Y))) 66 | }; 67 | } 68 | 69 | 70 | __forceinline__ __device__ float3 transformPoint4x3(const float3& p, const float* matrix) 71 | { 72 | float3 transformed = { 73 | matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12], 74 | matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13], 75 | matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14], 76 | }; 77 | return transformed; 78 | } 79 | 80 | __forceinline__ __device__ float4 transformPoint4x4(const float3& p, const float* matrix) 81 | { 82 | float4 transformed = { 83 | matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12], 84 | matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13], 85 | matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14], 86 | matrix[3] * p.x + matrix[7] * p.y + matrix[11] * p.z + matrix[15] 87 | }; 88 | return transformed; 89 | } 90 | 91 | __forceinline__ __device__ float3 transformVec4x3(const float3& p, const float* matrix) 92 | { 93 | float3 transformed = { 94 | matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z, 95 | matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z, 96 | matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z, 97 | }; 98 | return transformed; 99 | } 100 | 101 | __forceinline__ __device__ float3 transformVec4x3Transpose(const float3& p, const float* matrix) 102 | { 103 | float3 transformed = { 104 | matrix[0] * p.x + matrix[1] * p.y + matrix[2] * p.z, 105 | matrix[4] * p.x + matrix[5] * p.y + matrix[6] * p.z, 106 | matrix[8] * p.x + matrix[9] * p.y + matrix[10] * p.z, 107 | }; 108 | return transformed; 109 | } 110 | 111 | __forceinline__ __device__ float dnormvdz(float3 v, float3 dv) 112 | { 113 | float sum2 = v.x * v.x + v.y * v.y + v.z * v.z; 114 | float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2); 115 | float dnormvdz = (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) * invsum32; 116 | return dnormvdz; 117 | } 118 | 119 | __forceinline__ __device__ float3 dnormvdv(float3 v, float3 dv) 120 | { 121 | float sum2 = v.x * v.x + v.y * v.y + v.z * v.z; 122 | float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2); 123 | 124 | float3 dnormvdv; 125 | dnormvdv.x = ((+sum2 - v.x * v.x) * dv.x - v.y * v.x * dv.y - v.z * v.x * dv.z) * invsum32; 126 | dnormvdv.y = (-v.x * v.y * dv.x + (sum2 - v.y * v.y) * dv.y - v.z * v.y * dv.z) * invsum32; 127 | dnormvdv.z = (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) * invsum32; 128 | return dnormvdv; 129 | } 130 | 131 | __forceinline__ __device__ float4 dnormvdv(float4 v, float4 dv) 132 | { 133 | float sum2 = v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w; 134 | float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2); 135 | 136 | float4 vdv = { v.x * dv.x, v.y * dv.y, v.z * dv.z, v.w * dv.w }; 137 | float vdv_sum = vdv.x + vdv.y + vdv.z + vdv.w; 138 | float4 dnormvdv; 139 | dnormvdv.x = ((sum2 - v.x * v.x) * dv.x - v.x * (vdv_sum - vdv.x)) * invsum32; 140 | dnormvdv.y = ((sum2 - v.y * v.y) * dv.y - v.y * (vdv_sum - vdv.y)) * invsum32; 141 | dnormvdv.z = ((sum2 - v.z * v.z) * dv.z - v.z * (vdv_sum - vdv.z)) * invsum32; 142 | dnormvdv.w = ((sum2 - v.w * v.w) * dv.w - v.w * (vdv_sum - vdv.w)) * invsum32; 143 | return dnormvdv; 144 | } 145 | 146 | __forceinline__ __device__ float sigmoid(float x) 147 | { 148 | return 1.0f / (1.0f + expf(-x)); 149 | } 150 | 151 | __forceinline__ __device__ bool in_frustum(int idx, 152 | const float* orig_points, 153 | const float* viewmatrix, 154 | const float* projmatrix, 155 | bool prefiltered, 156 | float3& p_view) 157 | { 158 | float3 p_orig = { orig_points[3 * idx], orig_points[3 * idx + 1], orig_points[3 * idx + 2] }; 159 | 160 | // Bring points to screen space 161 | float4 p_hom = transformPoint4x4(p_orig, projmatrix); 162 | float p_w = 1.0f / (p_hom.w + 0.0000001f); 163 | float3 p_proj = { p_hom.x * p_w, p_hom.y * p_w, p_hom.z * p_w }; 164 | p_view = transformPoint4x3(p_orig, viewmatrix); 165 | 166 | if (p_view.z <= 0.2f)// || ((p_proj.x < -1.3 || p_proj.x > 1.3 || p_proj.y < -1.3 || p_proj.y > 1.3))) 167 | { 168 | if (prefiltered) 169 | { 170 | printf("Point is filtered although prefiltered is set. This shouldn't happen!"); 171 | __trap(); 172 | } 173 | return false; 174 | } 175 | return true; 176 | } 177 | 178 | #define CHECK_CUDA(A, debug) \ 179 | A; if(debug) { \ 180 | auto ret = cudaDeviceSynchronize(); \ 181 | if (ret != cudaSuccess) { \ 182 | std::cerr << "\n[CUDA ERROR] in " << __FILE__ << "\nLine " << __LINE__ << ": " << cudaGetErrorString(ret); \ 183 | throw std::runtime_error(cudaGetErrorString(ret)); \ 184 | } \ 185 | } 186 | 187 | #endif 188 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/backward.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef CUDA_RASTERIZER_BACKWARD_H_INCLUDED 13 | #define CUDA_RASTERIZER_BACKWARD_H_INCLUDED 14 | 15 | #include 16 | #include "cuda_runtime.h" 17 | #include "device_launch_parameters.h" 18 | #define GLM_FORCE_CUDA 19 | #include 20 | 21 | namespace BACKWARD 22 | { 23 | void render( 24 | const dim3 grid, dim3 block, 25 | const uint2* ranges, 26 | const uint32_t* point_list, 27 | int W, int H, 28 | const float* bg_color, 29 | const float2* means2D, 30 | const float4* conic_opacity, 31 | const float* colors, 32 | const float* depths, 33 | const float* final_Ts, 34 | const uint32_t* n_contrib, 35 | const float* dL_dpixels, 36 | const float* dL_invdepths, 37 | float3* dL_dmean2D, 38 | float4* dL_dconic2D, 39 | float* dL_dopacity, 40 | float* dL_dcolors, 41 | float* dL_dinvdepths); 42 | 43 | void preprocess( 44 | int P, int D, int M, 45 | const float3* means, 46 | const int* radii, 47 | const float* shs, 48 | const bool* clamped, 49 | const float* opacities, 50 | const glm::vec3* scales, 51 | const glm::vec4* rotations, 52 | const float scale_modifier, 53 | const float* cov3Ds, 54 | const float* view, 55 | const float* proj, 56 | const float focal_x, float focal_y, 57 | const float tan_fovx, float tan_fovy, 58 | const glm::vec3* campos, 59 | const float3* dL_dmean2D, 60 | const float* dL_dconics, 61 | const float* dL_dinvdepth, 62 | float* dL_dopacity, 63 | glm::vec3* dL_dmeans, 64 | float* dL_dcolor, 65 | float* dL_dcov3D, 66 | float* dL_dsh, 67 | glm::vec3* dL_dscale, 68 | glm::vec4* dL_drot, 69 | bool antialiasing); 70 | } 71 | 72 | #endif 73 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/config.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef CUDA_RASTERIZER_CONFIG_H_INCLUDED 13 | #define CUDA_RASTERIZER_CONFIG_H_INCLUDED 14 | 15 | #define NUM_CHANNELS 3 // Default 3, RGB 16 | #define BLOCK_X 16 17 | #define BLOCK_Y 16 18 | 19 | #endif -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/forward.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef CUDA_RASTERIZER_FORWARD_H_INCLUDED 13 | #define CUDA_RASTERIZER_FORWARD_H_INCLUDED 14 | 15 | #include 16 | #include "cuda_runtime.h" 17 | #include "device_launch_parameters.h" 18 | #define GLM_FORCE_CUDA 19 | #include 20 | 21 | namespace FORWARD 22 | { 23 | // Perform initial steps for each Gaussian prior to rasterization. 24 | void preprocess(int P, int D, int M, 25 | const float* orig_points, 26 | const glm::vec3* scales, 27 | const float scale_modifier, 28 | const glm::vec4* rotations, 29 | const float* opacities, 30 | const float* shs, 31 | bool* clamped, 32 | const float* cov3D_precomp, 33 | const float* colors_precomp, 34 | const float* viewmatrix, 35 | const float* projmatrix, 36 | const glm::vec3* cam_pos, 37 | const int W, int H, 38 | const float focal_x, float focal_y, 39 | const float tan_fovx, float tan_fovy, 40 | int* radii, 41 | float2* points_xy_image, 42 | float* depths, 43 | float* cov3Ds, 44 | float* colors, 45 | float4* conic_opacity, 46 | const dim3 grid, 47 | uint32_t* tiles_touched, 48 | bool prefiltered, 49 | bool antialiasing); 50 | 51 | // Main rasterization method. 52 | void render( 53 | const dim3 grid, dim3 block, 54 | const uint2* ranges, 55 | const uint32_t* point_list, 56 | int W, int H, 57 | const float2* points_xy_image, 58 | const float* features, 59 | const float4* conic_opacity, 60 | float* final_T, 61 | uint32_t* n_contrib, 62 | const float* bg_color, 63 | float* out_color, 64 | float* depths, 65 | float* depth); 66 | } 67 | 68 | 69 | #endif 70 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/rasterizer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef CUDA_RASTERIZER_H_INCLUDED 13 | #define CUDA_RASTERIZER_H_INCLUDED 14 | 15 | #include 16 | #include 17 | 18 | namespace CudaRasterizer 19 | { 20 | class Rasterizer 21 | { 22 | public: 23 | 24 | static void markVisible( 25 | int P, 26 | float* means3D, 27 | float* viewmatrix, 28 | float* projmatrix, 29 | bool* present); 30 | 31 | static int forward( 32 | std::function geometryBuffer, 33 | std::function binningBuffer, 34 | std::function imageBuffer, 35 | const int P, int D, int M, 36 | const float* background, 37 | const int width, int height, 38 | const float* means3D, 39 | const float* shs, 40 | const float* colors_precomp, 41 | const float* opacities, 42 | const float* scales, 43 | const float scale_modifier, 44 | const float* rotations, 45 | const float* cov3D_precomp, 46 | const float* viewmatrix, 47 | const float* projmatrix, 48 | const float* cam_pos, 49 | const float tan_fovx, float tan_fovy, 50 | const bool prefiltered, 51 | float* out_color, 52 | float* depth, 53 | bool antialiasing, 54 | int* radii = nullptr, 55 | bool debug = false); 56 | 57 | static void backward( 58 | const int P, int D, int M, int R, 59 | const float* background, 60 | const int width, int height, 61 | const float* means3D, 62 | const float* shs, 63 | const float* colors_precomp, 64 | const float* opacities, 65 | const float* scales, 66 | const float scale_modifier, 67 | const float* rotations, 68 | const float* cov3D_precomp, 69 | const float* viewmatrix, 70 | const float* projmatrix, 71 | const float* campos, 72 | const float tan_fovx, float tan_fovy, 73 | const int* radii, 74 | char* geom_buffer, 75 | char* binning_buffer, 76 | char* image_buffer, 77 | const float* dL_dpix, 78 | const float* dL_invdepths, 79 | float* dL_dmean2D, 80 | float* dL_dconic, 81 | float* dL_dopacity, 82 | float* dL_dcolor, 83 | float* dL_dinvdepth, 84 | float* dL_dmean3D, 85 | float* dL_dcov3D, 86 | float* dL_dsh, 87 | float* dL_dscale, 88 | float* dL_drot, 89 | bool antialiasing, 90 | bool debug); 91 | }; 92 | }; 93 | 94 | #endif 95 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/cuda_rasterizer/rasterizer_impl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #pragma once 13 | 14 | #include 15 | #include 16 | #include "rasterizer.h" 17 | #include 18 | 19 | namespace CudaRasterizer 20 | { 21 | template 22 | static void obtain(char*& chunk, T*& ptr, std::size_t count, std::size_t alignment) 23 | { 24 | std::size_t offset = (reinterpret_cast(chunk) + alignment - 1) & ~(alignment - 1); 25 | ptr = reinterpret_cast(offset); 26 | chunk = reinterpret_cast(ptr + count); 27 | } 28 | 29 | struct GeometryState 30 | { 31 | size_t scan_size; 32 | float* depths; 33 | char* scanning_space; 34 | bool* clamped; 35 | int* internal_radii; 36 | float2* means2D; 37 | float* cov3D; 38 | float4* conic_opacity; 39 | float* rgb; 40 | uint32_t* point_offsets; 41 | uint32_t* tiles_touched; 42 | 43 | static GeometryState fromChunk(char*& chunk, size_t P); 44 | }; 45 | 46 | struct ImageState 47 | { 48 | uint2* ranges; 49 | uint32_t* n_contrib; 50 | float* accum_alpha; 51 | 52 | static ImageState fromChunk(char*& chunk, size_t N); 53 | }; 54 | 55 | struct BinningState 56 | { 57 | size_t sorting_size; 58 | uint64_t* point_list_keys_unsorted; 59 | uint64_t* point_list_keys; 60 | uint32_t* point_list_unsorted; 61 | uint32_t* point_list; 62 | char* list_sorting_space; 63 | 64 | static BinningState fromChunk(char*& chunk, size_t P); 65 | }; 66 | 67 | template 68 | size_t required(size_t P) 69 | { 70 | char* size = nullptr; 71 | T::fromChunk(size, P); 72 | return ((size_t)size) + 128; 73 | } 74 | }; -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from typing import NamedTuple 13 | import torch.nn as nn 14 | import torch 15 | from . import _C 16 | 17 | def cpu_deep_copy_tuple(input_tuple): 18 | copied_tensors = [item.cpu().clone() if isinstance(item, torch.Tensor) else item for item in input_tuple] 19 | return tuple(copied_tensors) 20 | 21 | def rasterize_gaussians( 22 | means3D, 23 | means2D, 24 | sh, 25 | colors_precomp, 26 | opacities, 27 | scales, 28 | rotations, 29 | cov3Ds_precomp, 30 | raster_settings, 31 | ): 32 | return _RasterizeGaussians.apply( 33 | means3D, 34 | means2D, 35 | sh, 36 | colors_precomp, 37 | opacities, 38 | scales, 39 | rotations, 40 | cov3Ds_precomp, 41 | raster_settings, 42 | ) 43 | 44 | class _RasterizeGaussians(torch.autograd.Function): 45 | @staticmethod 46 | def forward( 47 | ctx, 48 | means3D, 49 | means2D, 50 | sh, 51 | colors_precomp, 52 | opacities, 53 | scales, 54 | rotations, 55 | cov3Ds_precomp, 56 | raster_settings, 57 | ): 58 | 59 | # Restructure arguments the way that the C++ lib expects them 60 | args = ( 61 | raster_settings.bg, 62 | means3D, 63 | colors_precomp, 64 | opacities, 65 | scales, 66 | rotations, 67 | raster_settings.scale_modifier, 68 | cov3Ds_precomp, 69 | raster_settings.viewmatrix, 70 | raster_settings.projmatrix, 71 | raster_settings.tanfovx, 72 | raster_settings.tanfovy, 73 | raster_settings.image_height, 74 | raster_settings.image_width, 75 | sh, 76 | raster_settings.sh_degree, 77 | raster_settings.campos, 78 | raster_settings.prefiltered, 79 | raster_settings.antialiasing, 80 | raster_settings.debug 81 | ) 82 | 83 | # Invoke C++/CUDA rasterizer 84 | num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer, invdepths = _C.rasterize_gaussians(*args) 85 | 86 | # Keep relevant tensors for backward 87 | ctx.raster_settings = raster_settings 88 | ctx.num_rendered = num_rendered 89 | ctx.save_for_backward(colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, opacities, geomBuffer, binningBuffer, imgBuffer) 90 | return color, radii, invdepths 91 | 92 | @staticmethod 93 | def backward(ctx, grad_out_color, _, grad_out_depth): 94 | 95 | # Restore necessary values from context 96 | num_rendered = ctx.num_rendered 97 | raster_settings = ctx.raster_settings 98 | colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, opacities, geomBuffer, binningBuffer, imgBuffer = ctx.saved_tensors 99 | 100 | # Restructure args as C++ method expects them 101 | args = (raster_settings.bg, 102 | means3D, 103 | radii, 104 | colors_precomp, 105 | opacities, 106 | scales, 107 | rotations, 108 | raster_settings.scale_modifier, 109 | cov3Ds_precomp, 110 | raster_settings.viewmatrix, 111 | raster_settings.projmatrix, 112 | raster_settings.tanfovx, 113 | raster_settings.tanfovy, 114 | grad_out_color, 115 | grad_out_depth, 116 | sh, 117 | raster_settings.sh_degree, 118 | raster_settings.campos, 119 | geomBuffer, 120 | num_rendered, 121 | binningBuffer, 122 | imgBuffer, 123 | raster_settings.antialiasing, 124 | raster_settings.debug) 125 | 126 | # Compute gradients for relevant tensors by invoking backward method 127 | grad_means2D, grad_colors_precomp, grad_opacities, grad_means3D, grad_cov3Ds_precomp, grad_sh, grad_scales, grad_rotations = _C.rasterize_gaussians_backward(*args) 128 | 129 | grads = ( 130 | grad_means3D, 131 | grad_means2D, 132 | grad_sh, 133 | grad_colors_precomp, 134 | grad_opacities, 135 | grad_scales, 136 | grad_rotations, 137 | grad_cov3Ds_precomp, 138 | None, 139 | ) 140 | 141 | return grads 142 | 143 | class GaussianRasterizationSettings(NamedTuple): 144 | image_height: int 145 | image_width: int 146 | tanfovx : float 147 | tanfovy : float 148 | bg : torch.Tensor 149 | scale_modifier : float 150 | viewmatrix : torch.Tensor 151 | projmatrix : torch.Tensor 152 | sh_degree : int 153 | campos : torch.Tensor 154 | prefiltered : bool 155 | debug : bool 156 | antialiasing : bool 157 | 158 | class GaussianRasterizer(nn.Module): 159 | def __init__(self, raster_settings): 160 | super().__init__() 161 | self.raster_settings = raster_settings 162 | 163 | def markVisible(self, positions): 164 | # Mark visible points (based on frustum culling for camera) with a boolean 165 | with torch.no_grad(): 166 | raster_settings = self.raster_settings 167 | visible = _C.mark_visible( 168 | positions, 169 | raster_settings.viewmatrix, 170 | raster_settings.projmatrix) 171 | 172 | return visible 173 | 174 | def forward(self, means3D, means2D, opacities, shs = None, colors_precomp = None, scales = None, rotations = None, cov3D_precomp = None): 175 | 176 | raster_settings = self.raster_settings 177 | 178 | if (shs is None and colors_precomp is None) or (shs is not None and colors_precomp is not None): 179 | raise Exception('Please provide excatly one of either SHs or precomputed colors!') 180 | 181 | if ((scales is None or rotations is None) and cov3D_precomp is None) or ((scales is not None or rotations is not None) and cov3D_precomp is not None): 182 | raise Exception('Please provide exactly one of either scale/rotation pair or precomputed 3D covariance!') 183 | 184 | if shs is None: 185 | shs = torch.Tensor([]) 186 | if colors_precomp is None: 187 | colors_precomp = torch.Tensor([]) 188 | 189 | if scales is None: 190 | scales = torch.Tensor([]) 191 | if rotations is None: 192 | rotations = torch.Tensor([]) 193 | if cov3D_precomp is None: 194 | cov3D_precomp = torch.Tensor([]) 195 | 196 | # Invoke C++/CUDA rasterization routine 197 | return rasterize_gaussians( 198 | means3D, 199 | means2D, 200 | shs, 201 | colors_precomp, 202 | opacities, 203 | scales, 204 | rotations, 205 | cov3D_precomp, 206 | raster_settings, 207 | ) 208 | 209 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/ext.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #include 13 | #include "rasterize_points.h" 14 | 15 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 16 | m.def("rasterize_gaussians", &RasterizeGaussiansCUDA); 17 | m.def("rasterize_gaussians_backward", &RasterizeGaussiansBackwardCUDA); 18 | m.def("mark_visible", &markVisible); 19 | } -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/rasterize_points.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include "cuda_rasterizer/config.h" 22 | #include "cuda_rasterizer/rasterizer.h" 23 | #include 24 | #include 25 | #include 26 | 27 | std::function resizeFunctional(torch::Tensor& t) { 28 | auto lambda = [&t](size_t N) { 29 | t.resize_({(long long)N}); 30 | return reinterpret_cast(t.contiguous().data_ptr()); 31 | }; 32 | return lambda; 33 | } 34 | 35 | std::tuple 36 | RasterizeGaussiansCUDA( 37 | const torch::Tensor& background, 38 | const torch::Tensor& means3D, 39 | const torch::Tensor& colors, 40 | const torch::Tensor& opacity, 41 | const torch::Tensor& scales, 42 | const torch::Tensor& rotations, 43 | const float scale_modifier, 44 | const torch::Tensor& cov3D_precomp, 45 | const torch::Tensor& viewmatrix, 46 | const torch::Tensor& projmatrix, 47 | const float tan_fovx, 48 | const float tan_fovy, 49 | const int image_height, 50 | const int image_width, 51 | const torch::Tensor& sh, 52 | const int degree, 53 | const torch::Tensor& campos, 54 | const bool prefiltered, 55 | const bool antialiasing, 56 | const bool debug) 57 | { 58 | if (means3D.ndimension() != 2 || means3D.size(1) != 3) { 59 | AT_ERROR("means3D must have dimensions (num_points, 3)"); 60 | } 61 | 62 | const int P = means3D.size(0); 63 | const int H = image_height; 64 | const int W = image_width; 65 | 66 | auto int_opts = means3D.options().dtype(torch::kInt32); 67 | auto float_opts = means3D.options().dtype(torch::kFloat32); 68 | 69 | torch::Tensor out_color = torch::full({NUM_CHANNELS, H, W}, 0.0, float_opts); 70 | torch::Tensor out_invdepth = torch::full({0, H, W}, 0.0, float_opts); 71 | float* out_invdepthptr = nullptr; 72 | 73 | out_invdepth = torch::full({1, H, W}, 0.0, float_opts).contiguous(); 74 | out_invdepthptr = out_invdepth.data(); 75 | 76 | torch::Tensor radii = torch::full({P}, 0, means3D.options().dtype(torch::kInt32)); 77 | 78 | torch::Device device(torch::kCUDA); 79 | torch::TensorOptions options(torch::kByte); 80 | torch::Tensor geomBuffer = torch::empty({0}, options.device(device)); 81 | torch::Tensor binningBuffer = torch::empty({0}, options.device(device)); 82 | torch::Tensor imgBuffer = torch::empty({0}, options.device(device)); 83 | std::function geomFunc = resizeFunctional(geomBuffer); 84 | std::function binningFunc = resizeFunctional(binningBuffer); 85 | std::function imgFunc = resizeFunctional(imgBuffer); 86 | 87 | int rendered = 0; 88 | if(P != 0) 89 | { 90 | int M = 0; 91 | if(sh.size(0) != 0) 92 | { 93 | M = sh.size(1); 94 | } 95 | 96 | rendered = CudaRasterizer::Rasterizer::forward( 97 | geomFunc, 98 | binningFunc, 99 | imgFunc, 100 | P, degree, M, 101 | background.contiguous().data(), 102 | W, H, 103 | means3D.contiguous().data(), 104 | sh.contiguous().data_ptr(), 105 | colors.contiguous().data(), 106 | opacity.contiguous().data(), 107 | scales.contiguous().data_ptr(), 108 | scale_modifier, 109 | rotations.contiguous().data_ptr(), 110 | cov3D_precomp.contiguous().data(), 111 | viewmatrix.contiguous().data(), 112 | projmatrix.contiguous().data(), 113 | campos.contiguous().data(), 114 | tan_fovx, 115 | tan_fovy, 116 | prefiltered, 117 | out_color.contiguous().data(), 118 | out_invdepthptr, 119 | antialiasing, 120 | radii.contiguous().data(), 121 | debug); 122 | } 123 | return std::make_tuple(rendered, out_color, radii, geomBuffer, binningBuffer, imgBuffer, out_invdepth); 124 | } 125 | 126 | std::tuple 127 | RasterizeGaussiansBackwardCUDA( 128 | const torch::Tensor& background, 129 | const torch::Tensor& means3D, 130 | const torch::Tensor& radii, 131 | const torch::Tensor& colors, 132 | const torch::Tensor& opacities, 133 | const torch::Tensor& scales, 134 | const torch::Tensor& rotations, 135 | const float scale_modifier, 136 | const torch::Tensor& cov3D_precomp, 137 | const torch::Tensor& viewmatrix, 138 | const torch::Tensor& projmatrix, 139 | const float tan_fovx, 140 | const float tan_fovy, 141 | const torch::Tensor& dL_dout_color, 142 | const torch::Tensor& dL_dout_invdepth, 143 | const torch::Tensor& sh, 144 | const int degree, 145 | const torch::Tensor& campos, 146 | const torch::Tensor& geomBuffer, 147 | const int R, 148 | const torch::Tensor& binningBuffer, 149 | const torch::Tensor& imageBuffer, 150 | const bool antialiasing, 151 | const bool debug) 152 | { 153 | const int P = means3D.size(0); 154 | const int H = dL_dout_color.size(1); 155 | const int W = dL_dout_color.size(2); 156 | 157 | int M = 0; 158 | if(sh.size(0) != 0) 159 | { 160 | M = sh.size(1); 161 | } 162 | 163 | torch::Tensor dL_dmeans3D = torch::zeros({P, 3}, means3D.options()); 164 | torch::Tensor dL_dmeans2D = torch::zeros({P, 3}, means3D.options()); 165 | torch::Tensor dL_dcolors = torch::zeros({P, NUM_CHANNELS}, means3D.options()); 166 | torch::Tensor dL_dconic = torch::zeros({P, 2, 2}, means3D.options()); 167 | torch::Tensor dL_dopacity = torch::zeros({P, 1}, means3D.options()); 168 | torch::Tensor dL_dcov3D = torch::zeros({P, 6}, means3D.options()); 169 | torch::Tensor dL_dsh = torch::zeros({P, M, 3}, means3D.options()); 170 | torch::Tensor dL_dscales = torch::zeros({P, 3}, means3D.options()); 171 | torch::Tensor dL_drotations = torch::zeros({P, 4}, means3D.options()); 172 | torch::Tensor dL_dinvdepths = torch::zeros({0, 1}, means3D.options()); 173 | 174 | float* dL_dinvdepthsptr = nullptr; 175 | float* dL_dout_invdepthptr = nullptr; 176 | if(dL_dout_invdepth.size(0) != 0) 177 | { 178 | dL_dinvdepths = torch::zeros({P, 1}, means3D.options()); 179 | dL_dinvdepths = dL_dinvdepths.contiguous(); 180 | dL_dinvdepthsptr = dL_dinvdepths.data(); 181 | dL_dout_invdepthptr = dL_dout_invdepth.data(); 182 | } 183 | 184 | if(P != 0) 185 | { 186 | CudaRasterizer::Rasterizer::backward(P, degree, M, R, 187 | background.contiguous().data(), 188 | W, H, 189 | means3D.contiguous().data(), 190 | sh.contiguous().data(), 191 | colors.contiguous().data(), 192 | opacities.contiguous().data(), 193 | scales.data_ptr(), 194 | scale_modifier, 195 | rotations.data_ptr(), 196 | cov3D_precomp.contiguous().data(), 197 | viewmatrix.contiguous().data(), 198 | projmatrix.contiguous().data(), 199 | campos.contiguous().data(), 200 | tan_fovx, 201 | tan_fovy, 202 | radii.contiguous().data(), 203 | reinterpret_cast(geomBuffer.contiguous().data_ptr()), 204 | reinterpret_cast(binningBuffer.contiguous().data_ptr()), 205 | reinterpret_cast(imageBuffer.contiguous().data_ptr()), 206 | dL_dout_color.contiguous().data(), 207 | dL_dout_invdepthptr, 208 | dL_dmeans2D.contiguous().data(), 209 | dL_dconic.contiguous().data(), 210 | dL_dopacity.contiguous().data(), 211 | dL_dcolors.contiguous().data(), 212 | dL_dinvdepthsptr, 213 | dL_dmeans3D.contiguous().data(), 214 | dL_dcov3D.contiguous().data(), 215 | dL_dsh.contiguous().data(), 216 | dL_dscales.contiguous().data(), 217 | dL_drotations.contiguous().data(), 218 | antialiasing, 219 | debug); 220 | } 221 | 222 | return std::make_tuple(dL_dmeans2D, dL_dcolors, dL_dopacity, dL_dmeans3D, dL_dcov3D, dL_dsh, dL_dscales, dL_drotations); 223 | } 224 | 225 | torch::Tensor markVisible( 226 | torch::Tensor& means3D, 227 | torch::Tensor& viewmatrix, 228 | torch::Tensor& projmatrix) 229 | { 230 | const int P = means3D.size(0); 231 | 232 | torch::Tensor present = torch::full({P}, false, means3D.options().dtype(at::kBool)); 233 | 234 | if(P != 0) 235 | { 236 | CudaRasterizer::Rasterizer::markVisible(P, 237 | means3D.contiguous().data(), 238 | viewmatrix.contiguous().data(), 239 | projmatrix.contiguous().data(), 240 | present.contiguous().data()); 241 | } 242 | 243 | return present; 244 | } 245 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/rasterize_points.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #pragma once 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | std::tuple 19 | RasterizeGaussiansCUDA( 20 | const torch::Tensor& background, 21 | const torch::Tensor& means3D, 22 | const torch::Tensor& colors, 23 | const torch::Tensor& opacity, 24 | const torch::Tensor& scales, 25 | const torch::Tensor& rotations, 26 | const float scale_modifier, 27 | const torch::Tensor& cov3D_precomp, 28 | const torch::Tensor& viewmatrix, 29 | const torch::Tensor& projmatrix, 30 | const float tan_fovx, 31 | const float tan_fovy, 32 | const int image_height, 33 | const int image_width, 34 | const torch::Tensor& sh, 35 | const int degree, 36 | const torch::Tensor& campos, 37 | const bool prefiltered, 38 | const bool antialiasing, 39 | const bool debug); 40 | 41 | std::tuple 42 | RasterizeGaussiansBackwardCUDA( 43 | const torch::Tensor& background, 44 | const torch::Tensor& means3D, 45 | const torch::Tensor& radii, 46 | const torch::Tensor& colors, 47 | const torch::Tensor& scales, 48 | const torch::Tensor& opacities, 49 | const torch::Tensor& rotations, 50 | const float scale_modifier, 51 | const torch::Tensor& cov3D_precomp, 52 | const torch::Tensor& viewmatrix, 53 | const torch::Tensor& projmatrix, 54 | const float tan_fovx, 55 | const float tan_fovy, 56 | const torch::Tensor& dL_dout_color, 57 | const torch::Tensor& dL_dout_invdepth, 58 | const torch::Tensor& sh, 59 | const int degree, 60 | const torch::Tensor& campos, 61 | const torch::Tensor& geomBuffer, 62 | const int R, 63 | const torch::Tensor& binningBuffer, 64 | const torch::Tensor& imageBuffer, 65 | const bool antialiasing, 66 | const bool debug); 67 | 68 | torch::Tensor markVisible( 69 | torch::Tensor& means3D, 70 | torch::Tensor& viewmatrix, 71 | torch::Tensor& projmatrix); 72 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/diff-gaussian-rasterization/setup.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from setuptools import setup 13 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 14 | import os 15 | os.path.dirname(os.path.abspath(__file__)) 16 | 17 | setup( 18 | name="diff_gaussian_rasterization", 19 | packages=['diff_gaussian_rasterization'], 20 | ext_modules=[ 21 | CUDAExtension( 22 | name="diff_gaussian_rasterization._C", 23 | sources=[ 24 | "cuda_rasterizer/rasterizer_impl.cu", 25 | "cuda_rasterizer/forward.cu", 26 | "cuda_rasterizer/backward.cu", 27 | "rasterize_points.cu", 28 | "ext.cpp"], 29 | extra_compile_args={"nvcc": ["-I" + os.path.join(os.path.dirname(os.path.abspath(__file__)), "third_party/glm/")]}) 30 | ], 31 | cmdclass={ 32 | 'build_ext': BuildExtension 33 | } 34 | ) 35 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/ext.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #include 13 | #include "spatial.h" 14 | 15 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 16 | m.def("distCUDA2", &distCUDA2); 17 | } 18 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/setup.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from setuptools import setup 13 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 14 | import os 15 | 16 | cxx_compiler_flags = [] 17 | 18 | if os.name == 'nt': 19 | cxx_compiler_flags.append("/wd4624") 20 | 21 | setup( 22 | name="simple_knn", 23 | ext_modules=[ 24 | CUDAExtension( 25 | name="simple_knn._C", 26 | sources=[ 27 | "spatial.cu", 28 | "simple_knn.cu", 29 | "ext.cpp"], 30 | extra_compile_args={"nvcc": [], "cxx": cxx_compiler_flags}) 31 | ], 32 | cmdclass={ 33 | 'build_ext': BuildExtension 34 | } 35 | ) 36 | -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/simple_knn.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #define BOX_SIZE 1024 13 | 14 | #include "cuda_runtime.h" 15 | #include "device_launch_parameters.h" 16 | #include "simple_knn.h" 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include 22 | #include 23 | #define __CUDACC__ 24 | #include 25 | #include 26 | 27 | namespace cg = cooperative_groups; 28 | 29 | struct CustomMin 30 | { 31 | __device__ __forceinline__ 32 | float3 operator()(const float3& a, const float3& b) const { 33 | return { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) }; 34 | } 35 | }; 36 | 37 | struct CustomMax 38 | { 39 | __device__ __forceinline__ 40 | float3 operator()(const float3& a, const float3& b) const { 41 | return { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) }; 42 | } 43 | }; 44 | 45 | __host__ __device__ uint32_t prepMorton(uint32_t x) 46 | { 47 | x = (x | (x << 16)) & 0x030000FF; 48 | x = (x | (x << 8)) & 0x0300F00F; 49 | x = (x | (x << 4)) & 0x030C30C3; 50 | x = (x | (x << 2)) & 0x09249249; 51 | return x; 52 | } 53 | 54 | __host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx) 55 | { 56 | uint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1)); 57 | uint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1)); 58 | uint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1)); 59 | 60 | return x | (y << 1) | (z << 2); 61 | } 62 | 63 | __global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes) 64 | { 65 | auto idx = cg::this_grid().thread_rank(); 66 | if (idx >= P) 67 | return; 68 | 69 | codes[idx] = coord2Morton(points[idx], minn, maxx); 70 | } 71 | 72 | struct MinMax 73 | { 74 | float3 minn; 75 | float3 maxx; 76 | }; 77 | 78 | __global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes) 79 | { 80 | auto idx = cg::this_grid().thread_rank(); 81 | 82 | MinMax me; 83 | if (idx < P) 84 | { 85 | me.minn = points[indices[idx]]; 86 | me.maxx = points[indices[idx]]; 87 | } 88 | else 89 | { 90 | me.minn = { FLT_MAX, FLT_MAX, FLT_MAX }; 91 | me.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX }; 92 | } 93 | 94 | __shared__ MinMax redResult[BOX_SIZE]; 95 | 96 | for (int off = BOX_SIZE / 2; off >= 1; off /= 2) 97 | { 98 | if (threadIdx.x < 2 * off) 99 | redResult[threadIdx.x] = me; 100 | __syncthreads(); 101 | 102 | if (threadIdx.x < off) 103 | { 104 | MinMax other = redResult[threadIdx.x + off]; 105 | me.minn.x = min(me.minn.x, other.minn.x); 106 | me.minn.y = min(me.minn.y, other.minn.y); 107 | me.minn.z = min(me.minn.z, other.minn.z); 108 | me.maxx.x = max(me.maxx.x, other.maxx.x); 109 | me.maxx.y = max(me.maxx.y, other.maxx.y); 110 | me.maxx.z = max(me.maxx.z, other.maxx.z); 111 | } 112 | __syncthreads(); 113 | } 114 | 115 | if (threadIdx.x == 0) 116 | boxes[blockIdx.x] = me; 117 | } 118 | 119 | __device__ __host__ float distBoxPoint(const MinMax& box, const float3& p) 120 | { 121 | float3 diff = { 0, 0, 0 }; 122 | if (p.x < box.minn.x || p.x > box.maxx.x) 123 | diff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x)); 124 | if (p.y < box.minn.y || p.y > box.maxx.y) 125 | diff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y)); 126 | if (p.z < box.minn.z || p.z > box.maxx.z) 127 | diff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z)); 128 | return diff.x * diff.x + diff.y * diff.y + diff.z * diff.z; 129 | } 130 | 131 | template 132 | __device__ void updateKBest(const float3& ref, const float3& point, float* knn) 133 | { 134 | float3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z }; 135 | float dist = d.x * d.x + d.y * d.y + d.z * d.z; 136 | for (int j = 0; j < K; j++) 137 | { 138 | if (knn[j] > dist) 139 | { 140 | float t = knn[j]; 141 | knn[j] = dist; 142 | dist = t; 143 | } 144 | } 145 | } 146 | 147 | __global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists) 148 | { 149 | int idx = cg::this_grid().thread_rank(); 150 | if (idx >= P) 151 | return; 152 | 153 | float3 point = points[indices[idx]]; 154 | float best[3] = { FLT_MAX, FLT_MAX, FLT_MAX }; 155 | 156 | for (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++) 157 | { 158 | if (i == idx) 159 | continue; 160 | updateKBest<3>(point, points[indices[i]], best); 161 | } 162 | 163 | float reject = best[2]; 164 | best[0] = FLT_MAX; 165 | best[1] = FLT_MAX; 166 | best[2] = FLT_MAX; 167 | 168 | for (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++) 169 | { 170 | MinMax box = boxes[b]; 171 | float dist = distBoxPoint(box, point); 172 | if (dist > reject || dist > best[2]) 173 | continue; 174 | 175 | for (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++) 176 | { 177 | if (i == idx) 178 | continue; 179 | updateKBest<3>(point, points[indices[i]], best); 180 | } 181 | } 182 | dists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f; 183 | } 184 | 185 | void SimpleKNN::knn(int P, float3* points, float* meanDists) 186 | { 187 | float3* result; 188 | cudaMalloc(&result, sizeof(float3)); 189 | size_t temp_storage_bytes; 190 | 191 | float3 init = { 0, 0, 0 }, minn, maxx; 192 | 193 | cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init); 194 | thrust::device_vector temp_storage(temp_storage_bytes); 195 | 196 | cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init); 197 | cudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost); 198 | 199 | cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init); 200 | cudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost); 201 | 202 | thrust::device_vector morton(P); 203 | thrust::device_vector morton_sorted(P); 204 | coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get()); 205 | 206 | thrust::device_vector indices(P); 207 | thrust::sequence(indices.begin(), indices.end()); 208 | thrust::device_vector indices_sorted(P); 209 | 210 | cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); 211 | temp_storage.resize(temp_storage_bytes); 212 | 213 | cub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); 214 | 215 | uint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE; 216 | thrust::device_vector boxes(num_boxes); 217 | boxMinMax << > > (P, points, indices_sorted.data().get(), boxes.data().get()); 218 | boxMeanDist << > > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists); 219 | 220 | cudaFree(result); 221 | } -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/simple_knn.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #ifndef SIMPLEKNN_H_INCLUDED 13 | #define SIMPLEKNN_H_INCLUDED 14 | 15 | class SimpleKNN 16 | { 17 | public: 18 | static void knn(int P, float3* points, float* meanDists); 19 | }; 20 | 21 | #endif -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/simple_knn/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/submodules/simple-knn/simple_knn/.gitkeep -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/spatial.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #include "spatial.h" 13 | #include "simple_knn.h" 14 | 15 | torch::Tensor 16 | distCUDA2(const torch::Tensor& points) 17 | { 18 | const int P = points.size(0); 19 | 20 | auto float_opts = points.options().dtype(torch::kFloat32); 21 | torch::Tensor means = torch::full({P}, 0.0, float_opts); 22 | 23 | SimpleKNN::knn(P, (float3*)points.contiguous().data(), means.contiguous().data()); 24 | 25 | return means; 26 | } -------------------------------------------------------------------------------- /gaussian_splatting/submodules/simple-knn/spatial.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023, Inria 3 | * GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | * All rights reserved. 5 | * 6 | * This software is free for non-commercial, research and evaluation use 7 | * under the terms of the LICENSE.md file. 8 | * 9 | * For inquiries contact george.drettakis@inria.fr 10 | */ 11 | 12 | #include 13 | 14 | torch::Tensor distCUDA2(const torch::Tensor& points); -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/camera_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/camera_utils.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/camera_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/camera_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/general_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/general_utils.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/general_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/general_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/graphics_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/graphics_utils.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/graphics_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/graphics_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/image_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/image_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/loss_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/loss_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/sh_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/sh_utils.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/sh_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/sh_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/system_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/system_utils.cpython-310.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/__pycache__/system_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/gaussian_splatting/utils/__pycache__/system_utils.cpython-38.pyc -------------------------------------------------------------------------------- /gaussian_splatting/utils/camera_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from scene.cameras import Camera 13 | import numpy as np 14 | from utils.general_utils import PILtoTorch 15 | from utils.graphics_utils import fov2focal 16 | 17 | WARNED = False 18 | 19 | def loadCam(args, id, cam_info, resolution_scale): 20 | orig_w, orig_h = cam_info.image.size 21 | 22 | if args.resolution in [1, 2, 4, 8]: 23 | resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution)) 24 | else: # should be a type that converts to float 25 | if args.resolution == -1: 26 | if orig_w > 1600: 27 | global WARNED 28 | if not WARNED: 29 | print("[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\n " 30 | "If this is not desired, please explicitly specify '--resolution/-r' as 1") 31 | WARNED = True 32 | global_down = orig_w / 1600 33 | else: 34 | global_down = 1 35 | else: 36 | global_down = orig_w / args.resolution 37 | 38 | scale = float(global_down) * float(resolution_scale) 39 | resolution = (int(orig_w / scale), int(orig_h / scale)) 40 | 41 | resized_image_rgb = PILtoTorch(cam_info.image, resolution) 42 | 43 | gt_image = resized_image_rgb[:3, ...] 44 | loaded_mask = None 45 | 46 | if resized_image_rgb.shape[1] == 4: 47 | loaded_mask = resized_image_rgb[3:4, ...] 48 | 49 | return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 50 | FoVx=cam_info.FovX, FoVy=cam_info.FovY, 51 | image=gt_image, gt_alpha_mask=loaded_mask, 52 | image_name=cam_info.image_name, uid=id, data_device=args.data_device) 53 | 54 | def cameraList_from_camInfos(cam_infos, resolution_scale, args): 55 | camera_list = [] 56 | 57 | for id, c in enumerate(cam_infos): 58 | camera_list.append(loadCam(args, id, c, resolution_scale)) 59 | 60 | return camera_list 61 | 62 | def camera_to_JSON(id, camera : Camera): 63 | Rt = np.zeros((4, 4)) 64 | Rt[:3, :3] = camera.R.transpose() 65 | Rt[:3, 3] = camera.T 66 | Rt[3, 3] = 1.0 67 | 68 | W2C = np.linalg.inv(Rt) 69 | pos = W2C[:3, 3] 70 | rot = W2C[:3, :3] 71 | serializable_array_2d = [x.tolist() for x in rot] 72 | camera_entry = { 73 | 'id' : id, 74 | 'img_name' : camera.image_name, 75 | 'width' : camera.width, 76 | 'height' : camera.height, 77 | 'position': pos.tolist(), 78 | 'rotation': serializable_array_2d, 79 | 'fy' : fov2focal(camera.FovY, camera.height), 80 | 'fx' : fov2focal(camera.FovX, camera.width) 81 | } 82 | return camera_entry 83 | -------------------------------------------------------------------------------- /gaussian_splatting/utils/general_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import sys 14 | from datetime import datetime 15 | import numpy as np 16 | import random 17 | 18 | def inverse_sigmoid(x): 19 | return torch.log(x/(1-x)) 20 | 21 | def PILtoTorch(pil_image, resolution): 22 | resized_image_PIL = pil_image.resize(resolution) 23 | resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 24 | if len(resized_image.shape) == 3: 25 | return resized_image.permute(2, 0, 1) 26 | else: 27 | return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) 28 | 29 | def get_expon_lr_func( 30 | lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 31 | ): 32 | """ 33 | Copied from Plenoxels 34 | 35 | Continuous learning rate decay function. Adapted from JaxNeRF 36 | The returned rate is lr_init when step=0 and lr_final when step=max_steps, and 37 | is log-linearly interpolated elsewhere (equivalent to exponential decay). 38 | If lr_delay_steps>0 then the learning rate will be scaled by some smooth 39 | function of lr_delay_mult, such that the initial learning rate is 40 | lr_init*lr_delay_mult at the beginning of optimization but will be eased back 41 | to the normal learning rate when steps>lr_delay_steps. 42 | :param conf: config subtree 'lr' or similar 43 | :param max_steps: int, the number of steps during optimization. 44 | :return HoF which takes step as input 45 | """ 46 | 47 | def helper(step): 48 | if step < 0 or (lr_init == 0.0 and lr_final == 0.0): 49 | # Disable this parameter 50 | return 0.0 51 | if lr_delay_steps > 0: 52 | # A kind of reverse cosine decay. 53 | delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( 54 | 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) 55 | ) 56 | else: 57 | delay_rate = 1.0 58 | t = np.clip(step / max_steps, 0, 1) 59 | log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) 60 | return delay_rate * log_lerp 61 | 62 | return helper 63 | 64 | def strip_lowerdiag(L): 65 | uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda") 66 | 67 | uncertainty[:, 0] = L[:, 0, 0] 68 | uncertainty[:, 1] = L[:, 0, 1] 69 | uncertainty[:, 2] = L[:, 0, 2] 70 | uncertainty[:, 3] = L[:, 1, 1] 71 | uncertainty[:, 4] = L[:, 1, 2] 72 | uncertainty[:, 5] = L[:, 2, 2] 73 | return uncertainty 74 | 75 | def strip_symmetric(sym): 76 | return strip_lowerdiag(sym) 77 | 78 | def build_rotation(r): 79 | norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3]) 80 | 81 | q = r / norm[:, None] 82 | 83 | R = torch.zeros((q.size(0), 3, 3), device='cuda') 84 | 85 | r = q[:, 0] 86 | x = q[:, 1] 87 | y = q[:, 2] 88 | z = q[:, 3] 89 | 90 | R[:, 0, 0] = 1 - 2 * (y*y + z*z) 91 | R[:, 0, 1] = 2 * (x*y - r*z) 92 | R[:, 0, 2] = 2 * (x*z + r*y) 93 | R[:, 1, 0] = 2 * (x*y + r*z) 94 | R[:, 1, 1] = 1 - 2 * (x*x + z*z) 95 | R[:, 1, 2] = 2 * (y*z - r*x) 96 | R[:, 2, 0] = 2 * (x*z - r*y) 97 | R[:, 2, 1] = 2 * (y*z + r*x) 98 | R[:, 2, 2] = 1 - 2 * (x*x + y*y) 99 | return R 100 | 101 | def build_scaling_rotation(s, r): 102 | L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda") 103 | R = build_rotation(r) 104 | 105 | L[:,0,0] = s[:,0] 106 | L[:,1,1] = s[:,1] 107 | L[:,2,2] = s[:,2] 108 | 109 | L = R @ L 110 | return L 111 | 112 | def safe_state(silent): 113 | old_f = sys.stdout 114 | class F: 115 | def __init__(self, silent): 116 | self.silent = silent 117 | 118 | def write(self, x): 119 | if not self.silent: 120 | if x.endswith("\n"): 121 | old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S"))))) 122 | else: 123 | old_f.write(x) 124 | 125 | def flush(self): 126 | old_f.flush() 127 | 128 | sys.stdout = F(silent) 129 | 130 | random.seed(0) 131 | np.random.seed(0) 132 | torch.manual_seed(0) 133 | torch.cuda.set_device(torch.device("cuda:0")) 134 | -------------------------------------------------------------------------------- /gaussian_splatting/utils/graphics_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import math 14 | import numpy as np 15 | from typing import NamedTuple 16 | 17 | class BasicPointCloud(NamedTuple): 18 | points : np.array 19 | colors : np.array 20 | normals : np.array 21 | 22 | def geom_transform_points(points, transf_matrix): 23 | P, _ = points.shape 24 | ones = torch.ones(P, 1, dtype=points.dtype, device=points.device) 25 | points_hom = torch.cat([points, ones], dim=1) 26 | points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0)) 27 | 28 | denom = points_out[..., 3:] + 0.0000001 29 | return (points_out[..., :3] / denom).squeeze(dim=0) 30 | 31 | def getWorld2View(R, t): 32 | Rt = np.zeros((4, 4)) 33 | Rt[:3, :3] = R.transpose() 34 | Rt[:3, 3] = t 35 | Rt[3, 3] = 1.0 36 | return np.float32(Rt) 37 | 38 | def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0): 39 | Rt = np.zeros((4, 4)) 40 | Rt[:3, :3] = R.transpose() 41 | Rt[:3, 3] = t 42 | Rt[3, 3] = 1.0 43 | 44 | C2W = np.linalg.inv(Rt) 45 | cam_center = C2W[:3, 3] 46 | cam_center = (cam_center + translate) * scale 47 | C2W[:3, 3] = cam_center 48 | Rt = np.linalg.inv(C2W) 49 | return np.float32(Rt) 50 | 51 | def getProjectionMatrix(znear, zfar, fovX, fovY): 52 | tanHalfFovY = math.tan((fovY / 2)) 53 | tanHalfFovX = math.tan((fovX / 2)) 54 | 55 | top = tanHalfFovY * znear 56 | bottom = -top 57 | right = tanHalfFovX * znear 58 | left = -right 59 | 60 | P = torch.zeros(4, 4) 61 | 62 | z_sign = 1.0 63 | 64 | P[0, 0] = 2.0 * znear / (right - left) 65 | P[1, 1] = 2.0 * znear / (top - bottom) 66 | P[0, 2] = (right + left) / (right - left) 67 | P[1, 2] = (top + bottom) / (top - bottom) 68 | P[3, 2] = z_sign 69 | P[2, 2] = z_sign * zfar / (zfar - znear) 70 | P[2, 3] = -(zfar * znear) / (zfar - znear) 71 | return P 72 | 73 | def fov2focal(fov, pixels): 74 | return pixels / (2 * math.tan(fov / 2)) 75 | 76 | def focal2fov(focal, pixels): 77 | return 2*math.atan(pixels/(2*focal)) -------------------------------------------------------------------------------- /gaussian_splatting/utils/image_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | 14 | def mse(img1, img2): 15 | return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 16 | 17 | def psnr(img1, img2): 18 | mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 19 | return 20 * torch.log10(1.0 / torch.sqrt(mse)) 20 | -------------------------------------------------------------------------------- /gaussian_splatting/utils/loss_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import torch.nn.functional as F 14 | from torch.autograd import Variable 15 | from math import exp 16 | 17 | def l1_loss(network_output, gt): 18 | return torch.abs((network_output - gt)).mean() 19 | 20 | def l2_loss(network_output, gt): 21 | return ((network_output - gt) ** 2).mean() 22 | 23 | def gaussian(window_size, sigma): 24 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 25 | return gauss / gauss.sum() 26 | 27 | def create_window(window_size, channel): 28 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 29 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 30 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 31 | return window 32 | 33 | def ssim(img1, img2, window_size=11, size_average=True): 34 | channel = img1.size(-3) 35 | window = create_window(window_size, channel) 36 | 37 | if img1.is_cuda: 38 | window = window.cuda(img1.get_device()) 39 | window = window.type_as(img1) 40 | 41 | return _ssim(img1, img2, window, window_size, channel, size_average) 42 | 43 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 44 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 45 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 46 | 47 | mu1_sq = mu1.pow(2) 48 | mu2_sq = mu2.pow(2) 49 | mu1_mu2 = mu1 * mu2 50 | 51 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 52 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 53 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 54 | 55 | C1 = 0.01 ** 2 56 | C2 = 0.03 ** 2 57 | 58 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 59 | 60 | if size_average: 61 | return ssim_map.mean() 62 | else: 63 | return ssim_map.mean(1).mean(1).mean(1) 64 | 65 | -------------------------------------------------------------------------------- /gaussian_splatting/utils/sh_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The PlenOctree Authors. 2 | # Redistribution and use in source and binary forms, with or without 3 | # modification, are permitted provided that the following conditions are met: 4 | # 5 | # 1. Redistributions of source code must retain the above copyright notice, 6 | # this list of conditions and the following disclaimer. 7 | # 8 | # 2. Redistributions in binary form must reproduce the above copyright notice, 9 | # this list of conditions and the following disclaimer in the documentation 10 | # and/or other materials provided with the distribution. 11 | # 12 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 13 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 14 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 15 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 16 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 17 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 18 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 19 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 20 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 21 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 22 | # POSSIBILITY OF SUCH DAMAGE. 23 | 24 | import torch 25 | 26 | C0 = 0.28209479177387814 27 | C1 = 0.4886025119029199 28 | C2 = [ 29 | 1.0925484305920792, 30 | -1.0925484305920792, 31 | 0.31539156525252005, 32 | -1.0925484305920792, 33 | 0.5462742152960396 34 | ] 35 | C3 = [ 36 | -0.5900435899266435, 37 | 2.890611442640554, 38 | -0.4570457994644658, 39 | 0.3731763325901154, 40 | -0.4570457994644658, 41 | 1.445305721320277, 42 | -0.5900435899266435 43 | ] 44 | C4 = [ 45 | 2.5033429417967046, 46 | -1.7701307697799304, 47 | 0.9461746957575601, 48 | -0.6690465435572892, 49 | 0.10578554691520431, 50 | -0.6690465435572892, 51 | 0.47308734787878004, 52 | -1.7701307697799304, 53 | 0.6258357354491761, 54 | ] 55 | 56 | 57 | def eval_sh(deg, sh, dirs): 58 | """ 59 | Evaluate spherical harmonics at unit directions 60 | using hardcoded SH polynomials. 61 | Works with torch/np/jnp. 62 | ... Can be 0 or more batch dimensions. 63 | Args: 64 | deg: int SH deg. Currently, 0-3 supported 65 | sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] 66 | dirs: jnp.ndarray unit directions [..., 3] 67 | Returns: 68 | [..., C] 69 | """ 70 | assert deg <= 4 and deg >= 0 71 | coeff = (deg + 1) ** 2 72 | assert sh.shape[-1] >= coeff 73 | 74 | result = C0 * sh[..., 0] 75 | if deg > 0: 76 | x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] 77 | result = (result - 78 | C1 * y * sh[..., 1] + 79 | C1 * z * sh[..., 2] - 80 | C1 * x * sh[..., 3]) 81 | 82 | if deg > 1: 83 | xx, yy, zz = x * x, y * y, z * z 84 | xy, yz, xz = x * y, y * z, x * z 85 | result = (result + 86 | C2[0] * xy * sh[..., 4] + 87 | C2[1] * yz * sh[..., 5] + 88 | C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + 89 | C2[3] * xz * sh[..., 7] + 90 | C2[4] * (xx - yy) * sh[..., 8]) 91 | 92 | if deg > 2: 93 | result = (result + 94 | C3[0] * y * (3 * xx - yy) * sh[..., 9] + 95 | C3[1] * xy * z * sh[..., 10] + 96 | C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + 97 | C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + 98 | C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + 99 | C3[5] * z * (xx - yy) * sh[..., 14] + 100 | C3[6] * x * (xx - 3 * yy) * sh[..., 15]) 101 | 102 | if deg > 3: 103 | result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + 104 | C4[1] * yz * (3 * xx - yy) * sh[..., 17] + 105 | C4[2] * xy * (7 * zz - 1) * sh[..., 18] + 106 | C4[3] * yz * (7 * zz - 3) * sh[..., 19] + 107 | C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + 108 | C4[5] * xz * (7 * zz - 3) * sh[..., 21] + 109 | C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + 110 | C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + 111 | C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) 112 | return result 113 | 114 | def RGB2SH(rgb): 115 | return (rgb - 0.5) / C0 116 | 117 | def SH2RGB(sh): 118 | return sh * C0 + 0.5 -------------------------------------------------------------------------------- /gaussian_splatting/utils/system_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from errno import EEXIST 13 | from os import makedirs, path 14 | import os 15 | 16 | def mkdir_p(folder_path): 17 | # Creates a directory. equivalent to using mkdir -p on the command line 18 | try: 19 | makedirs(folder_path) 20 | except OSError as exc: # Python >2.5 21 | if exc.errno == EEXIST and path.isdir(folder_path): 22 | pass 23 | else: 24 | raise 25 | 26 | def searchForMaxIteration(folder): 27 | saved_iters = [int(fname.split("_")[-1]) for fname in os.listdir(folder)] 28 | return max(saved_iters) 29 | -------------------------------------------------------------------------------- /lpipsPyTorch/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .modules.lpips import LPIPS 4 | 5 | 6 | def lpips(x: torch.Tensor, 7 | y: torch.Tensor, 8 | net_type: str = 'alex', 9 | version: str = '0.1'): 10 | r"""Function that measures 11 | Learned Perceptual Image Patch Similarity (LPIPS). 12 | 13 | Arguments: 14 | x, y (torch.Tensor): the input tensors to compare. 15 | net_type (str): the network type to compare the features: 16 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 17 | version (str): the version of LPIPS. Default: 0.1. 18 | """ 19 | device = x.device 20 | criterion = LPIPS(net_type, version).to(device) 21 | return criterion(x, y) 22 | -------------------------------------------------------------------------------- /lpipsPyTorch/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/lpipsPyTorch/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /lpipsPyTorch/modules/__pycache__/lpips.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/lpipsPyTorch/modules/__pycache__/lpips.cpython-38.pyc -------------------------------------------------------------------------------- /lpipsPyTorch/modules/__pycache__/networks.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/lpipsPyTorch/modules/__pycache__/networks.cpython-38.pyc -------------------------------------------------------------------------------- /lpipsPyTorch/modules/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/lpipsPyTorch/modules/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /lpipsPyTorch/modules/lpips.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .networks import get_network, LinLayers 5 | from .utils import get_state_dict 6 | 7 | 8 | class LPIPS(nn.Module): 9 | r"""Creates a criterion that measures 10 | Learned Perceptual Image Patch Similarity (LPIPS). 11 | 12 | Arguments: 13 | net_type (str): the network type to compare the features: 14 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 15 | version (str): the version of LPIPS. Default: 0.1. 16 | """ 17 | def __init__(self, net_type: str = 'alex', version: str = '0.1'): 18 | 19 | assert version in ['0.1'], 'v0.1 is only supported now' 20 | 21 | super(LPIPS, self).__init__() 22 | 23 | # pretrained network 24 | self.net = get_network(net_type) 25 | 26 | # linear layers 27 | self.lin = LinLayers(self.net.n_channels_list) 28 | self.lin.load_state_dict(get_state_dict(net_type, version)) 29 | 30 | def forward(self, x: torch.Tensor, y: torch.Tensor): 31 | feat_x, feat_y = self.net(x), self.net(y) 32 | 33 | diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] 34 | res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] 35 | 36 | return torch.sum(torch.cat(res, 0), 0, True) 37 | -------------------------------------------------------------------------------- /lpipsPyTorch/modules/networks.py: -------------------------------------------------------------------------------- 1 | from typing import Sequence 2 | 3 | from itertools import chain 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torchvision import models 8 | 9 | from .utils import normalize_activation 10 | 11 | 12 | def get_network(net_type: str): 13 | if net_type == 'alex': 14 | return AlexNet() 15 | elif net_type == 'squeeze': 16 | return SqueezeNet() 17 | elif net_type == 'vgg': 18 | return VGG16() 19 | else: 20 | raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') 21 | 22 | 23 | class LinLayers(nn.ModuleList): 24 | def __init__(self, n_channels_list: Sequence[int]): 25 | super(LinLayers, self).__init__([ 26 | nn.Sequential( 27 | nn.Identity(), 28 | nn.Conv2d(nc, 1, 1, 1, 0, bias=False) 29 | ) for nc in n_channels_list 30 | ]) 31 | 32 | for param in self.parameters(): 33 | param.requires_grad = False 34 | 35 | 36 | class BaseNet(nn.Module): 37 | def __init__(self): 38 | super(BaseNet, self).__init__() 39 | 40 | # register buffer 41 | self.register_buffer( 42 | 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) 43 | self.register_buffer( 44 | 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) 45 | 46 | def set_requires_grad(self, state: bool): 47 | for param in chain(self.parameters(), self.buffers()): 48 | param.requires_grad = state 49 | 50 | def z_score(self, x: torch.Tensor): 51 | return (x - self.mean) / self.std 52 | 53 | def forward(self, x: torch.Tensor): 54 | x = self.z_score(x) 55 | 56 | output = [] 57 | for i, (_, layer) in enumerate(self.layers._modules.items(), 1): 58 | x = layer(x) 59 | if i in self.target_layers: 60 | output.append(normalize_activation(x)) 61 | if len(output) == len(self.target_layers): 62 | break 63 | return output 64 | 65 | 66 | class SqueezeNet(BaseNet): 67 | def __init__(self): 68 | super(SqueezeNet, self).__init__() 69 | 70 | self.layers = models.squeezenet1_1(True).features 71 | self.target_layers = [2, 5, 8, 10, 11, 12, 13] 72 | self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] 73 | 74 | self.set_requires_grad(False) 75 | 76 | 77 | class AlexNet(BaseNet): 78 | def __init__(self): 79 | super(AlexNet, self).__init__() 80 | 81 | self.layers = models.alexnet(True).features 82 | self.target_layers = [2, 5, 8, 10, 12] 83 | self.n_channels_list = [64, 192, 384, 256, 256] 84 | 85 | self.set_requires_grad(False) 86 | 87 | 88 | class VGG16(BaseNet): 89 | def __init__(self): 90 | super(VGG16, self).__init__() 91 | 92 | # self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features 93 | self.layers = model = models.vgg16(pretrained=True).features # TODO: old version torchvision 94 | self.target_layers = [4, 9, 16, 23, 30] 95 | self.n_channels_list = [64, 128, 256, 512, 512] 96 | 97 | self.set_requires_grad(False) 98 | -------------------------------------------------------------------------------- /lpipsPyTorch/modules/utils.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | import torch 4 | 5 | 6 | def normalize_activation(x, eps=1e-10): 7 | norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) 8 | return x / (norm_factor + eps) 9 | 10 | 11 | def get_state_dict(net_type: str = 'alex', version: str = '0.1'): 12 | # build url 13 | url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ 14 | + f'master/lpips/weights/v{version}/{net_type}.pth' 15 | 16 | # download 17 | old_state_dict = torch.hub.load_state_dict_from_url( 18 | url, progress=True, 19 | map_location=None if torch.cuda.is_available() else torch.device('cpu') 20 | ) 21 | 22 | # rename keys 23 | new_state_dict = OrderedDict() 24 | for key, val in old_state_dict.items(): 25 | new_key = key 26 | new_key = new_key.replace('lin', '') 27 | new_key = new_key.replace('model.', '') 28 | new_state_dict[new_key] = val 29 | 30 | return new_state_dict 31 | -------------------------------------------------------------------------------- /media/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/media/overview.png -------------------------------------------------------------------------------- /metrics.py: -------------------------------------------------------------------------------- 1 | 2 | from pathlib import Path 3 | import os 4 | from PIL import Image 5 | import torch 6 | import torchvision.transforms.functional as tf 7 | from lpipsPyTorch import lpips 8 | import json 9 | from tqdm import tqdm 10 | from argparse import ArgumentParser 11 | import cv2 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | from math import exp 15 | 16 | 17 | def mse(img1, img2): 18 | return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 19 | 20 | def psnr(img1, img2): 21 | mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 22 | return 20 * torch.log10(1.0 / torch.sqrt(mse)) 23 | 24 | def gaussian(window_size, sigma): 25 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 26 | return gauss / gauss.sum() 27 | 28 | def create_window(window_size, channel): 29 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 30 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 31 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 32 | return window 33 | 34 | def ssim(img1, img2, window_size=11, size_average=True): 35 | channel = img1.size(-3) 36 | window = create_window(window_size, channel) 37 | 38 | if img1.is_cuda: 39 | window = window.cuda(img1.get_device()) 40 | window = window.type_as(img1) 41 | 42 | return _ssim(img1, img2, window, window_size, channel, size_average) 43 | 44 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 45 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 46 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 47 | 48 | mu1_sq = mu1.pow(2) 49 | mu2_sq = mu2.pow(2) 50 | mu1_mu2 = mu1 * mu2 51 | 52 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 53 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 54 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 55 | 56 | C1 = 0.01 ** 2 57 | C2 = 0.03 ** 2 58 | 59 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 60 | 61 | if size_average: 62 | return ssim_map.mean() 63 | else: 64 | return ssim_map.mean(1).mean(1).mean(1) 65 | 66 | 67 | 68 | def readImages(renders_dir, gt_dir): 69 | renders = [] 70 | gts = [] 71 | image_names = [] 72 | for fname in os.listdir(renders_dir): 73 | render = Image.open(renders_dir / fname) 74 | gt = Image.open(gt_dir / fname) 75 | renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda()) 76 | gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda()) 77 | image_names.append(fname) 78 | return renders, gts, image_names 79 | 80 | def evaluate(model_paths): 81 | 82 | full_dict = {} 83 | per_view_dict = {} 84 | full_dict_polytopeonly = {} 85 | per_view_dict_polytopeonly = {} 86 | print("") 87 | 88 | scene_dir = model_paths 89 | print("Scene:", scene_dir) 90 | full_dict[scene_dir] = {} 91 | per_view_dict[scene_dir] = {} 92 | full_dict_polytopeonly[scene_dir] = {} 93 | per_view_dict_polytopeonly[scene_dir] = {} 94 | 95 | test_dir = Path(scene_dir) / "test" 96 | 97 | method_dir = test_dir 98 | gt_dir = method_dir/ "gt" 99 | renders_dir = method_dir / "renders" 100 | renders, gts, image_names = readImages(renders_dir, gt_dir) 101 | 102 | ssims = [] 103 | psnrs = [] 104 | lpipss = [] 105 | 106 | for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"): 107 | s = ssim(renders[idx], gts[idx]) 108 | ssims.append(s) 109 | p = psnr(renders[idx], gts[idx]) 110 | # error = torch.abs(renders[idx] - gts[idx]).sum(0).sum(0).cpu().numpy() 111 | # error = error/(error.max()-error.min()) * 255 112 | psnrs.append(p) 113 | l = lpips(renders[idx], gts[idx], net_type='vgg') 114 | lpipss.append(l) 115 | print(s.item(), p.item(), l.item()) 116 | 117 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 118 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 119 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 120 | print("") 121 | 122 | full_dict[scene_dir].update({"SSIM": torch.tensor(ssims).mean().item(), 123 | "PSNR": torch.tensor(psnrs).mean().item(), 124 | "LPIPS": torch.tensor(lpipss).mean().item()}) 125 | per_view_dict[scene_dir].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 126 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 127 | "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}}) 128 | 129 | with open(scene_dir + "/results.json", 'w') as fp: 130 | json.dump(full_dict[scene_dir], fp, indent=True) 131 | with open(scene_dir + "/per_view.json", 'w') as fp: 132 | json.dump(per_view_dict[scene_dir], fp, indent=True) 133 | 134 | if __name__ == "__main__": 135 | device = torch.device("cuda:0") 136 | torch.cuda.set_device(device) 137 | 138 | # Set up command line argument parser 139 | parser = ArgumentParser(description="Training script parameters") 140 | parser.add_argument('--model_paths', '-m', required=True, type=str) 141 | args = parser.parse_args() 142 | evaluate(args.model_paths) 143 | -------------------------------------------------------------------------------- /np_utils/extensions/chamfer_dist/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Thibault GROUEIX 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-18 15:06:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import chamfer 11 | 12 | 13 | class ChamferFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, xyz1, xyz2): 16 | dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) 17 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 18 | 19 | return dist1, dist2 20 | 21 | @staticmethod 22 | def backward(ctx, grad_dist1, grad_dist2): 23 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 24 | grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) 25 | return grad_xyz1, grad_xyz2 26 | 27 | 28 | class ChamferDistanceL2(torch.nn.Module): 29 | f''' Chamder Distance L2 30 | ''' 31 | def __init__(self, ignore_zeros=False): 32 | super().__init__() 33 | self.ignore_zeros = ignore_zeros 34 | 35 | def forward(self, xyz1, xyz2): 36 | batch_size = xyz1.size(0) 37 | if batch_size == 1 and self.ignore_zeros: 38 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 39 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 40 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 41 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 42 | 43 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 44 | return torch.mean(dist1) + torch.mean(dist2) 45 | 46 | class ChamferDistanceL2_split(torch.nn.Module): 47 | f''' Chamder Distance L2 48 | ''' 49 | def __init__(self, ignore_zeros=False): 50 | super().__init__() 51 | self.ignore_zeros = ignore_zeros 52 | 53 | def forward(self, xyz1, xyz2): 54 | batch_size = xyz1.size(0) 55 | if batch_size == 1 and self.ignore_zeros: 56 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 57 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 58 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 59 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 60 | 61 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 62 | return torch.mean(dist1), torch.mean(dist2) 63 | 64 | class ChamferDistanceL1(torch.nn.Module): 65 | f''' Chamder Distance L1 66 | ''' 67 | def __init__(self, ignore_zeros=False): 68 | super().__init__() 69 | self.ignore_zeros = ignore_zeros 70 | 71 | def forward(self, xyz1, xyz2): 72 | batch_size = xyz1.size(0) 73 | if batch_size == 1 and self.ignore_zeros: 74 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 75 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 76 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 77 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 78 | 79 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 80 | # import pdb 81 | # pdb.set_trace() 82 | dist1 = torch.sqrt(dist1) 83 | dist2 = torch.sqrt(dist2) 84 | return (torch.mean(dist1) + torch.mean(dist2))/2 85 | 86 | -------------------------------------------------------------------------------- /np_utils/extensions/chamfer_dist/chamfer.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-08-07 20:54:24 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:58:55 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | __global__ void chamfer_dist_kernel(int batch_size, 16 | int n, 17 | const float* xyz1, 18 | int m, 19 | const float* xyz2, 20 | float* dist, 21 | int* indexes) { 22 | const int batch = 512; 23 | __shared__ float buf[batch * 3]; 24 | for (int i = blockIdx.x; i < batch_size; i += gridDim.x) { 25 | for (int k2 = 0; k2 < m; k2 += batch) { 26 | int end_k = min(m, k2 + batch) - k2; 27 | for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) { 28 | buf[j] = xyz2[(i * m + k2) * 3 + j]; 29 | } 30 | __syncthreads(); 31 | for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; 32 | j += blockDim.x * gridDim.y) { 33 | float x1 = xyz1[(i * n + j) * 3 + 0]; 34 | float y1 = xyz1[(i * n + j) * 3 + 1]; 35 | float z1 = xyz1[(i * n + j) * 3 + 2]; 36 | float best_dist = 0; 37 | int best_dist_index = 0; 38 | int end_ka = end_k - (end_k & 3); 39 | if (end_ka == batch) { 40 | for (int k = 0; k < batch; k += 4) { 41 | { 42 | float x2 = buf[k * 3 + 0] - x1; 43 | float y2 = buf[k * 3 + 1] - y1; 44 | float z2 = buf[k * 3 + 2] - z1; 45 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 46 | 47 | if (k == 0 || dist < best_dist) { 48 | best_dist = dist; 49 | best_dist_index = k + k2; 50 | } 51 | } 52 | { 53 | float x2 = buf[k * 3 + 3] - x1; 54 | float y2 = buf[k * 3 + 4] - y1; 55 | float z2 = buf[k * 3 + 5] - z1; 56 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 57 | if (dist < best_dist) { 58 | best_dist = dist; 59 | best_dist_index = k + k2 + 1; 60 | } 61 | } 62 | { 63 | float x2 = buf[k * 3 + 6] - x1; 64 | float y2 = buf[k * 3 + 7] - y1; 65 | float z2 = buf[k * 3 + 8] - z1; 66 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 67 | if (dist < best_dist) { 68 | best_dist = dist; 69 | best_dist_index = k + k2 + 2; 70 | } 71 | } 72 | { 73 | float x2 = buf[k * 3 + 9] - x1; 74 | float y2 = buf[k * 3 + 10] - y1; 75 | float z2 = buf[k * 3 + 11] - z1; 76 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 77 | if (dist < best_dist) { 78 | best_dist = dist; 79 | best_dist_index = k + k2 + 3; 80 | } 81 | } 82 | } 83 | } else { 84 | for (int k = 0; k < end_ka; k += 4) { 85 | { 86 | float x2 = buf[k * 3 + 0] - x1; 87 | float y2 = buf[k * 3 + 1] - y1; 88 | float z2 = buf[k * 3 + 2] - z1; 89 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 90 | if (k == 0 || dist < best_dist) { 91 | best_dist = dist; 92 | best_dist_index = k + k2; 93 | } 94 | } 95 | { 96 | float x2 = buf[k * 3 + 3] - x1; 97 | float y2 = buf[k * 3 + 4] - y1; 98 | float z2 = buf[k * 3 + 5] - z1; 99 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 100 | if (dist < best_dist) { 101 | best_dist = dist; 102 | best_dist_index = k + k2 + 1; 103 | } 104 | } 105 | { 106 | float x2 = buf[k * 3 + 6] - x1; 107 | float y2 = buf[k * 3 + 7] - y1; 108 | float z2 = buf[k * 3 + 8] - z1; 109 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 110 | if (dist < best_dist) { 111 | best_dist = dist; 112 | best_dist_index = k + k2 + 2; 113 | } 114 | } 115 | { 116 | float x2 = buf[k * 3 + 9] - x1; 117 | float y2 = buf[k * 3 + 10] - y1; 118 | float z2 = buf[k * 3 + 11] - z1; 119 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 120 | if (dist < best_dist) { 121 | best_dist = dist; 122 | best_dist_index = k + k2 + 3; 123 | } 124 | } 125 | } 126 | } 127 | for (int k = end_ka; k < end_k; k++) { 128 | float x2 = buf[k * 3 + 0] - x1; 129 | float y2 = buf[k * 3 + 1] - y1; 130 | float z2 = buf[k * 3 + 2] - z1; 131 | float dist = x2 * x2 + y2 * y2 + z2 * z2; 132 | if (k == 0 || dist < best_dist) { 133 | best_dist = dist; 134 | best_dist_index = k + k2; 135 | } 136 | } 137 | if (k2 == 0 || dist[(i * n + j)] > best_dist) { 138 | dist[(i * n + j)] = best_dist; 139 | indexes[(i * n + j)] = best_dist_index; 140 | } 141 | } 142 | __syncthreads(); 143 | } 144 | } 145 | } 146 | 147 | std::vector chamfer_cuda_forward(torch::Tensor xyz1, 148 | torch::Tensor xyz2) { 149 | const int batch_size = xyz1.size(0); 150 | const int n = xyz1.size(1); // num_points point cloud A 151 | const int m = xyz2.size(1); // num_points point cloud B 152 | torch::Tensor dist1 = 153 | torch::zeros({batch_size, n}, torch::CUDA(torch::kFloat)); 154 | torch::Tensor dist2 = 155 | torch::zeros({batch_size, m}, torch::CUDA(torch::kFloat)); 156 | torch::Tensor idx1 = torch::zeros({batch_size, n}, torch::CUDA(torch::kInt)); 157 | torch::Tensor idx2 = torch::zeros({batch_size, m}, torch::CUDA(torch::kInt)); 158 | 159 | chamfer_dist_kernel<<>>( 160 | batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), 161 | dist1.data_ptr(), idx1.data_ptr()); 162 | chamfer_dist_kernel<<>>( 163 | batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), 164 | dist2.data_ptr(), idx2.data_ptr()); 165 | 166 | cudaError_t err = cudaGetLastError(); 167 | if (err != cudaSuccess) { 168 | printf("Error in chamfer_cuda_forward: %s\n", cudaGetErrorString(err)); 169 | } 170 | return {dist1, dist2, idx1, idx2}; 171 | } 172 | 173 | __global__ void chamfer_dist_grad_kernel(int b, 174 | int n, 175 | const float* xyz1, 176 | int m, 177 | const float* xyz2, 178 | const float* grad_dist1, 179 | const int* idx1, 180 | float* grad_xyz1, 181 | float* grad_xyz2) { 182 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 183 | for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; 184 | j += blockDim.x * gridDim.y) { 185 | float x1 = xyz1[(i * n + j) * 3 + 0]; 186 | float y1 = xyz1[(i * n + j) * 3 + 1]; 187 | float z1 = xyz1[(i * n + j) * 3 + 2]; 188 | int j2 = idx1[i * n + j]; 189 | float x2 = xyz2[(i * m + j2) * 3 + 0]; 190 | float y2 = xyz2[(i * m + j2) * 3 + 1]; 191 | float z2 = xyz2[(i * m + j2) * 3 + 2]; 192 | float g = grad_dist1[i * n + j] * 2; 193 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 0]), g * (x1 - x2)); 194 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 1]), g * (y1 - y2)); 195 | atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 2]), g * (z1 - z2)); 196 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 0]), -(g * (x1 - x2))); 197 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 1]), -(g * (y1 - y2))); 198 | atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 2]), -(g * (z1 - z2))); 199 | } 200 | } 201 | } 202 | 203 | std::vector chamfer_cuda_backward(torch::Tensor xyz1, 204 | torch::Tensor xyz2, 205 | torch::Tensor idx1, 206 | torch::Tensor idx2, 207 | torch::Tensor grad_dist1, 208 | torch::Tensor grad_dist2) { 209 | const int batch_size = xyz1.size(0); 210 | const int n = xyz1.size(1); // num_points point cloud A 211 | const int m = xyz2.size(1); // num_points point cloud B 212 | torch::Tensor grad_xyz1 = torch::zeros_like(xyz1, torch::CUDA(torch::kFloat)); 213 | torch::Tensor grad_xyz2 = torch::zeros_like(xyz2, torch::CUDA(torch::kFloat)); 214 | 215 | chamfer_dist_grad_kernel<<>>( 216 | batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), 217 | grad_dist1.data_ptr(), idx1.data_ptr(), 218 | grad_xyz1.data_ptr(), grad_xyz2.data_ptr()); 219 | chamfer_dist_grad_kernel<<>>( 220 | batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), 221 | grad_dist2.data_ptr(), idx2.data_ptr(), 222 | grad_xyz2.data_ptr(), grad_xyz1.data_ptr()); 223 | 224 | cudaError_t err = cudaGetLastError(); 225 | if (err != cudaSuccess) { 226 | printf("Error in chamfer_cuda_backward: %s\n", cudaGetErrorString(err)); 227 | } 228 | return {grad_xyz1, grad_xyz2}; 229 | } 230 | -------------------------------------------------------------------------------- /np_utils/extensions/chamfer_dist/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-08-07 20:54:24 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2019-12-10 10:33:50 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | std::vector chamfer_cuda_forward(torch::Tensor xyz1, 13 | torch::Tensor xyz2); 14 | 15 | std::vector chamfer_cuda_backward(torch::Tensor xyz1, 16 | torch::Tensor xyz2, 17 | torch::Tensor idx1, 18 | torch::Tensor idx2, 19 | torch::Tensor grad_dist1, 20 | torch::Tensor grad_dist2); 21 | 22 | std::vector chamfer_forward(torch::Tensor xyz1, 23 | torch::Tensor xyz2) { 24 | return chamfer_cuda_forward(xyz1, xyz2); 25 | } 26 | 27 | std::vector chamfer_backward(torch::Tensor xyz1, 28 | torch::Tensor xyz2, 29 | torch::Tensor idx1, 30 | torch::Tensor idx2, 31 | torch::Tensor grad_dist1, 32 | torch::Tensor grad_dist2) { 33 | return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); 34 | } 35 | 36 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 37 | m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); 38 | m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); 39 | } 40 | -------------------------------------------------------------------------------- /np_utils/extensions/chamfer_dist/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-10 10:04:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='chamfer', 12 | version='2.0.0', 13 | ext_modules=[ 14 | CUDAExtension('chamfer', [ 15 | 'chamfer_cuda.cpp', 16 | 'chamfer.cu', 17 | ]), 18 | ], 19 | cmdclass={'build_ext': BuildExtension}) 20 | -------------------------------------------------------------------------------- /np_utils/extensions/chamfer_dist/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:38:01 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:21:36 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in chamfer.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | 17 | from torch.autograd import gradcheck 18 | 19 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 20 | from extensions.chamfer_dist import ChamferFunction 21 | 22 | 23 | class ChamferDistanceTestCase(unittest.TestCase): 24 | def test_chamfer_dist(self): 25 | x = torch.rand(4, 64, 3).double() 26 | y = torch.rand(4, 128, 3).double() 27 | x.requires_grad = True 28 | y.requires_grad = True 29 | print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) 30 | 31 | 32 | 33 | if __name__ == '__main__': 34 | # unittest.main() 35 | import pdb 36 | x = torch.rand(32,128,3) 37 | y = torch.rand(32,128,3) 38 | pdb.set_trace() 39 | -------------------------------------------------------------------------------- /np_utils/owndata.conf: -------------------------------------------------------------------------------- 1 | general { 2 | base_exp_dir = ./outs/, 3 | recording = [ 4 | ./, 5 | ./models, 6 | ] 7 | } 8 | 9 | dataset { 10 | data_dir = data/owndata/, 11 | } 12 | 13 | train { 14 | learning_rate = 0.001, 15 | step1_maxiter = 50000, 16 | step2_maxiter = 100000, 17 | warm_up_end = 1000, 18 | eval_num_points = 1000000, 19 | df_filter = 0.01, 20 | far = -1, 21 | outlier = 0.002, 22 | extra_points_rate = 1, 23 | low_range = 1.1, 24 | 25 | batch_size = 5000, 26 | batch_size_step2 = 20000, 27 | 28 | save_freq = 5000, 29 | val_freq = 2500, 30 | val_mesh_freq = 2500, 31 | report_freq = 5000, 32 | 33 | igr_weight = 0.1, 34 | mask_weight = 0.0, 35 | load_ckpt = none, 36 | } 37 | 38 | model { 39 | udf_network { 40 | d_out = 1, 41 | d_in = 3, 42 | d_hidden = 256, 43 | n_layers = 8, 44 | skip_in = [4], 45 | multires = 0, 46 | bias = 0.3, 47 | scale = 1.0, 48 | geometric_init = True, 49 | weight_norm = True, 50 | } 51 | 52 | } 53 | -------------------------------------------------------------------------------- /np_utils/train.py: -------------------------------------------------------------------------------- 1 | 2 | import time 3 | import torch 4 | import torch.nn.functional as F 5 | from tqdm import tqdm 6 | from np_utils.dataset import Dataset 7 | from np_utils.utils import CAPUDFNetwork 8 | import argparse 9 | import os 10 | from shutil import copyfile 11 | import numpy as np 12 | import trimesh 13 | from np_utils.extensions.chamfer_dist import ChamferDistanceL1, ChamferDistanceL2 14 | import math 15 | from pytorch3d.ops import knn_points 16 | from scipy.spatial import cKDTree 17 | import mcubes 18 | import warnings 19 | 20 | warnings.filterwarnings('ignore') 21 | 22 | 23 | class Runner: 24 | def __init__(self, path, pointcloud=None, part_num=4, iteration=9000): 25 | self.device = torch.device('cuda') 26 | self.part_num = part_num 27 | if pointcloud is not None: 28 | for i in range(1, self.part_num + 1): 29 | os.makedirs(os.path.join(path, 'query_data'), exist_ok=True) 30 | dataset_path = os.path.join(path, 'query_data', 'dataset_iter{}_part{}.pt'.format(iteration, i)) 31 | if os.path.exists(dataset_path): 32 | dataset = torch.load(dataset_path) 33 | else: 34 | raise NotImplementedError 35 | dataset = Dataset_multipart(pointcloud, part=i) 36 | torch.save(dataset, dataset_path) 37 | self.__setattr__('dataset' + str(i), dataset) 38 | self.ChamferDisL1 = ChamferDistanceL1().cuda() 39 | 40 | # Networks 41 | for i in range(1, part_num+1): 42 | self.__setattr__('sdf_network'+str(i), CAPUDFNetwork().to(self.device)) 43 | 44 | 45 | def reset_datasets(self, path, pointcloud, iteration=9000, scene_name='Barn'): 46 | for i in range(1, self.part_num + 1): 47 | os.makedirs(os.path.join(path, 'query_data'), exist_ok=True) 48 | dataset_path = os.path.join(path, 'query_data', 'dataset_iter{}_part{}.pt'.format(iteration, i)) 49 | if os.path.exists(dataset_path): 50 | dataset = torch.load(dataset_path) 51 | else: 52 | if not isinstance(pointcloud, np.ndarray): 53 | pointcloud = pointcloud.clone().detach().cpu().numpy() 54 | if hasattr(self, 'dataset' + str(i)): 55 | old_datset = self.__getattribute__('dataset' + str(i)) 56 | pc_bbx = old_datset.pc_bbx 57 | shape_center = old_datset.shape_center 58 | shape_scale = old_datset.shape_scale 59 | else: 60 | pc_bbx = shape_center = shape_scale = None 61 | dataset = Dataset(pointcloud, part=i, scene_name=scene_name, old_pc_bbx=pc_bbx, 62 | old_shape_center=shape_center, old_shape_scale=shape_scale) 63 | torch.save(dataset, dataset_path) 64 | self.__setattr__('dataset' + str(i), dataset) 65 | 66 | 67 | def get_learning_rate_at_iteration(self, iter_step, max_iter=60050): 68 | warn_up = 1000 69 | init_lr = 0.001 70 | lr = (iter_step / warn_up) if iter_step < warn_up else 0.5 * ( 71 | math.cos((iter_step - warn_up) / (max_iter - warn_up) * math.pi) + 1) 72 | lr = lr * init_lr 73 | return lr 74 | 75 | def load_checkpoint(self, path, checkpoint_name): 76 | checkpoint = torch.load(os.path.join(path, 'checkpoints', checkpoint_name), 77 | map_location=self.device) 78 | print(os.path.join(path, 'checkpoints', checkpoint_name)) 79 | for i in range(1, self.part_num+1): 80 | sdf_network = self.__getattribute__('sdf_network'+str(i)) 81 | sdf_network.load_state_dict(checkpoint['sdf_network'+str(i)]) 82 | 83 | def save_checkpoint(self, path, iter_step): 84 | checkpoint = {} 85 | for i in range(1, self.part_num+1): 86 | sdf_network = self.__getattribute__('sdf_network'+str(i)) 87 | checkpoint.update({'sdf_network'+str(i): sdf_network.state_dict()}) 88 | 89 | os.makedirs(os.path.join(path, 'checkpoints'), exist_ok=True) 90 | torch.save(checkpoint, os.path.join(path, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(iter_step))) 91 | 92 | 93 | if __name__ == '__main__': 94 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 95 | parser = argparse.ArgumentParser() 96 | parser.add_argument('--conf', type=str, default='./confs/ndf.conf') 97 | parser.add_argument('--mcube_resolution', type=int, default=256) 98 | parser.add_argument('--gpu', type=int, default=0) 99 | parser.add_argument('--dir', type=str, default='test') 100 | parser.add_argument('--dataname', type=str, default='demo') 101 | args = parser.parse_args() 102 | 103 | torch.cuda.set_device(args.gpu) 104 | runner = Runner(args, args.conf) 105 | # runner.load_checkpoint('ckpt_100000.pth') 106 | # runner.iter_step = 100000-1 107 | 108 | runner.train() 109 | # runner.extract_mesh_meshudf(resolution=512, dist_threshold_ratio=1.0) 110 | # runner.extract_mesh(resolution=512) 111 | for threshold in [0.001, 0.002, 0.003, 0.004, 0.005]: 112 | runner.extract_mesh_sdf(resolution=512, threshold=threshold) -------------------------------------------------------------------------------- /np_utils/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | 6 | class CAPUDFNetwork(nn.Module): 7 | def __init__(self, 8 | d_in=3, 9 | d_out=1, 10 | d_hidden=256, 11 | n_layers=8, 12 | skip_in=(4,), 13 | multires=0, 14 | bias=0.25, 15 | scale=1, 16 | geometric_init=True, 17 | weight_norm=True, 18 | inside_outside=False): 19 | super(CAPUDFNetwork, self).__init__() 20 | 21 | dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out] 22 | 23 | self.embed_fn_fine = None 24 | 25 | if multires > 0: 26 | embed_fn, input_ch = get_embedder(multires, input_dims=d_in) 27 | self.embed_fn_fine = embed_fn 28 | dims[0] = input_ch 29 | 30 | self.num_layers = len(dims) 31 | self.skip_in = skip_in 32 | self.scale = scale 33 | 34 | for l in range(0, self.num_layers - 1): 35 | if l + 1 in self.skip_in: 36 | out_dim = dims[l + 1] - dims[0] 37 | else: 38 | out_dim = dims[l + 1] 39 | 40 | lin = nn.Linear(dims[l], out_dim) 41 | 42 | if geometric_init: 43 | if l == self.num_layers - 2: 44 | if not inside_outside: 45 | torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001) 46 | torch.nn.init.constant_(lin.bias, -bias) 47 | else: 48 | torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001) 49 | torch.nn.init.constant_(lin.bias, bias) 50 | elif multires > 0 and l == 0: 51 | torch.nn.init.constant_(lin.bias, 0.0) 52 | torch.nn.init.constant_(lin.weight[:, 3:], 0.0) 53 | torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim)) 54 | elif multires > 0 and l in self.skip_in: 55 | torch.nn.init.constant_(lin.bias, 0.0) 56 | torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim)) 57 | torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0) 58 | else: 59 | torch.nn.init.constant_(lin.bias, 0.0) 60 | torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim)) 61 | 62 | if weight_norm: 63 | lin = nn.utils.weight_norm(lin) 64 | 65 | setattr(self, "lin" + str(l), lin) 66 | 67 | #self.activation = nn.Softplus(beta=100) 68 | self.activation = nn.ReLU() 69 | 70 | self.act_last = nn.Sigmoid() 71 | 72 | def forward(self, inputs): 73 | inputs = inputs * self.scale 74 | if self.embed_fn_fine is not None: 75 | inputs = self.embed_fn_fine(inputs) 76 | 77 | x = inputs 78 | for l in range(0, self.num_layers - 1): 79 | lin = getattr(self, "lin" + str(l)) 80 | 81 | if l in self.skip_in: 82 | x = torch.cat([x, inputs], 1) / np.sqrt(2) 83 | 84 | x = lin(x) 85 | 86 | if l < self.num_layers - 2: 87 | x = self.activation(x) 88 | 89 | # res = torch.abs(x) # todo: udf is better for Barn 90 | res = x 91 | return res / self.scale 92 | 93 | def sdf(self, x): 94 | return self.forward(x) 95 | 96 | def udf_hidden_appearance(self, x): 97 | return self.forward(x) 98 | 99 | def gradient(self, x): 100 | x.requires_grad_(True) 101 | y = self.sdf(x) 102 | # y.requires_grad_(True) 103 | d_output = torch.ones_like(y, requires_grad=False, device=y.device) 104 | gradients = torch.autograd.grad( 105 | outputs=y, 106 | inputs=x, 107 | grad_outputs=d_output, 108 | create_graph=True, 109 | retain_graph=True, 110 | only_inputs=True)[0] 111 | return gradients.unsqueeze(1) 112 | 113 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.26.1 2 | torch==1.13.0 3 | torchvision==1.14.0 4 | tqdm 5 | trimesh 6 | pytorch3d==0.7.8 7 | lpips 8 | imageio 9 | Pillow==9.5.0 10 | json 11 | plyfile 12 | opencv-python 13 | tensorboardX 14 | PyMCubes -------------------------------------------------------------------------------- /sugar_scene/__pycache__/cameras.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_scene/__pycache__/cameras.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_scene/__pycache__/gs_model.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_scene/__pycache__/gs_model.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_scene/__pycache__/sugar_densifier.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_scene/__pycache__/sugar_densifier.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_scene/__pycache__/sugar_model.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_scene/__pycache__/sugar_model.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_scene/__pycache__/sugar_optimizer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_scene/__pycache__/sugar_optimizer.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_scene/sugar_optimizer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from sugar_utils.general_utils import get_expon_lr_func 4 | from .sugar_model import SuGaR 5 | 6 | 7 | class OptimizationParams(): 8 | def __init__(self, 9 | iterations:int=30_000, 10 | position_lr_init:float=0.00016, 11 | position_lr_final:float=0.0000016, 12 | position_lr_delay_mult:float=0.01, 13 | position_lr_max_steps:int=30_000, 14 | feature_lr:float=0.0025, 15 | opacity_lr:float=0.05, 16 | scaling_lr:float=0.005, 17 | rotation_lr:float=0.001, 18 | ): 19 | 20 | # Basic Gaussian Splatting 21 | self.iterations = iterations 22 | self.position_lr_init = position_lr_init 23 | self.position_lr_final = position_lr_final 24 | self.position_lr_delay_mult = position_lr_delay_mult 25 | self.position_lr_max_steps = position_lr_max_steps 26 | self.feature_lr = feature_lr 27 | self.opacity_lr = opacity_lr 28 | self.scaling_lr = scaling_lr 29 | self.rotation_lr = rotation_lr 30 | 31 | def __str__(self): 32 | return f"""OptimizationParams( 33 | iterations={self.iterations}, 34 | position_lr_init={self.position_lr_init}, 35 | position_lr_final={self.position_lr_final}, 36 | position_lr_delay_mult={self.position_lr_delay_mult}, 37 | position_lr_max_steps={self.position_lr_max_steps}, 38 | feature_lr={self.feature_lr}, 39 | opacity_lr={self.opacity_lr}, 40 | scaling_lr={self.scaling_lr}, 41 | rotation_lr={self.rotation_lr}, 42 | )""" 43 | 44 | 45 | class SuGaROptimizer(): 46 | """Wrapper of the Adam optimizer used for SuGaR optimization. 47 | Largely inspired by the original implementation of the 3D Gaussian Splatting paper: 48 | https://github.com/graphdeco-inria/gaussian-splatting 49 | """ 50 | def __init__( 51 | self, 52 | model:SuGaR, 53 | opt:OptimizationParams=None, 54 | spatial_lr_scale:float=None, 55 | ) -> None: 56 | 57 | self.current_iteration = 0 58 | self.num_iterations = opt.iterations 59 | 60 | if opt is None: 61 | opt = OptimizationParams() 62 | 63 | if spatial_lr_scale is None: 64 | spatial_lr_scale = model.get_cameras_spatial_extent() 65 | self.spatial_lr_scale = spatial_lr_scale 66 | 67 | if (not model.binded_to_surface_mesh and model.learn_positions) or (model.binded_to_surface_mesh and model.learn_surface_mesh_positions): 68 | l = [{'params': [model._points], 'lr': opt.position_lr_init * spatial_lr_scale, "name": "points"}] 69 | else: 70 | l = [] 71 | 72 | if not model.freeze_gaussians: 73 | l = l + [ 74 | {'params': [model._sh_coordinates_dc], 'lr': opt.feature_lr, "name": "sh_coordinates_dc"}, 75 | {'params': [model._sh_coordinates_rest], 'lr': opt.feature_lr / 20.0, "name": "sh_coordinates_rest"} 76 | ] 77 | 78 | if model.learn_opacities: 79 | l = l + [{'params': [model.all_densities], 'lr': opt.opacity_lr, "name": "all_densities"}] 80 | if (model.binded_to_surface_mesh and model.learn_surface_mesh_scales) or (not model.binded_to_surface_mesh and model.learn_scales): 81 | l = l + [{'params': [model._scales], 'lr': opt.scaling_lr, "name": "scales"}] 82 | if (model.binded_to_surface_mesh and model.learn_surface_mesh_scales) or (not model.binded_to_surface_mesh and model.learn_quaternions): 83 | l = l + [{'params': [model._quaternions], 'lr': opt.rotation_lr, "name": "quaternions"}] 84 | # ours: +np 85 | for i in range(1, model.part_num+1): 86 | sdf_network = getattr(model.neus, 'sdf_network'+str(i)) 87 | l = l + [{'params': list(sdf_network.parameters()), 'lr': 1e-4, "name": "sdfnetwork"+str(i)}] 88 | 89 | self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15) 90 | 91 | self.position_sheduler_func = get_expon_lr_func( 92 | lr_init=opt.position_lr_init * spatial_lr_scale, 93 | lr_final=opt.position_lr_final * spatial_lr_scale, 94 | lr_delay_mult=opt.position_lr_delay_mult, 95 | max_steps=opt.position_lr_max_steps 96 | ) 97 | 98 | def step(self): 99 | self.optimizer.step() 100 | self.current_iteration += 1 101 | 102 | def zero_grad(self, set_to_none:bool=True): 103 | self.optimizer.zero_grad(set_to_none=set_to_none) 104 | 105 | def update_learning_rate(self, iteration:int=None, sdfnet_lr=1e-4): 106 | if iteration is None: 107 | iteration = self.current_iteration 108 | lr = 0. 109 | for param_group in self.optimizer.param_groups: 110 | if param_group["name"] == "points": 111 | lr = self.position_sheduler_func(iteration) 112 | param_group['lr'] = lr 113 | elif "sdfnetwork" in param_group["name"]: 114 | param_group['lr'] = sdfnet_lr 115 | return lr 116 | 117 | def add_param_group(self, new_param_group): 118 | self.optimizer.add_param_group(new_param_group) 119 | 120 | def state_dict(self): 121 | return self.optimizer.state_dict() 122 | 123 | def load_state_dict(self, state_dict): 124 | self.optimizer.load_state_dict(state_dict) 125 | -------------------------------------------------------------------------------- /sugar_trainers/__pycache__/coarse_density.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_trainers/__pycache__/coarse_density.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_trainers/__pycache__/coarse_sdf.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_trainers/__pycache__/coarse_sdf.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_trainers/__pycache__/refine.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_trainers/__pycache__/refine.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_utils/__pycache__/general_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_utils/__pycache__/general_utils.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_utils/__pycache__/graphics_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_utils/__pycache__/graphics_utils.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_utils/__pycache__/loss_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_utils/__pycache__/loss_utils.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_utils/__pycache__/spherical_harmonics.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wen-yuan-zhang/GS-Pull/997da91b72932830cae7848f3a34d02b53370fec/sugar_utils/__pycache__/spherical_harmonics.cpython-310.pyc -------------------------------------------------------------------------------- /sugar_utils/general_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | def PILtoTorch(pil_image, resolution): 5 | resized_image_PIL = pil_image.resize(resolution) 6 | resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 7 | if len(resized_image.shape) == 3: 8 | return resized_image.permute(2, 0, 1) 9 | else: 10 | return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) 11 | 12 | def inverse_sigmoid(x): 13 | return torch.log(x/(1-x)) 14 | 15 | def str2bool(v): 16 | if isinstance(v, bool): 17 | return v 18 | if v.lower() in ('yes', 'true', 't', 'y', '1'): 19 | return True 20 | elif v.lower() in ('no', 'false', 'f', 'n', '0'): 21 | return False 22 | 23 | def get_expon_lr_func( 24 | lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 25 | ): 26 | """ 27 | Copied from 3D Gaussian Splatting (which copied from Plenoxels) 28 | 29 | Continuous learning rate decay function. Adapted from JaxNeRF 30 | The returned rate is lr_init when step=0 and lr_final when step=max_steps, and 31 | is log-linearly interpolated elsewhere (equivalent to exponential decay). 32 | If lr_delay_steps>0 then the learning rate will be scaled by some smooth 33 | function of lr_delay_mult, such that the initial learning rate is 34 | lr_init*lr_delay_mult at the beginning of optimization but will be eased back 35 | to the normal learning rate when steps>lr_delay_steps. 36 | :param conf: config subtree 'lr' or similar 37 | :param max_steps: int, the number of steps during optimization. 38 | :return HoF which takes step as input 39 | """ 40 | 41 | def helper(step): 42 | if step < 0 or (lr_init == 0.0 and lr_final == 0.0): 43 | # Disable this parameter 44 | return 0.0 45 | if lr_delay_steps > 0: 46 | # A kind of reverse cosine decay. 47 | delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( 48 | 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) 49 | ) 50 | else: 51 | delay_rate = 1.0 52 | t = np.clip(step / max_steps, 0, 1) 53 | log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) 54 | return delay_rate * log_lerp 55 | 56 | return helper -------------------------------------------------------------------------------- /sugar_utils/graphics_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import math 14 | import numpy as np 15 | from typing import NamedTuple 16 | 17 | class BasicPointCloud(NamedTuple): 18 | points : np.array 19 | colors : np.array 20 | normals : np.array 21 | 22 | def geom_transform_points(points, transf_matrix): 23 | P, _ = points.shape 24 | ones = torch.ones(P, 1, dtype=points.dtype, device=points.device) 25 | points_hom = torch.cat([points, ones], dim=1) 26 | points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0)) 27 | 28 | denom = points_out[..., 3:] + 0.0000001 29 | return (points_out[..., :3] / denom).squeeze(dim=0) 30 | 31 | # def getWorld2View(R, t): 32 | # Rt = np.zeros((4, 4)) 33 | # Rt[:3, :3] = R.transpose() 34 | # Rt[:3, 3] = t 35 | # Rt[3, 3] = 1.0 36 | # return np.float32(Rt) 37 | 38 | def getWorld2View(R, t, tensor=False): 39 | if tensor: 40 | Rt = torch.zeros(4, 4, device=R.device) 41 | Rt[..., :3, :3] = R.transpose(-1, -2) 42 | Rt[..., :3, 3] = t 43 | Rt[..., 3, 3] = 1.0 44 | return Rt 45 | else: 46 | Rt = np.zeros((4, 4)) 47 | Rt[:3, :3] = R.transpose() 48 | Rt[:3, 3] = t 49 | Rt[3, 3] = 1.0 50 | return np.float32(Rt) 51 | 52 | def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0): 53 | Rt = np.zeros((4, 4)) 54 | Rt[:3, :3] = R.transpose() 55 | Rt[:3, 3] = t 56 | Rt[3, 3] = 1.0 57 | 58 | C2W = np.linalg.inv(Rt) 59 | cam_center = C2W[:3, 3] 60 | cam_center = (cam_center + translate) * scale 61 | C2W[:3, 3] = cam_center 62 | Rt = np.linalg.inv(C2W) 63 | return np.float32(Rt) 64 | 65 | def getProjectionMatrix(znear, zfar, fovX, fovY): 66 | tanHalfFovY = math.tan((fovY / 2)) 67 | tanHalfFovX = math.tan((fovX / 2)) 68 | 69 | top = tanHalfFovY * znear 70 | bottom = -top 71 | right = tanHalfFovX * znear 72 | left = -right 73 | 74 | P = torch.zeros(4, 4) 75 | 76 | z_sign = 1.0 77 | 78 | P[0, 0] = 2.0 * znear / (right - left) 79 | P[1, 1] = 2.0 * znear / (top - bottom) 80 | P[0, 2] = (right + left) / (right - left) 81 | P[1, 2] = (top + bottom) / (top - bottom) 82 | P[3, 2] = z_sign 83 | P[2, 2] = z_sign * zfar / (zfar - znear) 84 | P[2, 3] = -(zfar * znear) / (zfar - znear) 85 | return P 86 | 87 | def fov2focal(fov, pixels): 88 | return pixels / (2 * math.tan(fov / 2)) 89 | 90 | def focal2fov(focal, pixels): 91 | return 2*math.atan(pixels/(2*focal)) -------------------------------------------------------------------------------- /sugar_utils/loss_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import torch.nn.functional as F 14 | from torch.autograd import Variable 15 | from math import exp 16 | 17 | def l1_loss(network_output, gt): 18 | return torch.abs((network_output - gt)).mean() 19 | 20 | def l2_loss(network_output, gt): 21 | return ((network_output - gt) ** 2).mean() 22 | 23 | def gaussian(window_size, sigma): 24 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 25 | return gauss / gauss.sum() 26 | 27 | def create_window(window_size, channel): 28 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 29 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 30 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 31 | return window 32 | 33 | def ssim(img1, img2, window_size=11, size_average=True): 34 | channel = img1.size(-3) 35 | window = create_window(window_size, channel) 36 | 37 | if img1.is_cuda: 38 | window = window.cuda(img1.get_device()) 39 | window = window.type_as(img1) 40 | 41 | return _ssim(img1, img2, window, window_size, channel, size_average) 42 | 43 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 44 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 45 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 46 | 47 | mu1_sq = mu1.pow(2) 48 | mu2_sq = mu2.pow(2) 49 | mu1_mu2 = mu1 * mu2 50 | 51 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 52 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 53 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 54 | 55 | C1 = 0.01 ** 2 56 | C2 = 0.03 ** 2 57 | 58 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 59 | 60 | if size_average: 61 | return ssim_map.mean() 62 | else: 63 | return ssim_map.mean(1).mean(1).mean(1) -------------------------------------------------------------------------------- /sugar_utils/spherical_harmonics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | 5 | C0 = 0.28209479177387814 6 | C1 = 0.4886025119029199 7 | C2 = [ 8 | 1.0925484305920792, 9 | -1.0925484305920792, 10 | 0.31539156525252005, 11 | -1.0925484305920792, 12 | 0.5462742152960396 13 | ] 14 | C3 = [ 15 | -0.5900435899266435, 16 | 2.890611442640554, 17 | -0.4570457994644658, 18 | 0.3731763325901154, 19 | -0.4570457994644658, 20 | 1.445305721320277, 21 | -0.5900435899266435 22 | ] 23 | C4 = [ 24 | 2.5033429417967046, 25 | -1.7701307697799304, 26 | 0.9461746957575601, 27 | -0.6690465435572892, 28 | 0.10578554691520431, 29 | -0.6690465435572892, 30 | 0.47308734787878004, 31 | -1.7701307697799304, 32 | 0.6258357354491761, 33 | ] 34 | 35 | 36 | def get_cartesian_coords(r, elev, azim, in_degrees=False): 37 | """ 38 | Returns the cartesian coordinates of 3D points written in spherical coordinates. 39 | :param r: (Tensor) Radius tensor of 3D points, with shape (N). 40 | :param elev: (Tensor) Elevation tensor of 3D points, with shape (N). 41 | :param azim: (Tensor) Azimuth tensor of 3D points, with shape (N). 42 | :param in_degrees: (bool) In True, elevation and azimuth are written in degrees. 43 | Else, in radians. 44 | :return: (Tensor) Cartesian coordinates tensor with shape (N, 3). 45 | """ 46 | factor = 1 47 | if in_degrees: 48 | factor *= np.pi / 180. 49 | X = torch.stack(( 50 | torch.cos(factor * elev) * torch.sin(factor * azim), 51 | torch.sin(factor * elev), 52 | torch.cos(factor * elev) * torch.cos(factor * azim) 53 | ), dim=2) 54 | 55 | return r * X.view(-1, 3) 56 | 57 | 58 | def get_spherical_coords(X): 59 | """ 60 | Returns the spherical coordinates of 3D points written in cartesian coordinates 61 | :param X: (Tensor) Tensor with shape (N, 3) that represents 3D points in cartesian coordinates. 62 | :return: (3-tuple of Tensors) r_x, elev_x and azim_x are Tensors with shape (N) that corresponds 63 | to radius, elevation and azimuths of all 3D points. 64 | """ 65 | r_x = torch.linalg.norm(X, dim=1) 66 | 67 | elev_x = torch.asin(X[:, 1] / r_x) # between -pi/2 and pi/2 68 | elev_x[X[:, 1] / r_x <= -1] = -np.pi / 2 69 | elev_x[X[:, 1] / r_x >= 1] = np.pi / 2 70 | 71 | azim_x = torch.acos(X[:, 2] / (r_x * torch.cos(elev_x))) 72 | azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) <= -1] = np.pi 73 | azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) >= 1] = 0. 74 | azim_x[X[:, 0] < 0] *= -1 75 | 76 | return r_x, elev_x, azim_x 77 | 78 | 79 | def get_samples_on_sphere(device, pole_samples=False, n_elev=10, n_azim=2*10): 80 | """ 81 | Returns cameras candidate positions, sampled on a sphere. 82 | :param params: (Params) The dictionary of parameters. 83 | :param device: 84 | :return: A tuple of Tensors (X_cam, candidate_dist, candidate_elev, candidate_azim) 85 | X_cam has shape (n_camera_candidate, 3) 86 | All other tensors have shape (n_camera candidate, ) 87 | """ 88 | n_camera = n_elev * n_azim 89 | if pole_samples: 90 | n_camera += 2 91 | 92 | candidate_dist = torch.Tensor([1. for i in range(n_camera)]).to(device) 93 | 94 | candidate_elev = [-90. + (i + 1) / (n_elev + 1) * 180. 95 | for i in range(n_elev) 96 | for j in range(n_azim)] 97 | 98 | candidate_azim = [360. * j / n_azim 99 | for i in range(n_elev) 100 | for j in range(n_azim)] 101 | 102 | if pole_samples: 103 | candidate_elev = [-89.9] + candidate_elev + [89.9] 104 | candidate_azim = [0.] + candidate_azim + [0.] 105 | 106 | candidate_elev = torch.Tensor(candidate_elev).to(device).view(-1, 1) 107 | candidate_azim = torch.Tensor(candidate_azim).to(device).view(-1, 1) 108 | 109 | X_cam = get_cartesian_coords(r=candidate_dist.view(-1, 1), 110 | elev=candidate_elev, 111 | azim=candidate_azim, 112 | in_degrees=True) 113 | 114 | return X_cam, candidate_dist, candidate_elev, candidate_azim 115 | 116 | 117 | def eval_sh(deg, sh, dirs): 118 | """ 119 | Evaluate spherical harmonics at unit directions 120 | using hardcoded SH polynomials. 121 | Works with torch/np/jnp. 122 | ... Can be 0 or more batch dimensions. 123 | Args: 124 | deg: int SH deg. Currently, 0-3 supported 125 | sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] 126 | dirs: jnp.ndarray unit directions [..., 3] 127 | Returns: 128 | [..., C] 129 | """ 130 | assert deg <= 4 and deg >= 0 131 | coeff = (deg + 1) ** 2 132 | assert sh.shape[-1] >= coeff 133 | 134 | result = C0 * sh[..., 0] 135 | if deg > 0: 136 | x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] 137 | result = (result - 138 | C1 * y * sh[..., 1] + 139 | C1 * z * sh[..., 2] - 140 | C1 * x * sh[..., 3]) 141 | 142 | if deg > 1: 143 | xx, yy, zz = x * x, y * y, z * z 144 | xy, yz, xz = x * y, y * z, x * z 145 | result = (result + 146 | C2[0] * xy * sh[..., 4] + 147 | C2[1] * yz * sh[..., 5] + 148 | C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + 149 | C2[3] * xz * sh[..., 7] + 150 | C2[4] * (xx - yy) * sh[..., 8]) 151 | 152 | if deg > 2: 153 | result = (result + 154 | C3[0] * y * (3 * xx - yy) * sh[..., 9] + 155 | C3[1] * xy * z * sh[..., 10] + 156 | C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + 157 | C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + 158 | C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + 159 | C3[5] * z * (xx - yy) * sh[..., 14] + 160 | C3[6] * x * (xx - 3 * yy) * sh[..., 15]) 161 | 162 | if deg > 3: 163 | result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + 164 | C4[1] * yz * (3 * xx - yy) * sh[..., 17] + 165 | C4[2] * xy * (7 * zz - 1) * sh[..., 18] + 166 | C4[3] * yz * (7 * zz - 3) * sh[..., 19] + 167 | C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + 168 | C4[5] * xz * (7 * zz - 3) * sh[..., 21] + 169 | C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + 170 | C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + 171 | C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) 172 | return result 173 | 174 | def RGB2SH(rgb): 175 | return (rgb - 0.5) / C0 176 | 177 | def SH2RGB(sh): 178 | return sh * C0 + 0.5 -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import numpy as np 4 | import random 5 | from sugar_utils.general_utils import str2bool 6 | from sugar_trainers.coarse_sdf import coarse_training_with_sdf_regularization 7 | 8 | 9 | class AttrDict(dict): 10 | def __init__(self, *args, **kwargs): 11 | super().__init__(*args, **kwargs) 12 | self.__dict__ = self 13 | 14 | 15 | if __name__ == "__main__": 16 | # ours: fixed random seed 17 | # seed = 1 18 | # torch.manual_seed(seed) 19 | # torch.cuda.manual_seed(seed) 20 | # np.random.seed(seed) 21 | # random.seed(seed) 22 | # torch.backends.cudnn.deterministic = True 23 | # ----- Parser ----- 24 | parser = argparse.ArgumentParser(description='Script to optimize a full SuGaR model.') 25 | 26 | # Data and vanilla 3DGS checkpoint 27 | parser.add_argument('-s', '--scene_path', 28 | type=str, 29 | help='(Required) path to the scene data to use.') 30 | parser.add_argument('-c', '--checkpoint_path', 31 | type=str, 32 | help='(Required) path to the vanilla 3D Gaussian Splatting Checkpoint to load.') 33 | parser.add_argument('-i', '--iteration_to_load', 34 | type=int, default=7000, 35 | help='iteration to load.') 36 | 37 | # Regularization for coarse SuGaR 38 | parser.add_argument('-r', '--regularization_type', type=str, default='sdf', 39 | help='(Required) Type of regularization to use for coarse SuGaR. Can be "sdf" or "density". ' 40 | 'For reconstructing detailed objects centered in the scene with 360° coverage, "density" provides a better foreground mesh. ' 41 | 'For a stronger regularization and a better balance between foreground and background, choose "sdf".') 42 | 43 | # Extract mesh 44 | parser.add_argument('-l', '--surface_level', type=float, default=0.3, 45 | help='Surface level to extract the mesh at. Default is 0.3') 46 | parser.add_argument('-v', '--n_vertices_in_mesh', type=int, default=1_000_000, 47 | help='Number of vertices in the extracted mesh.') 48 | parser.add_argument('-b', '--bboxmin', type=str, default=None, 49 | help='Min coordinates to use for foreground.') 50 | parser.add_argument('-B', '--bboxmax', type=str, default=None, 51 | help='Max coordinates to use for foreground.') 52 | parser.add_argument('--center_bbox', type=str2bool, default=True, 53 | help='If True, center the bbox. Default is False.') 54 | 55 | # Parameters for refined SuGaR 56 | parser.add_argument('-g', '--gaussians_per_triangle', type=int, default=1, 57 | help='Number of gaussians per triangle.') 58 | parser.add_argument('-f', '--refinement_iterations', type=int, default=15_000, 59 | help='Number of refinement iterations.') 60 | 61 | # (Optional) Parameters for textured mesh extraction 62 | parser.add_argument('-t', '--export_uv_textured_mesh', type=str2bool, default=True, 63 | help='If True, will export a textured mesh as an .obj file from the refined SuGaR model. ' 64 | 'Computing a traditional colored UV texture should take less than 10 minutes.') 65 | parser.add_argument('--square_size', 66 | default=10, type=int, help='Size of the square to use for the UV texture.') 67 | parser.add_argument('--postprocess_mesh', type=str2bool, default=False, 68 | help='If True, postprocess the mesh by removing border triangles with low-density. ' 69 | 'This step takes a few minutes and is not needed in general, as it can also be risky. ' 70 | 'However, it increases the quality of the mesh in some cases, especially when an object is visible only from one side.') 71 | parser.add_argument('--postprocess_density_threshold', type=float, default=0.1, 72 | help='Threshold to use for postprocessing the mesh.') 73 | parser.add_argument('--postprocess_iterations', type=int, default=5, 74 | help='Number of iterations to use for postprocessing the mesh.') 75 | 76 | # (Optional) Default configurations 77 | parser.add_argument('--low_poly', type=str2bool, default=False, 78 | help='Use standard config for a low poly mesh, with 200k vertices and 6 Gaussians per triangle.') 79 | parser.add_argument('--high_poly', type=str2bool, default=False, 80 | help='Use standard config for a high poly mesh, with 1M vertices and 1 Gaussians per triangle.') 81 | parser.add_argument('--refinement_time', type=str, default=None, 82 | help="Default configs for time to spend on refinement. Can be 'short', 'medium' or 'long'.") 83 | 84 | # Evaluation split 85 | parser.add_argument('--eval', type=str2bool, default=False, help='Use eval split.') 86 | 87 | # GPU 88 | parser.add_argument('--gpu', type=int, default=0, help='Index of GPU device to use.') 89 | 90 | # ours 91 | parser.add_argument('--dataset_name', default="real360", help='blender, real360, relight3d') 92 | parser.add_argument('-w', '--white_bg', action='store_true', help="Use white background (default: False)") # dtu black; tnt white; mipnerf360 white 93 | parser.add_argument('--output', type=str, default='output/mipnerf360', help='output directory(do not include experiment name') # TODO: change this 94 | parser.add_argument('--resolution', type=int, default=1, help='image resolution. (Courthouse 2 especially)') # TODO: change this 95 | 96 | # Parse arguments 97 | args = parser.parse_args() 98 | if args.low_poly: 99 | args.n_vertices_in_mesh = 200_000 100 | args.gaussians_per_triangle = 6 101 | print('Using low poly config.') 102 | if args.high_poly: 103 | args.n_vertices_in_mesh = 1_000_000 104 | args.gaussians_per_triangle = 1 105 | print('Using high poly config.') 106 | if args.refinement_time == 'short': 107 | args.refinement_iterations = 2_000 108 | print('Using short refinement time.') 109 | if args.refinement_time == 'medium': 110 | args.refinement_iterations = 7_000 111 | print('Using medium refinement time.') 112 | if args.refinement_time == 'long': 113 | args.refinement_iterations = 15_000 114 | print('Using long refinement time.') 115 | if args.export_uv_textured_mesh: 116 | print('Will export a UV-textured mesh as an .obj file.') 117 | 118 | # ----- Optimize coarse SuGaR ----- 119 | coarse_args = AttrDict({ 120 | 'checkpoint_path': args.checkpoint_path, 121 | 'scene_path': args.scene_path, 122 | 'iteration_to_load': args.iteration_to_load, 123 | 'output_dir': args.output, 124 | 'eval': args.eval, 125 | 'estimation_factor': 0.2, 126 | 'normal_factor': 0.0, 127 | 'gpu': args.gpu, 128 | 'dataset_name': args.dataset_name, 129 | 'white_bg': args.white_bg, 130 | 'image_resolution': args.resolution, 131 | }) 132 | coarse_sugar_path = coarse_training_with_sdf_regularization(coarse_args) 133 | --------------------------------------------------------------------------------