├── imgs ├── Fig1.png ├── Fig2.png ├── Fig3.png ├── Fig5.png ├── Fig6.png ├── Table1.png ├── wechat_code.jpg └── 3DV_demo_cover.png ├── cpp_wrappers ├── compile_wrappers.sh ├── cpp_subsampling │ ├── setup.py │ ├── grid_subsampling │ │ ├── grid_subsampling.h │ │ └── grid_subsampling.cpp │ └── wrapper.cpp └── cpp_utils │ └── cloud │ ├── cloud.cpp │ └── cloud.h ├── compile_op.sh ├── utils ├── cpp_wrappers │ ├── compile_wrappers.sh │ ├── cpp_subsampling │ │ ├── setup.py │ │ ├── grid_subsampling │ │ │ ├── grid_subsampling.h │ │ │ └── grid_subsampling.cpp │ │ └── wrapper.cpp │ └── cpp_utils │ │ └── cloud │ │ ├── cloud.cpp │ │ └── cloud.h └── nearest_neighbors │ ├── test.py │ ├── setup.py │ ├── knn_.h │ ├── knn.pyx │ ├── knn_.cxx │ └── KDTreeTableAdaptor.h ├── helper_requirements.txt ├── LICENSE ├── .gitignore ├── input_preparation.py ├── tester_SensatUrban.py ├── README.md ├── helper_ply.py ├── tool.py ├── main_SensatUrban.py ├── RandLANet.py └── tf_util.py /imgs/Fig1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Fig1.png -------------------------------------------------------------------------------- /imgs/Fig2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Fig2.png -------------------------------------------------------------------------------- /imgs/Fig3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Fig3.png -------------------------------------------------------------------------------- /imgs/Fig5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Fig5.png -------------------------------------------------------------------------------- /imgs/Fig6.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Fig6.png -------------------------------------------------------------------------------- /imgs/Table1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/Table1.png -------------------------------------------------------------------------------- /imgs/wechat_code.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/wechat_code.jpg -------------------------------------------------------------------------------- /imgs/3DV_demo_cover.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/SensatUrban/HEAD/imgs/3DV_demo_cover.png -------------------------------------------------------------------------------- /cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /compile_op.sh: -------------------------------------------------------------------------------- 1 | cd utils/nearest_neighbors 2 | python setup.py install --home="." 3 | cd ../../ 4 | 5 | cd utils/cpp_wrappers 6 | sh compile_wrappers.sh 7 | cd ../../../ -------------------------------------------------------------------------------- /utils/cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /helper_requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.16.1 2 | h5py==2.10.0 3 | cython==0.29.15 4 | open3d-python==0.3.0 5 | pandas==0.25.3 6 | scikit-learn==0.21.3 7 | scipy==1.4.1 8 | PyYAML==5.4 9 | 10 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lib.python.nearest_neighbors as nearest_neighbors 3 | import time 4 | 5 | batch_size = 16 6 | num_points = 81920 7 | K = 16 8 | pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) 9 | 10 | # nearest neighbours 11 | start = time.time() 12 | neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) 13 | print(time.time() - start) 14 | 15 | 16 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | from Cython.Distutils import build_ext 4 | import numpy 5 | 6 | 7 | 8 | ext_modules = [Extension( 9 | "nearest_neighbors", 10 | sources=["knn.pyx", "knn_.cxx",], # source file(s) 11 | include_dirs=["./", numpy.get_include()], 12 | language="c++", 13 | extra_compile_args = [ "-std=c++11", "-fopenmp",], 14 | extra_link_args=["-std=c++11", '-fopenmp'], 15 | )] 16 | 17 | setup( 18 | name = "KNN NanoFLANN", 19 | ext_modules = ext_modules, 20 | cmdclass = {'build_ext': build_ext}, 21 | ) 22 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Qingyong 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include 4 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 5 | const float* queries, const size_t nqueries, 6 | const size_t K, long* indices); 7 | 8 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 9 | const float* queries, const size_t nqueries, 10 | const size_t K, long* indices); 11 | 12 | 13 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* batch_indices); 16 | 17 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices); 20 | 21 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices); 24 | 25 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* batch_queries, const size_t nqueries, 27 | const size_t K, long* batch_indices); -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /input_preparation.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname 2 | from sklearn.neighbors import KDTree 3 | from tool import DataProcessing as DP 4 | from helper_ply import write_ply 5 | import numpy as np 6 | import os, pickle, argparse 7 | 8 | if __name__ == '__main__': 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--dataset_path', type=str, required=True, help='the number of GPUs to use [default: 0]') 11 | FLAGS = parser.parse_args() 12 | dataset_name = 'SensatUrban' 13 | dataset_path = FLAGS.dataset_path 14 | preparation_types = ['grid'] # Grid sampling & Random sampling 15 | grid_size = 0.2 16 | random_sample_ratio = 10 17 | train_files = np.sort([join(dataset_path, 'train', i) for i in os.listdir(join(dataset_path, 'train'))]) 18 | test_files = np.sort([join(dataset_path, 'test', i) for i in os.listdir(join(dataset_path, 'test'))]) 19 | files = np.sort(np.hstack((train_files, test_files))) 20 | 21 | for sample_type in preparation_types: 22 | for pc_path in files: 23 | cloud_name = pc_path.split('/')[-1][:-4] 24 | print('start to process:', cloud_name) 25 | 26 | # create output directory 27 | out_folder = join(dirname(dataset_path), sample_type + '_{:.3f}'.format(grid_size)) 28 | os.makedirs(out_folder) if not exists(out_folder) else None 29 | 30 | # check if it has already calculated 31 | if exists(join(out_folder, cloud_name + '_KDTree.pkl')): 32 | print(cloud_name, 'already exists, skipped') 33 | continue 34 | 35 | if pc_path in train_files: 36 | xyz, rgb, labels = DP.read_ply_data(pc_path, with_rgb=True) 37 | else: 38 | xyz, rgb = DP.read_ply_data(pc_path, with_rgb=True, with_label=False) 39 | labels = np.zeros(len(xyz), dtype=np.uint8) 40 | 41 | sub_ply_file = join(out_folder, cloud_name + '.ply') 42 | if sample_type == 'grid': 43 | sub_xyz, sub_rgb, sub_labels = DP.grid_sub_sampling(xyz, rgb, labels, grid_size) 44 | else: 45 | sub_xyz, sub_rgb, sub_labels = DP.random_sub_sampling(xyz, rgb, labels, random_sample_ratio) 46 | 47 | sub_rgb = sub_rgb / 255.0 48 | sub_labels = np.squeeze(sub_labels) 49 | write_ply(sub_ply_file, [sub_xyz, sub_rgb, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 50 | 51 | search_tree = KDTree(sub_xyz, leaf_size=50) 52 | kd_tree_file = join(out_folder, cloud_name + '_KDTree.pkl') 53 | with open(kd_tree_file, 'wb') as f: 54 | pickle.dump(search_tree, f) 55 | 56 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 57 | proj_idx = proj_idx.astype(np.int32) 58 | proj_save = join(out_folder, cloud_name + '_proj.pkl') 59 | with open(proj_save, 'wb') as f: 60 | pickle.dump([proj_idx, labels], f) 61 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initiate variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initiate variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 101 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 102 | } 103 | } 104 | 105 | return; 106 | } 107 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initiate variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initiate variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 101 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 102 | } 103 | } 104 | 105 | return; 106 | } 107 | -------------------------------------------------------------------------------- /tester_SensatUrban.py: -------------------------------------------------------------------------------- 1 | from os import makedirs, system 2 | from os.path import exists, join, dirname, abspath 3 | import tensorflow as tf 4 | import numpy as np 5 | import time 6 | 7 | 8 | def log_out(out_str, log_f_out): 9 | log_f_out.write(out_str + '\n') 10 | log_f_out.flush() 11 | print(out_str) 12 | 13 | 14 | class ModelTester: 15 | def __init__(self, model, dataset, restore_snap=None): 16 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 17 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 18 | self.Log_file = open('log_test_' + dataset.name + '.txt', 'a') 19 | 20 | # Create a session for running Ops on the Graph. 21 | on_cpu = False 22 | if on_cpu: 23 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 24 | else: 25 | c_proto = tf.ConfigProto() 26 | c_proto.gpu_options.allow_growth = True 27 | self.sess = tf.Session(config=c_proto) 28 | self.sess.run(tf.global_variables_initializer()) 29 | 30 | # Load trained model 31 | if restore_snap is not None: 32 | self.saver.restore(self.sess, restore_snap) 33 | print("Model restored from " + restore_snap) 34 | 35 | self.prob_logits = tf.nn.softmax(model.logits) 36 | 37 | # Initiate global prediction over all test clouds 38 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 39 | for l in dataset.input_labels['test']] 40 | 41 | def test(self, model, dataset, num_votes=100): 42 | 43 | # Smoothing parameter for votes 44 | test_smooth = 0.95 45 | 46 | # Initialise iterator with validation/test data 47 | self.sess.run(dataset.test_init_op) 48 | 49 | # Test saving path 50 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 51 | test_path = join('test', saving_path.split('/')[-1]) 52 | makedirs(test_path) if not exists(test_path) else None 53 | makedirs(join(test_path, 'test_preds')) if not exists(join(test_path, 'test_preds')) else None 54 | 55 | step_id = 0 56 | epoch_id = 0 57 | last_min = -0.5 58 | 59 | while last_min < num_votes: 60 | try: 61 | ops = (self.prob_logits, 62 | model.labels, 63 | model.inputs['input_inds'], 64 | model.inputs['cloud_inds']) 65 | 66 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 67 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 68 | model.config.num_classes]) 69 | 70 | for j in range(np.shape(stacked_probs)[0]): 71 | probs = stacked_probs[j, :, :] 72 | p_idx = point_idx[j, :] 73 | c_i = cloud_idx[j][0] 74 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 75 | step_id += 1 76 | 77 | except tf.errors.OutOfRangeError: 78 | 79 | new_min = np.min(dataset.min_possibility['test']) 80 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 81 | 82 | if last_min + 1 < new_min: 83 | 84 | # Update last_min 85 | last_min += 1 86 | 87 | # Show vote results (On subcloud so it is not the good values here) 88 | log_out('\nConfusion on sub clouds', self.Log_file) 89 | num_test = len(dataset.input_labels['test']) 90 | 91 | # Project predictions 92 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 93 | proj_probs_list = [] 94 | 95 | for i_test in range(num_test): 96 | # Reproject probs back to the evaluations points 97 | proj_idx = dataset.test_proj[i_test] 98 | probs = self.test_probs[i_test][proj_idx, :] 99 | proj_probs_list += [probs] 100 | 101 | # Show vote results 102 | log_out('Confusion on full clouds', self.Log_file) 103 | for i_test in range(num_test): 104 | # Get the predicted labels 105 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 106 | save_name = join(test_path, 'test_preds', dataset.input_names['test'][i_test] + '.label') 107 | preds = preds.astype(np.uint8) 108 | preds.tofile(save_name) 109 | 110 | # creat submission files 111 | base_dir = dirname(abspath(__file__)) 112 | results_path = join(base_dir, test_path, 'test_preds') 113 | system('cd %s && zip -r %s/submission.zip *.label' % (results_path, results_path)) 114 | return 115 | 116 | self.sess.run(dataset.test_init_op) 117 | epoch_id += 1 118 | step_id = 0 119 | continue 120 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | # distutils: sources = knn.cxx 3 | 4 | import numpy as np 5 | cimport numpy as np 6 | import cython 7 | 8 | cdef extern from "knn_.h": 9 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 10 | const float* queries, const size_t nqueries, 11 | const size_t K, long* indices) 12 | 13 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* indices) 16 | 17 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices) 20 | 21 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | const float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices) 24 | 25 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* queries, const size_t nqueries, 27 | const size_t K, long* batch_indices) 28 | 29 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 30 | float* batch_queries, const size_t nqueries, 31 | const size_t K, long* batch_indices) 32 | 33 | def knn(pts, queries, K, omp=False): 34 | 35 | # define shape parameters 36 | cdef int npts 37 | cdef int dim 38 | cdef int K_cpp 39 | cdef int nqueries 40 | 41 | # define tables 42 | cdef np.ndarray[np.float32_t, ndim=2] pts_cpp 43 | cdef np.ndarray[np.float32_t, ndim=2] queries_cpp 44 | cdef np.ndarray[np.int64_t, ndim=2] indices_cpp 45 | 46 | # set shape values 47 | npts = pts.shape[0] 48 | nqueries = queries.shape[0] 49 | dim = pts.shape[1] 50 | K_cpp = K 51 | 52 | # create indices tensor 53 | indices = np.zeros((queries.shape[0], K), dtype=np.int64) 54 | 55 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 56 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 57 | indices_cpp = indices 58 | 59 | # normal estimation 60 | if omp: 61 | cpp_knn_omp( pts_cpp.data, npts, dim, 62 | queries_cpp.data, nqueries, 63 | K_cpp, indices_cpp.data) 64 | else: 65 | cpp_knn( pts_cpp.data, npts, dim, 66 | queries_cpp.data, nqueries, 67 | K_cpp, indices_cpp.data) 68 | 69 | return indices 70 | 71 | def knn_batch(pts, queries, K, omp=False): 72 | 73 | # define shape parameters 74 | cdef int batch_size 75 | cdef int npts 76 | cdef int nqueries 77 | cdef int K_cpp 78 | cdef int dim 79 | 80 | # define tables 81 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 82 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 83 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 84 | 85 | # set shape values 86 | batch_size = pts.shape[0] 87 | npts = pts.shape[1] 88 | dim = pts.shape[2] 89 | nqueries = queries.shape[1] 90 | K_cpp = K 91 | 92 | # create indices tensor 93 | indices = np.zeros((pts.shape[0], queries.shape[1], K), dtype=np.int64) 94 | 95 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 96 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 97 | indices_cpp = indices 98 | 99 | # normal estimation 100 | if omp: 101 | cpp_knn_batch_omp( pts_cpp.data, batch_size, npts, dim, 102 | queries_cpp.data, nqueries, 103 | K_cpp, indices_cpp.data) 104 | else: 105 | cpp_knn_batch( pts_cpp.data, batch_size, npts, dim, 106 | queries_cpp.data, nqueries, 107 | K_cpp, indices_cpp.data) 108 | 109 | return indices 110 | 111 | def knn_batch_distance_pick(pts, nqueries, K, omp=False): 112 | 113 | # define shape parameters 114 | cdef int batch_size 115 | cdef int npts 116 | cdef int nqueries_cpp 117 | cdef int K_cpp 118 | cdef int dim 119 | 120 | # define tables 121 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 122 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 123 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 124 | 125 | # set shape values 126 | batch_size = pts.shape[0] 127 | npts = pts.shape[1] 128 | dim = pts.shape[2] 129 | nqueries_cpp = nqueries 130 | K_cpp = K 131 | 132 | # create indices tensor 133 | indices = np.zeros((pts.shape[0], nqueries, K), dtype=np.long) 134 | queries = np.zeros((pts.shape[0], nqueries, dim), dtype=np.float32) 135 | 136 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 137 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 138 | indices_cpp = indices 139 | 140 | if omp: 141 | cpp_knn_batch_distance_pick_omp( pts_cpp.data, batch_size, npts, dim, 142 | queries_cpp.data, nqueries, 143 | K_cpp, indices_cpp.data) 144 | else: 145 | cpp_knn_batch_distance_pick( pts_cpp.data, batch_size, npts, dim, 146 | queries_cpp.data, nqueries, 147 | K_cpp, indices_cpp.data) 148 | 149 | return indices, queries -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.cxx: -------------------------------------------------------------------------------- 1 | 2 | #include "knn_.h" 3 | #include "nanoflann.hpp" 4 | using namespace nanoflann; 5 | 6 | #include "KDTreeTableAdaptor.h" 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | using namespace std; 19 | 20 | 21 | 22 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 23 | const float* queries, const size_t nqueries, 24 | const size_t K, long* indices){ 25 | 26 | // create the kdtree 27 | typedef KDTreeTableAdaptor< float, float> KDTree; 28 | KDTree mat_index(npts, dim, points, 10); 29 | mat_index.index->buildIndex(); 30 | 31 | std::vector out_dists_sqr(K); 32 | std::vector out_ids(K); 33 | 34 | // iterate over the points 35 | for(size_t i=0; i resultSet(K); 38 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 39 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 40 | for(size_t j=0; j KDTree; 52 | KDTree mat_index(npts, dim, points, 10); 53 | mat_index.index->buildIndex(); 54 | 55 | 56 | // iterate over the points 57 | # pragma omp parallel for 58 | for(size_t i=0; i out_ids(K); 60 | std::vector out_dists_sqr(K); 61 | 62 | nanoflann::KNNResultSet resultSet(K); 63 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 64 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 65 | for(size_t j=0; j KDTree; 83 | KDTree mat_index(npts, dim, points, 10); 84 | 85 | mat_index.index->buildIndex(); 86 | 87 | std::vector out_dists_sqr(K); 88 | std::vector out_ids(K); 89 | 90 | // iterate over the points 91 | for(size_t i=0; i resultSet(K); 93 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 94 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 95 | for(size_t j=0; j KDTree; 116 | KDTree mat_index(npts, dim, points, 10); 117 | 118 | mat_index.index->buildIndex(); 119 | 120 | std::vector out_dists_sqr(K); 121 | std::vector out_ids(K); 122 | 123 | // iterate over the points 124 | for(size_t i=0; i resultSet(K); 126 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 127 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 128 | for(size_t j=0; j KDTree; 153 | KDTree tree(npts, dim, points, 10); 154 | tree.index->buildIndex(); 155 | 156 | vector used(npts, 0); 157 | int current_id = 0; 158 | for(size_t ptid=0; ptid possible_ids; 162 | while(possible_ids.size() == 0){ 163 | for(size_t i=0; i query(3); 178 | for(size_t i=0; i dists(K); 183 | std::vector ids(K); 184 | nanoflann::KNNResultSet resultSet(K); 185 | resultSet.init(&ids[0], &dists[0] ); 186 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 187 | 188 | for(size_t i=0; i KDTree; 221 | KDTree tree(npts, dim, points, 10); 222 | tree.index->buildIndex(); 223 | 224 | vector used(npts, 0); 225 | int current_id = 0; 226 | for(size_t ptid=0; ptid possible_ids; 230 | while(possible_ids.size() == 0){ 231 | for(size_t i=0; i query(3); 246 | for(size_t i=0; i dists(K); 251 | std::vector ids(K); 252 | nanoflann::KNNResultSet resultSet(K); 253 | resultSet.init(&ids[0], &dists[0] ); 254 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 255 | 256 | for(size_t i=0; i 34 | 35 | // ===== This example shows how to use nanoflann with these types of containers: ======= 36 | //typedef std::vector > my_vector_of_vectors_t; 37 | //typedef std::vector my_vector_of_vectors_t; // This requires #include 38 | // ===================================================================================== 39 | 40 | 41 | /** A simple vector-of-vectors adaptor for nanoflann, without duplicating the storage. 42 | * The i'th vector represents a point in the state space. 43 | * 44 | * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality for the points in the data set, allowing more compiler optimizations. 45 | * \tparam num_t The type of the point coordinates (typically, double or float). 46 | * \tparam Distance The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. 47 | * \tparam IndexType The type for indices in the KD-tree index (typically, size_t of int) 48 | */ 49 | // template 50 | // struct KDTreeVectorAdaptor 51 | // { 52 | // typedef KDTreeVectorAdaptor self_t; 53 | // typedef typename Distance::template traits::distance_t metric_t; 54 | // typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 55 | 56 | // index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 57 | // size_t dims; 58 | 59 | // /// Constructor: takes a const ref to the vector of vectors object with the data points 60 | // KDTreeVectorAdaptor(const size_t dims /* dimensionality */, const VectorType &mat, const int leaf_max_size = 10) : m_data(mat) 61 | // { 62 | // assert(mat.size() != 0); 63 | // this->dims= dims; 64 | // index = new index_t( static_cast(dims), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 65 | // index->buildIndex(); 66 | // } 67 | 68 | // ~KDTreeVectorAdaptor() { 69 | // delete index; 70 | // } 71 | 72 | // const VectorType &m_data; 73 | 74 | // /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 75 | // * Note that this is a short-cut method for index->findNeighbors(). 76 | // * The user can also call index->... methods as desired. 77 | // * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 78 | // */ 79 | // inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 80 | // { 81 | // nanoflann::KNNResultSet resultSet(num_closest); 82 | // resultSet.init(out_indices, out_distances_sq); 83 | // index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 84 | // } 85 | 86 | // /** @name Interface expected by KDTreeSingleIndexAdaptor 87 | // * @{ */ 88 | 89 | // const self_t & derived() const { 90 | // return *this; 91 | // } 92 | // self_t & derived() { 93 | // return *this; 94 | // } 95 | 96 | // // Must return the number of data points 97 | // inline size_t kdtree_get_point_count() const { 98 | // return m_data.size()/this->dims; 99 | // } 100 | 101 | // // Returns the dim'th component of the idx'th point in the class: 102 | // inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const { 103 | // return m_data[idx*this->dims + dim]; 104 | // } 105 | 106 | // // Optional bounding-box computation: return false to default to a standard bbox computation loop. 107 | // // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 108 | // // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 109 | // template 110 | // bool kdtree_get_bbox(BBOX & /*bb*/) const { 111 | // return false; 112 | // } 113 | 114 | // /** @} */ 115 | 116 | // }; // end of KDTreeVectorOfVectorsAdaptor 117 | 118 | 119 | 120 | 121 | template 122 | struct KDTreeTableAdaptor 123 | { 124 | typedef KDTreeTableAdaptor self_t; 125 | typedef typename Distance::template traits::distance_t metric_t; 126 | typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 127 | 128 | index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 129 | size_t dim; 130 | size_t npts; 131 | const TableType* m_data; 132 | 133 | /// Constructor: takes a const ref to the vector of vectors object with the data points 134 | KDTreeTableAdaptor(const size_t npts, const size_t dim, const TableType* mat, const int leaf_max_size = 10) : m_data(mat), dim(dim), npts(npts) 135 | { 136 | assert(npts != 0); 137 | index = new index_t( static_cast(dim), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 138 | index->buildIndex(); 139 | } 140 | 141 | ~KDTreeTableAdaptor() { 142 | delete index; 143 | } 144 | 145 | 146 | /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 147 | * Note that this is a short-cut method for index->findNeighbors(). 148 | * The user can also call index->... methods as desired. 149 | * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 150 | */ 151 | inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 152 | { 153 | nanoflann::KNNResultSet resultSet(num_closest); 154 | resultSet.init(out_indices, out_distances_sq); 155 | index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 156 | } 157 | 158 | /** @name Interface expected by KDTreeSingleIndexAdaptor 159 | * @{ */ 160 | 161 | const self_t & derived() const { 162 | return *this; 163 | } 164 | self_t & derived() { 165 | return *this; 166 | } 167 | 168 | // Must return the number of data points 169 | inline size_t kdtree_get_point_count() const { 170 | return this->npts; 171 | } 172 | 173 | // Returns the dim'th component of the idx'th point in the class: 174 | inline num_t kdtree_get_pt(const size_t pts_id, const size_t coord_id) const { 175 | return m_data[pts_id*this->dim + coord_id]; 176 | } 177 | 178 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 179 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 180 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 181 | template 182 | bool kdtree_get_bbox(BBOX & /*bb*/) const { 183 | return false; 184 | } 185 | 186 | /** @} */ 187 | 188 | }; // end of KDTreeVectorOfVectorsAdaptor 189 | 190 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![arXiv](https://img.shields.io/badge/arXiv-2009.03137-b31b1b.svg)](https://arxiv.org/abs/2009.03137) 2 | [![GitHub Stars](https://img.shields.io/github/stars/QingyongHu/SensatUrban?style=social)](https://github.com/QingyongHu/SensatUrban) 3 | ![visitors](https://visitor-badge.glitch.me/badge?page_id=QingyongHu/SensatUrban) 4 | 5 | # Towards Semantic Segmentation of Urban-Scale 3D Point Clouds: A Dataset, Benchmarks and Challenges 6 | 7 | This is the official repository of the **SensatUrban** dataset. For technical details, please refer to: 8 | 9 | **Towards Semantic Segmentation of Urban-Scale 3D Point Clouds: A Dataset, Benchmarks and Challenges**
10 | [Qingyong Hu](https://qingyonghu.github.io/), [Bo Yang*](https://yang7879.github.io/), [Sheikh Khalid](https://uk.linkedin.com/in/fakharkhalid), 11 | [Wen Xiao](https://www.ncl.ac.uk/engineering/staff/profile/wenxiao.html), [Niki Trigoni](https://www.cs.ox.ac.uk/people/niki.trigoni/), [Andrew Markham](https://www.cs.ox.ac.uk/people/andrew.markham/).
12 | **[[Paper](http://arxiv.org/abs/2009.03137)] [[Blog](https://zhuanlan.zhihu.com/p/259208850)] [[Video](https://www.youtube.com/watch?v=IG0tTdqB3L8)] [[Project page](https://github.com/QingyongHu/SensatUrban)] [[Download](https://forms.gle/m4HJiqZxnq8rmjc8A)] 13 | [[Evaluation](https://competitions.codalab.org/competitions/31519#participate-submit_results)] 14 | [[Urban3D workshop@ECCV2022](https://urban3dchallenge.github.io/)]**
15 | 16 | **We are hosting the 2nd Urban3D Challenge at ECCV 2022, please refer to [this page](https://urban3dchallenge.github.io/) for more details.** 17 | 18 | 19 | ### (1) Dataset 20 | 21 | #### 1.1 Overview 22 | 23 | This dataset is an urban-scale photogrammetric point cloud dataset with nearly three billion richly annotated points, 24 | which is five times the number of labeled points than the existing largest point cloud dataset. 25 | Our dataset consists of large areas from two UK cities, covering about 6 km^2 of the city landscape. 26 | In the dataset, each 3D point is labeled as one of 13 semantic classes, such as *ground*, *vegetation*, 27 | *car*, *etc.*. 28 | 29 |

30 |

31 | 32 | #### 1.2 Data Collection 33 | 34 | The 3D point clouds are generated from high-quality aerial images captured by a 35 | professional-grade UAV mapping system. In order to fully and evenly cover the survey area, 36 | all flight paths are pre-planned in a grid fashion and automated by the flight control system (e-Motion). 37 | 38 |

39 | 40 | #### 1.3 Semantic Annotations 41 | 42 |

43 | 44 | - Ground: including impervious surfaces, grass, terrain 45 | - Vegetation: including trees, shrubs, hedges, bushes 46 | - Building: including commercial / residential buildings 47 | - Wall: including fence, highway barriers, walls 48 | - Bridge: road bridges 49 | - Parking: parking lots 50 | - Rail: railroad tracks 51 | - Traffic Road: including main streets, highways 52 | - Street Furniture: including benches, poles, lights 53 | - Car: including cars, trucks, HGVs 54 | - Footpath: including walkway, alley 55 | - Bike: bikes / bicyclists 56 | - Water: rivers / water canals 57 | 58 | 59 | #### 1.4 Statistics 60 |

61 | 62 | 63 | ### (2) Benchmarks 64 | We extensively evaluate the performance of state-of-the-art algorithms on our dataset 65 | and provide a comprehensive analysis of the results. In particular, we identify several key challenges 66 | towards urban-scale point cloud understanding. 67 | 68 |

69 | 70 | 71 | ### (3) Demo 72 | 73 |

74 | 75 | 76 | ### (4) Training and Evaluation 77 | Here we provide the training and evaluation script of [RandLA-Net](https://github.com/QingyongHu/RandLA-Net) for your reference. 78 | - Download the dataset 79 | 80 | Download the files named "data_release.zip" [here](https://forms.gle/m4HJiqZxnq8rmjc8A). Uncompress the folder and move it to `/Dataset/SensatUrban`. 81 | 82 | - Setup the environment 83 | ``` 84 | conda create -n randlanet python=3.5 85 | source activate randlanet 86 | pip install -r helper_requirements.txt 87 | sh compile_op.sh 88 | ``` 89 | 90 | - Preparing the dataset 91 | ``` 92 | python input_preparation.py --dataset_path $YOURPATH 93 | cd $YOURPATH; 94 | cd ../; mkdir original_block_ply; mv data_release/train/* original_block_ply; mv data_release/test/* original_block_ply; 95 | mv data_release/grid* ./ 96 | ``` 97 | The data should organized in the following format: 98 | ``` 99 | /Dataset/SensatUrban/ 100 | └── original_block_ply/ 101 | ├── birmingham_block_0.ply 102 | ├── birmingham_block_1.ply  103 | ... 104 | └── cambridge_block_34.ply  105 | └── grid_0.200/ 106 | ├── birmingham_block_0_KDTree.pkl 107 | ├── birmingham_block_0.ply 108 | ├── birmingham_block_0_proj.pkl 109 | ... 110 | └── cambridge_block_34.ply  111 | ``` 112 | 113 | - Start training: (Please first modified the root_path) 114 | ``` 115 | python main_SensatUrban.py --mode train --gpu 0 116 | ``` 117 | - Evaluation: 118 | ``` 119 | python main_SensatUrban.py --mode test --gpu 0 120 | ``` 121 | - Submit the results to the server: 122 | The compressed results can be found in `/test/Log_*/test_preds/submission.zip`. Then, feel free to submit this results to the 123 | [evaluation server](https://competitions.codalab.org/competitions/31519#participate-submit_results). 124 | 125 | - The Urban3D Challenge@ICCV2021 Forum: 126 | Please scan the code to join our wechat group or drop a message [here](https://competitions.codalab.org/forums/28215/): 127 |

128 | 129 | 130 | ### Citation 131 | If you find our work useful in your research, please consider citing: 132 | 133 | @inproceedings{hu2020towards, 134 | title={Towards Semantic Segmentation of Urban-Scale 3D Point Clouds: A Dataset, Benchmarks and Challenges}, 135 | author={Hu, Qingyong and Yang, Bo and Khalid, Sheikh and Xiao, Wen and Trigoni, Niki and Markham, Andrew}, 136 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 137 | year={2021} 138 | } 139 | 140 | 141 | @article{hu2022sensaturban, 142 | title={Sensaturban: Learning semantics from urban-scale photogrammetric point clouds}, 143 | author={Hu, Qingyong and Yang, Bo and Khalid, Sheikh and Xiao, Wen and Trigoni, Niki and Markham, Andrew}, 144 | journal={International Journal of Computer Vision}, 145 | volume={130}, 146 | number={2}, 147 | pages={316--343}, 148 | year={2022}, 149 | publisher={Springer} 150 | } 151 | 152 | 153 | ### Updates 154 | * 03/25/2022: we are organizing the [Urban3D@ECCV2022 - The 2nd Challenge on Large-Scale Point Clouds Analysis for Urban Scenes Understanding](https://urban3dchallenge.github.io/)! 155 | * 01/03/2021: The SensatUrban has been accepted by CVPR 2021! 156 | * 11/02/2021: The dataset is available for download! 157 | * 07/09/2020: Initial release! 158 | 159 | 160 | 161 | ## Related Repos 162 | 1. [RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds](https://github.com/QingyongHu/RandLA-Net) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/RandLA-Net.svg?style=flat&label=Star) 163 | 2. [SoTA-Point-Cloud: Deep Learning for 3D Point Clouds: A Survey](https://github.com/QingyongHu/SoTA-Point-Cloud) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SoTA-Point-Cloud.svg?style=flat&label=Star) 164 | 3. [3D-BoNet: Learning Object Bounding Boxes for 3D Instance Segmentation on Point Clouds](https://github.com/Yang7879/3D-BoNet) ![GitHub stars](https://img.shields.io/github/stars/Yang7879/3D-BoNet.svg?style=flat&label=Star) 165 | 4. [SpinNet: Learning a General Surface Descriptor for 3D Point Cloud Registration](https://github.com/QingyongHu/SpinNet) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SpinNet.svg?style=flat&label=Star) 166 | 5. [SQN: Weakly-Supervised Semantic Segmentation of Large-Scale 3D Point Clouds](https://github.com/QingyongHu/SQN) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SQN.svg?style=flat&label=Star) 167 | 6. [Not All Points Are Equal: Learning Highly Efficient Point-based Detectors for 3D LiDAR Point Clouds](https://github.com/yifanzhang713/IA-SSD) ![GitHub stars](https://img.shields.io/github/stars/yifanzhang713/IA-SSD.svg?style=flat&label=Star) 168 | 7. [STPLS3D: A Large-Scale Synthetic and Real Aerial Photogrammetry 3D Point Cloud Dataset](https://github.com/meidachen/STPLS3D) ![GitHub stars](https://img.shields.io/github/stars/meidachen/STPLS3D.svg?style=flat&label=Star) 169 | 170 | 171 | 172 | 173 | 174 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "grid_subsampling/grid_subsampling.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 12 | 13 | static char compute_docstring[] = "function subsampling a pointcloud"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "grid_subsampling", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Actual wrapper 56 | // ************** 57 | 58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject *points_obj = NULL; 66 | PyObject *features_obj = NULL; 67 | PyObject *classes_obj = NULL; 68 | 69 | // Keywords containers 70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 71 | float sampleDl = 0.1; 72 | const char *method_buffer = "barycenters"; 73 | int verbose = 0; 74 | 75 | // Parse the input 76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 77 | { 78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 79 | return NULL; 80 | } 81 | 82 | // Get the method argument 83 | string method(method_buffer); 84 | 85 | // Interpret method 86 | if (method.compare("barycenters") && method.compare("voxelcenters")) 87 | { 88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 89 | return NULL; 90 | } 91 | 92 | // Check if using features or classes 93 | bool use_feature = true, use_classes = true; 94 | if (features_obj == NULL) 95 | use_feature = false; 96 | if (classes_obj == NULL) 97 | use_classes = false; 98 | 99 | // Interpret the input objects as numpy arrays. 100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 101 | PyObject *features_array = NULL; 102 | PyObject *classes_array = NULL; 103 | if (use_feature) 104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 105 | if (use_classes) 106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 107 | 108 | // Verify data was load correctly. 109 | if (points_array == NULL) 110 | { 111 | Py_XDECREF(points_array); 112 | Py_XDECREF(classes_array); 113 | Py_XDECREF(features_array); 114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 115 | return NULL; 116 | } 117 | if (use_feature && features_array == NULL) 118 | { 119 | Py_XDECREF(points_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (use_classes && classes_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(classes_array); 129 | Py_XDECREF(features_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 136 | { 137 | Py_XDECREF(points_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 141 | return NULL; 142 | } 143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(classes_array); 147 | Py_XDECREF(features_array); 148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 149 | return NULL; 150 | } 151 | 152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 153 | { 154 | Py_XDECREF(points_array); 155 | Py_XDECREF(classes_array); 156 | Py_XDECREF(features_array); 157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 158 | return NULL; 159 | } 160 | 161 | // Number of points 162 | int N = (int)PyArray_DIM(points_array, 0); 163 | 164 | // Dimension of the features 165 | int fdim = 0; 166 | if (use_feature) 167 | fdim = (int)PyArray_DIM(features_array, 1); 168 | 169 | //Dimension of labels 170 | int ldim = 1; 171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 172 | ldim = (int)PyArray_DIM(classes_array, 1); 173 | 174 | // Check that the input array respect the number of points 175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 176 | { 177 | Py_XDECREF(points_array); 178 | Py_XDECREF(classes_array); 179 | Py_XDECREF(features_array); 180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 181 | return NULL; 182 | } 183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 184 | { 185 | Py_XDECREF(points_array); 186 | Py_XDECREF(classes_array); 187 | Py_XDECREF(features_array); 188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 189 | return NULL; 190 | } 191 | 192 | 193 | // Call the C++ function 194 | // ********************* 195 | 196 | // Create pyramid 197 | if (verbose > 0) 198 | cout << "Computing cloud pyramid with support points: " << endl; 199 | 200 | 201 | // Convert PyArray to Cloud C++ class 202 | vector original_points; 203 | vector original_features; 204 | vector original_classes; 205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 206 | if (use_feature) 207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 208 | if (use_classes) 209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 210 | 211 | // Subsample 212 | vector subsampled_points; 213 | vector subsampled_features; 214 | vector subsampled_classes; 215 | grid_subsampling(original_points, 216 | subsampled_points, 217 | original_features, 218 | subsampled_features, 219 | original_classes, 220 | subsampled_classes, 221 | sampleDl, 222 | verbose); 223 | 224 | // Check result 225 | if (subsampled_points.size() < 1) 226 | { 227 | PyErr_SetString(PyExc_RuntimeError, "Error"); 228 | return NULL; 229 | } 230 | 231 | // Manage outputs 232 | // ************** 233 | 234 | // Dimension of input containers 235 | npy_intp* point_dims = new npy_intp[2]; 236 | point_dims[0] = subsampled_points.size(); 237 | point_dims[1] = 3; 238 | npy_intp* feature_dims = new npy_intp[2]; 239 | feature_dims[0] = subsampled_points.size(); 240 | feature_dims[1] = fdim; 241 | npy_intp* classes_dims = new npy_intp[2]; 242 | classes_dims[0] = subsampled_points.size(); 243 | classes_dims[1] = ldim; 244 | 245 | // Create output array 246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 247 | PyObject *res_features_obj = NULL; 248 | PyObject *res_classes_obj = NULL; 249 | PyObject *ret = NULL; 250 | 251 | // Fill output array with values 252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 254 | if (use_feature) 255 | { 256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 259 | } 260 | if (use_classes) 261 | { 262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 265 | } 266 | 267 | 268 | // Merge results 269 | if (use_feature && use_classes) 270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 271 | else if (use_feature) 272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 273 | else if (use_classes) 274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 275 | else 276 | ret = Py_BuildValue("N", res_points_obj); 277 | 278 | // Clean up 279 | // ******** 280 | 281 | Py_DECREF(points_array); 282 | Py_XDECREF(features_array); 283 | Py_XDECREF(classes_array); 284 | 285 | return ret; 286 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "grid_subsampling/grid_subsampling.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 12 | 13 | static char compute_docstring[] = "function subsampling a pointcloud"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "grid_subsampling", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Actual wrapper 56 | // ************** 57 | 58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject *points_obj = NULL; 66 | PyObject *features_obj = NULL; 67 | PyObject *classes_obj = NULL; 68 | 69 | // Keywords containers 70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 71 | float sampleDl = 0.1; 72 | const char *method_buffer = "barycenters"; 73 | int verbose = 0; 74 | 75 | // Parse the input 76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 77 | { 78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 79 | return NULL; 80 | } 81 | 82 | // Get the method argument 83 | string method(method_buffer); 84 | 85 | // Interpret method 86 | if (method.compare("barycenters") && method.compare("voxelcenters")) 87 | { 88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 89 | return NULL; 90 | } 91 | 92 | // Check if using features or classes 93 | bool use_feature = true, use_classes = true; 94 | if (features_obj == NULL) 95 | use_feature = false; 96 | if (classes_obj == NULL) 97 | use_classes = false; 98 | 99 | // Interpret the input objects as numpy arrays. 100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 101 | PyObject *features_array = NULL; 102 | PyObject *classes_array = NULL; 103 | if (use_feature) 104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 105 | if (use_classes) 106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 107 | 108 | // Verify data was load correctly. 109 | if (points_array == NULL) 110 | { 111 | Py_XDECREF(points_array); 112 | Py_XDECREF(classes_array); 113 | Py_XDECREF(features_array); 114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 115 | return NULL; 116 | } 117 | if (use_feature && features_array == NULL) 118 | { 119 | Py_XDECREF(points_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (use_classes && classes_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(classes_array); 129 | Py_XDECREF(features_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 136 | { 137 | Py_XDECREF(points_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 141 | return NULL; 142 | } 143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(classes_array); 147 | Py_XDECREF(features_array); 148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 149 | return NULL; 150 | } 151 | 152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 153 | { 154 | Py_XDECREF(points_array); 155 | Py_XDECREF(classes_array); 156 | Py_XDECREF(features_array); 157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 158 | return NULL; 159 | } 160 | 161 | // Number of points 162 | int N = (int)PyArray_DIM(points_array, 0); 163 | 164 | // Dimension of the features 165 | int fdim = 0; 166 | if (use_feature) 167 | fdim = (int)PyArray_DIM(features_array, 1); 168 | 169 | //Dimension of labels 170 | int ldim = 1; 171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 172 | ldim = (int)PyArray_DIM(classes_array, 1); 173 | 174 | // Check that the input array respect the number of points 175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 176 | { 177 | Py_XDECREF(points_array); 178 | Py_XDECREF(classes_array); 179 | Py_XDECREF(features_array); 180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 181 | return NULL; 182 | } 183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 184 | { 185 | Py_XDECREF(points_array); 186 | Py_XDECREF(classes_array); 187 | Py_XDECREF(features_array); 188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 189 | return NULL; 190 | } 191 | 192 | 193 | // Call the C++ function 194 | // ********************* 195 | 196 | // Create pyramid 197 | if (verbose > 0) 198 | cout << "Computing cloud pyramid with support points: " << endl; 199 | 200 | 201 | // Convert PyArray to Cloud C++ class 202 | vector original_points; 203 | vector original_features; 204 | vector original_classes; 205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 206 | if (use_feature) 207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 208 | if (use_classes) 209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 210 | 211 | // Subsample 212 | vector subsampled_points; 213 | vector subsampled_features; 214 | vector subsampled_classes; 215 | grid_subsampling(original_points, 216 | subsampled_points, 217 | original_features, 218 | subsampled_features, 219 | original_classes, 220 | subsampled_classes, 221 | sampleDl, 222 | verbose); 223 | 224 | // Check result 225 | if (subsampled_points.size() < 1) 226 | { 227 | PyErr_SetString(PyExc_RuntimeError, "Error"); 228 | return NULL; 229 | } 230 | 231 | // Manage outputs 232 | // ************** 233 | 234 | // Dimension of input containers 235 | npy_intp* point_dims = new npy_intp[2]; 236 | point_dims[0] = subsampled_points.size(); 237 | point_dims[1] = 3; 238 | npy_intp* feature_dims = new npy_intp[2]; 239 | feature_dims[0] = subsampled_points.size(); 240 | feature_dims[1] = fdim; 241 | npy_intp* classes_dims = new npy_intp[2]; 242 | classes_dims[0] = subsampled_points.size(); 243 | classes_dims[1] = ldim; 244 | 245 | // Create output array 246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 247 | PyObject *res_features_obj = NULL; 248 | PyObject *res_classes_obj = NULL; 249 | PyObject *ret = NULL; 250 | 251 | // Fill output array with values 252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 254 | if (use_feature) 255 | { 256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 259 | } 260 | if (use_classes) 261 | { 262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 265 | } 266 | 267 | 268 | // Merge results 269 | if (use_feature && use_classes) 270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 271 | else if (use_feature) 272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 273 | else if (use_classes) 274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 275 | else 276 | ret = Py_BuildValue("N", res_points_obj); 277 | 278 | // Clean up 279 | // ******** 280 | 281 | Py_DECREF(points_array); 282 | Py_XDECREF(features_array); 283 | Py_XDECREF(classes_array); 284 | 285 | return ret; 286 | } -------------------------------------------------------------------------------- /helper_ply.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0===============================0 4 | # | PLY files reader/writer | 5 | # 0===============================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # function to read/write .ply files 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 10/02/2017 15 | # 16 | 17 | 18 | # ---------------------------------------------------------------------------------------------------------------------- 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Basic libs 26 | import numpy as np 27 | import sys 28 | 29 | 30 | # Define PLY types 31 | ply_dtypes = dict([ 32 | (b'int8', 'i1'), 33 | (b'char', 'i1'), 34 | (b'uint8', 'u1'), 35 | (b'uchar', 'u1'), 36 | (b'int16', 'i2'), 37 | (b'short', 'i2'), 38 | (b'uint16', 'u2'), 39 | (b'ushort', 'u2'), 40 | (b'int32', 'i4'), 41 | (b'int', 'i4'), 42 | (b'uint32', 'u4'), 43 | (b'uint', 'u4'), 44 | (b'float32', 'f4'), 45 | (b'float', 'f4'), 46 | (b'float64', 'f8'), 47 | (b'double', 'f8') 48 | ]) 49 | 50 | # Numpy reader format 51 | valid_formats = {'ascii': '', 'binary_big_endian': '>', 52 | 'binary_little_endian': '<'} 53 | 54 | 55 | # ---------------------------------------------------------------------------------------------------------------------- 56 | # 57 | # Functions 58 | # \***************/ 59 | # 60 | 61 | 62 | def parse_header(plyfile, ext): 63 | # Variables 64 | line = [] 65 | properties = [] 66 | num_points = None 67 | 68 | while b'end_header' not in line and line != b'': 69 | line = plyfile.readline() 70 | 71 | if b'element' in line: 72 | line = line.split() 73 | num_points = int(line[2]) 74 | 75 | elif b'property' in line: 76 | line = line.split() 77 | properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 78 | 79 | return num_points, properties 80 | 81 | 82 | def parse_mesh_header(plyfile, ext): 83 | # Variables 84 | line = [] 85 | vertex_properties = [] 86 | num_points = None 87 | num_faces = None 88 | current_element = None 89 | 90 | 91 | while b'end_header' not in line and line != b'': 92 | line = plyfile.readline() 93 | 94 | # Find point element 95 | if b'element vertex' in line: 96 | current_element = 'vertex' 97 | line = line.split() 98 | num_points = int(line[2]) 99 | 100 | elif b'element face' in line: 101 | current_element = 'face' 102 | line = line.split() 103 | num_faces = int(line[2]) 104 | 105 | elif b'property' in line: 106 | if current_element == 'vertex': 107 | line = line.split() 108 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 109 | elif current_element == 'vertex': 110 | if not line.startswith('property list uchar int'): 111 | raise ValueError('Unsupported faces property : ' + line) 112 | 113 | return num_points, num_faces, vertex_properties 114 | 115 | 116 | def read_ply(filename, triangular_mesh=False): 117 | """ 118 | Read ".ply" files 119 | 120 | Parameters 121 | ---------- 122 | filename : string 123 | the name of the file to read. 124 | 125 | Returns 126 | ------- 127 | result : array 128 | data stored in the file 129 | 130 | Examples 131 | -------- 132 | Store data in file 133 | 134 | >>> points = np.random.rand(5, 3) 135 | >>> values = np.random.randint(2, size=10) 136 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 137 | 138 | Read the file 139 | 140 | >>> data = read_ply('example.ply') 141 | >>> values = data['values'] 142 | array([0, 0, 1, 1, 0]) 143 | 144 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 145 | array([[ 0.466 0.595 0.324] 146 | [ 0.538 0.407 0.654] 147 | [ 0.850 0.018 0.988] 148 | [ 0.395 0.394 0.363] 149 | [ 0.873 0.996 0.092]]) 150 | 151 | """ 152 | 153 | with open(filename, 'rb') as plyfile: 154 | 155 | 156 | # Check if the file start with ply 157 | if b'ply' not in plyfile.readline(): 158 | raise ValueError('The file does not start whith the word ply') 159 | 160 | # get binary_little/big or ascii 161 | fmt = plyfile.readline().split()[1].decode() 162 | if fmt == "ascii": 163 | raise ValueError('The file is not binary') 164 | 165 | # get extension for building the numpy dtypes 166 | ext = valid_formats[fmt] 167 | 168 | # PointCloud reader vs mesh reader 169 | if triangular_mesh: 170 | 171 | # Parse header 172 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 173 | 174 | # Get point data 175 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 176 | 177 | # Get face data 178 | face_properties = [('k', ext + 'u1'), 179 | ('v1', ext + 'i4'), 180 | ('v2', ext + 'i4'), 181 | ('v3', ext + 'i4')] 182 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 183 | 184 | # Return vertex data and concatenated faces 185 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 186 | data = [vertex_data, faces] 187 | 188 | else: 189 | 190 | # Parse header 191 | num_points, properties = parse_header(plyfile, ext) 192 | 193 | # Get data 194 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 195 | 196 | return data 197 | 198 | 199 | def header_properties(field_list, field_names): 200 | 201 | # List of lines to write 202 | lines = [] 203 | 204 | # First line describing element vertex 205 | lines.append('element vertex %d' % field_list[0].shape[0]) 206 | 207 | # Properties lines 208 | i = 0 209 | for fields in field_list: 210 | for field in fields.T: 211 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 212 | i += 1 213 | 214 | return lines 215 | 216 | 217 | def write_ply(filename, field_list, field_names, triangular_faces=None): 218 | """ 219 | Write ".ply" files 220 | 221 | Parameters 222 | ---------- 223 | filename : string 224 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 225 | file name if it does no already have one. 226 | 227 | field_list : list, tuple, numpy array 228 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 229 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 230 | as one field. 231 | 232 | field_names : list 233 | the name of each fields as a list of strings. Has to be the same length as the number of 234 | fields. 235 | 236 | Examples 237 | -------- 238 | >>> points = np.random.rand(10, 3) 239 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 240 | 241 | >>> values = np.random.randint(2, size=10) 242 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 243 | 244 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 245 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values'] 246 | >>> write_ply('example3.ply', [points, colors, values], field_names) 247 | 248 | """ 249 | 250 | # Format list input to the right form 251 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 252 | for i, field in enumerate(field_list): 253 | if field.ndim < 2: 254 | field_list[i] = field.reshape(-1, 1) 255 | if field.ndim > 2: 256 | print('fields have more than 2 dimensions') 257 | return False 258 | 259 | # check all fields have the same number of data 260 | n_points = [field.shape[0] for field in field_list] 261 | if not np.all(np.equal(n_points, n_points[0])): 262 | print('wrong field dimensions') 263 | return False 264 | 265 | # Check if field_names and field_list have same nb of column 266 | n_fields = np.sum([field.shape[1] for field in field_list]) 267 | if (n_fields != len(field_names)): 268 | print('wrong number of field names') 269 | return False 270 | 271 | # Add extension if not there 272 | if not filename.endswith('.ply'): 273 | filename += '.ply' 274 | 275 | # open in text mode to write the header 276 | with open(filename, 'w') as plyfile: 277 | 278 | # First magical word 279 | header = ['ply'] 280 | 281 | # Encoding format 282 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 283 | 284 | # Points properties description 285 | header.extend(header_properties(field_list, field_names)) 286 | 287 | # Add faces if needded 288 | if triangular_faces is not None: 289 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 290 | header.append('property list uchar int vertex_indices') 291 | 292 | # End of header 293 | header.append('end_header') 294 | 295 | # Write all lines 296 | for line in header: 297 | plyfile.write("%s\n" % line) 298 | 299 | # open in binary/append to use tofile 300 | with open(filename, 'ab') as plyfile: 301 | 302 | # Create a structured array 303 | i = 0 304 | type_list = [] 305 | for fields in field_list: 306 | for field in fields.T: 307 | type_list += [(field_names[i], field.dtype.str)] 308 | i += 1 309 | data = np.empty(field_list[0].shape[0], dtype=type_list) 310 | i = 0 311 | for fields in field_list: 312 | for field in fields.T: 313 | data[field_names[i]] = field 314 | i += 1 315 | 316 | data.tofile(plyfile) 317 | 318 | if triangular_faces is not None: 319 | triangular_faces = triangular_faces.astype(np.int32) 320 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 321 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 322 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 323 | data['0'] = triangular_faces[:, 0] 324 | data['1'] = triangular_faces[:, 1] 325 | data['2'] = triangular_faces[:, 2] 326 | data.tofile(plyfile) 327 | 328 | return True 329 | 330 | 331 | def describe_element(name, df): 332 | """ Takes the columns of the dataframe and builds a ply-like description 333 | 334 | Parameters 335 | ---------- 336 | name: str 337 | df: pandas DataFrame 338 | 339 | Returns 340 | ------- 341 | element: list[str] 342 | """ 343 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 344 | element = ['element ' + name + ' ' + str(len(df))] 345 | 346 | if name == 'face': 347 | element.append("property list uchar int points_indices") 348 | 349 | else: 350 | for i in range(len(df.columns)): 351 | # get first letter of dtype to infer format 352 | f = property_formats[str(df.dtypes[i])[0]] 353 | element.append('property ' + f + ' ' + df.columns.values[i]) 354 | 355 | return element 356 | 357 | -------------------------------------------------------------------------------- /tool.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname, abspath 2 | import numpy as np 3 | import colorsys, random, os, sys 4 | import open3d as o3d 5 | from helper_ply import read_ply, write_ply 6 | 7 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | 11 | sys.path.append(BASE_DIR) 12 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 13 | 14 | import cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 15 | import nearest_neighbors.lib.python.nearest_neighbors as nearest_neighbors 16 | 17 | 18 | class ConfigSensatUrban: 19 | k_n = 16 # KNN 20 | num_layers = 5 # Number of layers 21 | num_points = 65536 # Number of input points 22 | num_classes = 13 # Number of valid classes 23 | sub_grid_size = 0.2 # preprocess_parameter 24 | 25 | batch_size = 4 # batch_size during training 26 | val_batch_size = 14 # batch_size during validation and test 27 | train_steps = 500 # Number of steps per epochs 28 | val_steps = 100 # Number of validation steps per epoch 29 | 30 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 31 | d_out = [16, 64, 128, 256, 512] # feature dimension 32 | 33 | noise_init = 3.5 # noise initial parameter 34 | max_epoch = 100 # maximum epoch during training 35 | learning_rate = 1e-2 # initial learning rate 36 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 37 | 38 | train_sum_dir = 'train_log_SensatUrban' 39 | saving = True 40 | saving_path = None 41 | 42 | 43 | class DataProcessing: 44 | 45 | @staticmethod 46 | def get_num_class_from_label(labels, total_class): 47 | num_pts_per_class = np.zeros(total_class, dtype=np.int32) 48 | # original class distribution 49 | val_list, counts = np.unique(labels, return_counts=True) 50 | for idx, val in enumerate(val_list): 51 | num_pts_per_class[val] += counts[idx] 52 | # for idx, nums in enumerate(num_pts_per_class): 53 | # print(idx, ':', nums) 54 | return num_pts_per_class 55 | 56 | @staticmethod 57 | def knn_search(support_pts, query_pts, k): 58 | """ 59 | :param support_pts: points you have, B*N1*3 60 | :param query_pts: points you want to know the neighbour index, B*N2*3 61 | :param k: Number of neighbours in knn search 62 | :return: neighbor_idx: neighboring points indexes, B*N2*k 63 | """ 64 | 65 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 66 | return neighbor_idx.astype(np.int32) 67 | 68 | @staticmethod 69 | def data_aug(xyz, color, labels, idx, num_out): 70 | num_in = len(xyz) 71 | dup = np.random.choice(num_in, num_out - num_in) 72 | xyz_dup = xyz[dup, ...] 73 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) 74 | color_dup = color[dup, ...] 75 | color_aug = np.concatenate([color, color_dup], 0) 76 | idx_dup = list(range(num_in)) + list(dup) 77 | idx_aug = idx[idx_dup] 78 | label_aug = labels[idx_dup] 79 | return xyz_aug, color_aug, idx_aug, label_aug 80 | 81 | @staticmethod 82 | def shuffle_idx(x): 83 | # random shuffle the index 84 | idx = np.arange(len(x)) 85 | np.random.shuffle(idx) 86 | return x[idx] 87 | 88 | @staticmethod 89 | def shuffle_list(data_list): 90 | indices = np.arange(np.shape(data_list)[0]) 91 | np.random.shuffle(indices) 92 | data_list = data_list[indices] 93 | return data_list 94 | 95 | @staticmethod 96 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 97 | """ 98 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 99 | :param points: (N, 3) matrix of input points 100 | :param features: optional (N, d) matrix of features (floating number) 101 | :param labels: optional (N,) matrix of integer labels 102 | :param grid_size: parameter defining the size of grid voxels 103 | :param verbose: 1 to display 104 | :return: sub_sampled points, with features and/or labels depending of the input 105 | """ 106 | 107 | if (features is None) and (labels is None): 108 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 109 | elif labels is None: 110 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 111 | elif features is None: 112 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 113 | else: 114 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 115 | verbose=verbose) 116 | 117 | @staticmethod 118 | def IoU_from_confusions(confusions): 119 | """ 120 | Computes IoU from confusion matrices. 121 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 122 | the last axes. n_c = number of classes 123 | :return: ([..., n_c] np.float32) IoU score 124 | """ 125 | 126 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 127 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 128 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 129 | TP_plus_FN = np.sum(confusions, axis=-1) 130 | TP_plus_FP = np.sum(confusions, axis=-2) 131 | 132 | # Compute IoU 133 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 134 | 135 | # Compute mIoU with only the actual classes 136 | mask = TP_plus_FN < 1e-3 137 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 138 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 139 | 140 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 141 | IoU += mask * mIoU 142 | return IoU 143 | 144 | @staticmethod 145 | def read_ply_data(path, with_rgb=True, with_label=True): 146 | data = read_ply(path) 147 | xyz = np.vstack((data['x'], data['y'], data['z'])).T 148 | if with_rgb and with_label: 149 | rgb = np.vstack((data['red'], data['green'], data['blue'])).T 150 | labels = data['class'] 151 | return xyz.astype(np.float32), rgb.astype(np.uint8), labels.astype(np.uint8) 152 | elif with_rgb and not with_label: 153 | rgb = np.vstack((data['red'], data['green'], data['blue'])).T 154 | return xyz.astype(np.float32), rgb.astype(np.uint8) 155 | elif not with_rgb and with_label: 156 | labels = data['class'] 157 | return xyz.astype(np.float32), labels.astype(np.uint8) 158 | elif not with_rgb and not with_label: 159 | return xyz.astype(np.float32) 160 | 161 | @staticmethod 162 | def random_sub_sampling(points, features=None, labels=None, sub_ratio=10, verbose=0): 163 | num_input = np.shape(points)[0] 164 | num_output = num_input // sub_ratio 165 | idx = np.random.choice(num_input, num_output) 166 | if (features is None) and (labels is None): 167 | return points[idx] 168 | elif labels is None: 169 | return points[idx], features[idx] 170 | elif features is None: 171 | return points[idx], labels[idx] 172 | else: 173 | return points[idx], features[idx], labels[idx] 174 | 175 | @staticmethod 176 | def get_class_weights(num_per_class, name='sqrt'): 177 | # # pre-calculate the number of points in each category 178 | frequency = num_per_class / float(sum(num_per_class)) 179 | if name == 'sqrt' or name == 'lovas': 180 | ce_label_weight = 1 / np.sqrt(frequency) 181 | elif name == 'wce': 182 | ce_label_weight = 1 / (frequency + 0.02) 183 | else: 184 | raise ValueError('Only support sqrt and wce') 185 | return np.expand_dims(ce_label_weight, axis=0) 186 | 187 | 188 | class Plot: 189 | @staticmethod 190 | def random_colors(N, bright=True, seed=0): 191 | brightness = 1.0 if bright else 0.7 192 | hsv = [(0.15 + i / float(N), 1, brightness) for i in range(N)] 193 | colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)) 194 | random.seed(seed) 195 | random.shuffle(colors) 196 | return colors 197 | 198 | @staticmethod 199 | def draw_pc(pc_xyzrgb): 200 | pc = o3d.geometry.PointCloud() 201 | pc.points = o3d.utility.Vector3dVector(pc_xyzrgb[:, 0:3]) 202 | if pc_xyzrgb.shape[1] == 3: 203 | o3d.visualization.draw_geometries([pc]) 204 | return 0 205 | if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255 206 | pc.colors = o3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.) 207 | else: 208 | pc.colors = o3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6]) 209 | 210 | o3d.geometry.PointCloud.estimate_normals(pc) 211 | o3d.visualization.draw_geometries([pc], width=1000, height=1000) 212 | return 0 213 | 214 | @staticmethod 215 | def draw_pc_sem_ins(pc_xyz, pc_sem_ins, plot_colors=None): 216 | # only visualize a number of points to save memory 217 | if plot_colors is not None: 218 | ins_colors = plot_colors 219 | else: 220 | # ins_colors = Plot.random_colors(len(np.unique(pc_sem_ins)) + 1, seed=1) 221 | ins_colors = [[85, 107, 47], # ground -> OliveDrab 222 | [0, 255, 0], # tree -> Green 223 | [255, 165, 0], # building -> orange 224 | [41, 49, 101], # Walls -> darkblue 225 | [0, 0, 0], # Bridge -> black 226 | [0, 0, 255], # parking -> blue 227 | [255, 0, 255], # rail -> Magenta 228 | [200, 200, 200], # traffic Roads -> grey 229 | [89, 47, 95], # Street Furniture -> DimGray 230 | [255, 0, 0], # cars -> red 231 | [255, 255, 0], # Footpath -> deeppink 232 | [0, 255, 255], # bikes -> cyan 233 | [0, 191, 255] # water -> skyblue 234 | ] 235 | 236 | ############################## 237 | sem_ins_labels = np.unique(pc_sem_ins) 238 | sem_ins_bbox = [] 239 | Y_colors = np.zeros((pc_sem_ins.shape[0], 3)) 240 | for id, semins in enumerate(sem_ins_labels): 241 | valid_ind = np.argwhere(pc_sem_ins == semins)[:, 0] 242 | if semins <= -1: 243 | tp = [0, 0, 0] 244 | else: 245 | if plot_colors is not None: 246 | tp = ins_colors[semins] 247 | else: 248 | tp = ins_colors[id] 249 | 250 | Y_colors[valid_ind] = tp 251 | 252 | ### bbox 253 | valid_xyz = pc_xyz[valid_ind] 254 | 255 | xmin = np.min(valid_xyz[:, 0]); 256 | xmax = np.max(valid_xyz[:, 0]) 257 | ymin = np.min(valid_xyz[:, 1]); 258 | ymax = np.max(valid_xyz[:, 1]) 259 | zmin = np.min(valid_xyz[:, 2]); 260 | zmax = np.max(valid_xyz[:, 2]) 261 | sem_ins_bbox.append( 262 | [[xmin, ymin, zmin], [xmax, ymax, zmax], [min(tp[0], 1.), min(tp[1], 1.), min(tp[2], 1.)]]) 263 | 264 | Y_semins = np.concatenate([pc_xyz[:, 0:3], Y_colors], axis=-1) 265 | Plot.draw_pc(Y_semins) 266 | return Y_semins 267 | 268 | @staticmethod 269 | def save_ply_o3d(data, save_name): 270 | pcd = o3d.geometry.PointCloud() 271 | pcd.points = o3d.utility.Vector3dVector(data[:, 0:3]) 272 | if np.shape(data)[1] == 3: 273 | o3d.io.write_point_cloud(save_name, pcd) 274 | elif np.shape(data)[1] == 6: 275 | if np.max(data[:, 3:6]) > 20: 276 | pcd.colors = o3d.utility.Vector3dVector(data[:, 3:6] / 255.) 277 | else: 278 | pcd.colors = o3d.utility.Vector3dVector(data[:, 3:6]) 279 | o3d.io.write_point_cloud(save_name, pcd) 280 | return 281 | 282 | -------------------------------------------------------------------------------- /main_SensatUrban.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname, abspath 2 | from RandLANet import Network 3 | from tester_SensatUrban import ModelTester 4 | from helper_ply import read_ply 5 | from tool import ConfigSensatUrban as cfg 6 | from tool import DataProcessing as DP 7 | from tool import Plot 8 | import tensorflow as tf 9 | import numpy as np 10 | import time, pickle, argparse, glob, os, shutil 11 | 12 | 13 | class SensatUrban: 14 | def __init__(self): 15 | self.name = 'SensatUrban' 16 | root_path = '/media/qingyong/data/Dataset' # path to the dataset 17 | self.path = join(root_path, self.name) 18 | self.label_to_names = {0: 'Ground', 1: 'High Vegetation', 2: 'Buildings', 3: 'Walls', 19 | 4: 'Bridge', 5: 'Parking', 6: 'Rail', 7: 'traffic Roads', 8: 'Street Furniture', 20 | 9: 'Cars', 10: 'Footpath', 11: 'Bikes', 12: 'Water'} 21 | self.num_classes = len(self.label_to_names) 22 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 23 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 24 | self.ignored_labels = np.array([]) 25 | 26 | self.all_files = np.sort(glob.glob(join(self.path, 'original_block_ply', '*.ply'))) 27 | self.val_file_name = ['birmingham_block_1', 28 | 'birmingham_block_5', 29 | 'cambridge_block_10', 30 | 'cambridge_block_7'] 31 | self.test_file_name = ['birmingham_block_2', 'birmingham_block_8', 32 | 'cambridge_block_15', 'cambridge_block_22', 33 | 'cambridge_block_16', 'cambridge_block_27'] 34 | self.use_val = True # whether use validation set or not 35 | 36 | # initialize 37 | self.num_per_class = np.zeros(self.num_classes) 38 | self.val_proj = [] 39 | self.val_labels = [] 40 | self.test_proj = [] 41 | self.test_labels = [] 42 | self.possibility = {} 43 | self.min_possibility = {} 44 | self.input_trees = {'training': [], 'validation': [], 'test': []} 45 | self.input_colors = {'training': [], 'validation': [], 'test': []} 46 | self.input_labels = {'training': [], 'validation': [], 'test': []} 47 | self.input_names = {'training': [], 'validation': [], 'test': []} 48 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 49 | for ignore_label in self.ignored_labels: 50 | self.num_per_class = np.delete(self.num_per_class, ignore_label) 51 | 52 | def load_sub_sampled_clouds(self, sub_grid_size): 53 | tree_path = join(self.path, 'grid_{:.3f}'.format(sub_grid_size)) 54 | 55 | for i, file_path in enumerate(self.all_files): 56 | t0 = time.time() 57 | cloud_name = file_path.split('/')[-1][:-4] 58 | if cloud_name in self.test_file_name: 59 | cloud_split = 'test' 60 | elif cloud_name in self.val_file_name: 61 | cloud_split = 'validation' 62 | else: 63 | cloud_split = 'training' 64 | 65 | # Name of the input files 66 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 67 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 68 | 69 | data = read_ply(sub_ply_file) 70 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 71 | sub_labels = data['class'] 72 | 73 | # compute num_per_class in training set 74 | if cloud_split == 'training': 75 | self.num_per_class += DP.get_num_class_from_label(sub_labels, self.num_classes) 76 | 77 | # Read pkl with search tree 78 | with open(kd_tree_file, 'rb') as f: 79 | search_tree = pickle.load(f) 80 | 81 | self.input_trees[cloud_split] += [search_tree] 82 | self.input_colors[cloud_split] += [sub_colors] 83 | self.input_labels[cloud_split] += [sub_labels] 84 | self.input_names[cloud_split] += [cloud_name] 85 | 86 | size = sub_colors.shape[0] * 4 * 7 87 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 88 | 89 | print('\nPreparing reprojected indices for testing') 90 | 91 | # Get validation and test reprojected indices 92 | 93 | for i, file_path in enumerate(self.all_files): 94 | t0 = time.time() 95 | cloud_name = file_path.split('/')[-1][:-4] 96 | 97 | # val projection and labels 98 | if cloud_name in self.val_file_name: 99 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 100 | with open(proj_file, 'rb') as f: 101 | proj_idx, labels = pickle.load(f) 102 | self.val_proj += [proj_idx] 103 | self.val_labels += [labels] 104 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 105 | 106 | # test projection and labels 107 | if cloud_name in self.test_file_name: 108 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 109 | with open(proj_file, 'rb') as f: 110 | proj_idx, labels = pickle.load(f) 111 | self.test_proj += [proj_idx] 112 | self.test_labels += [labels] 113 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 114 | 115 | def get_batch_gen(self, split): 116 | if split == 'training': 117 | num_per_epoch = cfg.train_steps * cfg.batch_size 118 | elif split == 'validation': 119 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 120 | else: 121 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 122 | 123 | # Reset possibility 124 | self.possibility[split] = [] 125 | self.min_possibility[split] = [] 126 | for i, tree in enumerate(self.input_colors[split]): 127 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 128 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 129 | 130 | def spatially_regular_gen(): 131 | # Generator loop 132 | for i in range(num_per_epoch): # num_per_epoch 133 | 134 | # Choose a random cloud 135 | cloud_idx = int(np.argmin(self.min_possibility[split])) 136 | 137 | # choose the point with the minimum of possibility as query point 138 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 139 | 140 | # Get points from tree structure 141 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 142 | 143 | # Center point of input region 144 | center_point = points[point_ind, :].reshape(1, -1) 145 | 146 | # Add noise to the center point 147 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 148 | pick_point = center_point + noise.astype(center_point.dtype) 149 | 150 | if len(points) < cfg.num_points: 151 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 152 | else: 153 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 154 | 155 | queried_idx = DP.shuffle_idx(queried_idx) 156 | # Collect points and colors 157 | queried_pc_xyz = points[queried_idx] 158 | queried_pc_xyz = queried_pc_xyz - pick_point 159 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 160 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 161 | 162 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 163 | delta = np.square(1 - dists / np.max(dists)) 164 | self.possibility[split][cloud_idx][queried_idx] += delta 165 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 166 | 167 | if len(points) < cfg.num_points: 168 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 169 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 170 | 171 | if True: 172 | yield (queried_pc_xyz.astype(np.float32), 173 | queried_pc_colors.astype(np.float32), 174 | queried_pc_labels, 175 | queried_idx.astype(np.int32), 176 | np.array([cloud_idx], dtype=np.int32)) 177 | 178 | gen_func = spatially_regular_gen 179 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 180 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 181 | return gen_func, gen_types, gen_shapes 182 | 183 | @staticmethod 184 | def get_tf_mapping2(): 185 | 186 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 187 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 188 | input_points = [] 189 | input_neighbors = [] 190 | input_pools = [] 191 | input_up_samples = [] 192 | 193 | for i in range(cfg.num_layers): 194 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 195 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 196 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 197 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 198 | input_points.append(batch_xyz) 199 | input_neighbors.append(neighbour_idx) 200 | input_pools.append(pool_i) 201 | input_up_samples.append(up_i) 202 | batch_xyz = sub_points 203 | 204 | input_list = input_points + input_neighbors + input_pools + input_up_samples 205 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 206 | 207 | return input_list 208 | 209 | return tf_map 210 | 211 | def init_input_pipeline(self): 212 | print('Initiating input pipelines') 213 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 214 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 215 | gen_function_val, _, _ = self.get_batch_gen('validation') 216 | gen_function_test, _, _ = self.get_batch_gen('test') 217 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 218 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 219 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 220 | 221 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 222 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 223 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 224 | map_func = self.get_tf_mapping2() 225 | 226 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 227 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 228 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 229 | 230 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 231 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 232 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 233 | 234 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 235 | self.flat_inputs = iter.get_next() 236 | self.train_init_op = iter.make_initializer(self.batch_train_data) 237 | self.val_init_op = iter.make_initializer(self.batch_val_data) 238 | self.test_init_op = iter.make_initializer(self.batch_test_data) 239 | 240 | 241 | if __name__ == '__main__': 242 | parser = argparse.ArgumentParser() 243 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 244 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 245 | FLAGS = parser.parse_args() 246 | 247 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 248 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 249 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 250 | Mode = FLAGS.mode 251 | 252 | shutil.rmtree('__pycache__') if exists('__pycache__') else None 253 | if Mode == 'train': 254 | shutil.rmtree('results') if exists('results') else None 255 | shutil.rmtree('train_log') if exists('train_log') else None 256 | for f in os.listdir(dirname(abspath(__file__))): 257 | if f.startswith('log_'): 258 | os.remove(f) 259 | 260 | dataset = SensatUrban() 261 | dataset.init_input_pipeline() 262 | 263 | if Mode == 'train': 264 | model = Network(dataset, cfg) 265 | model.train(dataset) 266 | elif Mode == 'test': 267 | cfg.saving = False 268 | model = Network(dataset, cfg) 269 | chosen_snapshot = -1 270 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 271 | chosen_folder = logs[-1] 272 | snap_path = join(chosen_folder, 'snapshots') 273 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 274 | chosen_step = np.sort(snap_steps)[-1] 275 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 276 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 277 | tester.test(model, dataset) 278 | shutil.rmtree('train_log') if exists('train_log') else None 279 | 280 | else: 281 | 282 | with tf.Session() as sess: 283 | sess.run(tf.global_variables_initializer()) 284 | sess.run(dataset.train_init_op) 285 | while True: 286 | data_list = sess.run(dataset.flat_inputs) 287 | xyz = data_list[0] 288 | sub_xyz = data_list[1] 289 | label = data_list[21] 290 | Plot.draw_pc_sem_ins(xyz[0, :, :], label[0, :]) 291 | -------------------------------------------------------------------------------- /RandLANet.py: -------------------------------------------------------------------------------- 1 | from os.path import exists, join 2 | from os import makedirs 3 | from sklearn.metrics import confusion_matrix 4 | from tool import DataProcessing as DP 5 | import tensorflow as tf 6 | import numpy as np 7 | import tf_util 8 | import time 9 | 10 | 11 | def log_out(out_str, f_out): 12 | f_out.write(out_str + '\n') 13 | f_out.flush() 14 | print(out_str) 15 | 16 | 17 | class Network: 18 | def __init__(self, dataset, config): 19 | flat_inputs = dataset.flat_inputs 20 | self.config = config 21 | # Path of the result folder 22 | if self.config.saving: 23 | if self.config.saving_path is None: 24 | self.saving_path = time.strftime('results/Log_%Y-%m-%d', time.gmtime()) 25 | self.saving_path = self.saving_path + '_' + dataset.name 26 | else: 27 | self.saving_path = self.config.saving_path 28 | makedirs(self.saving_path) if not exists(self.saving_path) else None 29 | 30 | with tf.variable_scope('inputs'): 31 | self.inputs = dict() 32 | num_layers = self.config.num_layers 33 | self.inputs['xyz'] = flat_inputs[:num_layers] 34 | self.inputs['neigh_idx'] = flat_inputs[num_layers: 2 * num_layers] 35 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 36 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 37 | self.inputs['features'] = flat_inputs[4 * num_layers] 38 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 39 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 40 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 41 | 42 | self.labels = self.inputs['labels'] 43 | self.is_training = tf.placeholder(tf.bool, shape=()) 44 | self.training_step = 1 45 | self.training_epoch = 0 46 | self.correct_prediction = 0 47 | self.accuracy = 0 48 | self.mIou_list = [0] 49 | self.loss_type = 'sqrt' # wce, lovas 50 | self.class_weights = DP.get_class_weights(dataset.num_per_class, self.loss_type) 51 | self.Log_file = open('log_train_' + dataset.name + '.txt', 'a') 52 | 53 | with tf.variable_scope('layers'): 54 | self.logits = self.inference(self.inputs, self.is_training) 55 | 56 | with tf.variable_scope('loss'): 57 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 58 | self.labels = tf.reshape(self.labels, [-1]) 59 | 60 | # Boolean mask of points that should be ignored 61 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 62 | for ign_label in self.config.ignored_label_inds: 63 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 64 | 65 | # Collect logits and labels that are not ignored 66 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 67 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 68 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 69 | 70 | # Reduce label values in the range of logit shape 71 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 72 | inserted_value = tf.zeros((1,), dtype=tf.int32) 73 | for ign_label in self.config.ignored_label_inds: 74 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 75 | valid_labels = tf.gather(reducing_list, valid_labels_init) 76 | 77 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 78 | 79 | with tf.variable_scope('optimizer'): 80 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 81 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 82 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 83 | 84 | with tf.variable_scope('results'): 85 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 86 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 87 | self.prob_logits = tf.nn.softmax(self.logits) 88 | 89 | tf.summary.scalar('learning_rate', self.learning_rate) 90 | tf.summary.scalar('loss', self.loss) 91 | tf.summary.scalar('accuracy', self.accuracy) 92 | 93 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 94 | self.saver = tf.train.Saver(my_vars, max_to_keep=1) 95 | c_proto = tf.ConfigProto() 96 | c_proto.gpu_options.allow_growth = True 97 | self.sess = tf.Session(config=c_proto) 98 | self.merged = tf.summary.merge_all() 99 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 100 | self.sess.run(tf.global_variables_initializer()) 101 | 102 | def inference(self, inputs, is_training): 103 | 104 | d_out = self.config.d_out 105 | feature = inputs['features'] 106 | feature = tf.layers.dense(feature, 8, activation=None, name='fc0') 107 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 108 | feature = tf.expand_dims(feature, axis=2) 109 | 110 | # ###########################Encoder############################ 111 | f_encoder_list = [] 112 | for i in range(self.config.num_layers): 113 | f_encoder_i = self.dilated_res_block(feature, inputs['xyz'][i], inputs['neigh_idx'][i], d_out[i], 114 | 'Encoder_layer_' + str(i), is_training) 115 | f_sampled_i = self.random_sample(f_encoder_i, inputs['sub_idx'][i]) 116 | feature = f_sampled_i 117 | if i == 0: 118 | f_encoder_list.append(f_encoder_i) 119 | f_encoder_list.append(f_sampled_i) 120 | # ###########################Encoder############################ 121 | 122 | feature = tf_util.conv2d(f_encoder_list[-1], f_encoder_list[-1].get_shape()[3].value, [1, 1], 123 | 'decoder_0', [1, 1], 'VALID', True, is_training) 124 | 125 | # ###########################Decoder############################ 126 | f_decoder_list = [] 127 | for j in range(self.config.num_layers): 128 | f_interp_i = self.nearest_interpolation(feature, inputs['interp_idx'][-j - 1]) 129 | f_decoder_i = tf_util.conv2d_transpose(tf.concat([f_encoder_list[-j - 2], f_interp_i], axis=3), 130 | f_encoder_list[-j - 2].get_shape()[-1].value, [1, 1], 131 | 'Decoder_layer_' + str(j), [1, 1], 'VALID', bn=True, 132 | is_training=is_training) 133 | feature = f_decoder_i 134 | f_decoder_list.append(f_decoder_i) 135 | # ###########################Decoder############################ 136 | 137 | f_layer_fc1 = tf_util.conv2d(f_decoder_list[-1], 64, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 138 | f_layer_fc2 = tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 139 | f_layer_drop = tf_util.dropout(f_layer_fc2, keep_prob=0.5, is_training=is_training, scope='dp1') 140 | f_layer_fc3 = tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, 141 | is_training, activation_fn=None) 142 | f_out = tf.squeeze(f_layer_fc3, [2]) 143 | return f_out 144 | 145 | def train(self, dataset): 146 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 147 | self.sess.run(dataset.train_init_op) 148 | while self.training_epoch < self.config.max_epoch: 149 | t_start = time.time() 150 | try: 151 | ops = [self.train_op, 152 | self.extra_update_ops, 153 | self.merged, 154 | self.loss, 155 | self.logits, 156 | self.labels, 157 | self.accuracy] 158 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 159 | self.train_writer.add_summary(summary, self.training_step) 160 | t_end = time.time() 161 | if self.training_step % 50 == 0: 162 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 163 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 164 | self.training_step += 1 165 | 166 | except tf.errors.OutOfRangeError: 167 | 168 | if dataset.use_val and self.training_epoch % 2 == 0: 169 | m_iou = self.evaluate(dataset) 170 | if m_iou > np.max(self.mIou_list): 171 | # Save the best model 172 | snapshot_directory = join(self.saving_path, 'snapshots') 173 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 174 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 175 | self.mIou_list.append(m_iou) 176 | log_out('Best m_IoU of {} is: {:5.3f}'.format(dataset.name, max(self.mIou_list)), self.Log_file) 177 | else: 178 | snapshot_directory = join(self.saving_path, 'snapshots') 179 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 180 | self.saver.save(self.sess, snapshot_directory + '/snap', self.training_step) 181 | 182 | self.training_epoch += 1 183 | self.sess.run(dataset.train_init_op) 184 | # Update learning rate 185 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, 186 | self.config.lr_decays[self.training_epoch])) 187 | self.sess.run(op) 188 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 189 | 190 | except tf.errors.InvalidArgumentError as e: 191 | 192 | print('Caught a NaN error :') 193 | print(e.error_code) 194 | print(e.message) 195 | print(e.op) 196 | print(e.op.name) 197 | print([t.name for t in e.op.inputs]) 198 | print([t.name for t in e.op.outputs]) 199 | 200 | a = 1 / 0 201 | 202 | print('finished') 203 | self.sess.close() 204 | 205 | def evaluate(self, dataset): 206 | 207 | # Initialise iterator with validation data 208 | self.sess.run(dataset.val_init_op) 209 | 210 | gt_classes = [0 for _ in range(self.config.num_classes)] 211 | positive_classes = [0 for _ in range(self.config.num_classes)] 212 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 213 | val_total_correct = 0 214 | val_total_seen = 0 215 | 216 | for step_id in range(self.config.val_steps): 217 | if step_id % 50 == 0: 218 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 219 | try: 220 | ops = (self.prob_logits, self.labels, self.accuracy) 221 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 222 | pred = np.argmax(stacked_prob, 1) 223 | if not self.config.ignored_label_inds: 224 | pred_valid = pred 225 | labels_valid = labels 226 | else: 227 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 228 | labels_valid = np.delete(labels, invalid_idx) 229 | labels_valid = labels_valid - 1 230 | pred_valid = np.delete(pred, invalid_idx) 231 | 232 | correct = np.sum(pred_valid == labels_valid) 233 | val_total_correct += correct 234 | val_total_seen += len(labels_valid) 235 | 236 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 237 | gt_classes += np.sum(conf_matrix, axis=1) 238 | positive_classes += np.sum(conf_matrix, axis=0) 239 | true_positive_classes += np.diagonal(conf_matrix) 240 | 241 | except tf.errors.OutOfRangeError: 242 | break 243 | 244 | iou_list = [] 245 | for n in range(0, self.config.num_classes, 1): 246 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n] + 0.1) 247 | iou_list.append(iou) 248 | mean_iou = sum(iou_list) / float(self.config.num_classes) 249 | 250 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 251 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 252 | 253 | mean_iou = 100 * mean_iou 254 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 255 | s = '{:5.2f} | '.format(mean_iou) 256 | for IoU in iou_list: 257 | s += '{:5.2f} '.format(100 * IoU) 258 | log_out('-' * len(s), self.Log_file) 259 | log_out(s, self.Log_file) 260 | log_out('-' * len(s) + '\n', self.Log_file) 261 | return mean_iou 262 | 263 | def get_loss(self, logits, labels, pre_cal_weights): 264 | # calculate the weighted cross entropy according to the inverse frequency 265 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 266 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 267 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 268 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits_v2(logits=logits, labels=one_hot_labels) 269 | weighted_losses = unweighted_losses * weights 270 | output_loss = tf.reduce_mean(weighted_losses) 271 | return output_loss 272 | 273 | def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training): 274 | f_pc = tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 275 | f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 276 | f_pc = tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, 277 | activation_fn=None) 278 | shortcut = tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', 279 | activation_fn=None, bn=True, is_training=is_training) 280 | return tf.nn.leaky_relu(f_pc + shortcut) 281 | 282 | def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training): 283 | d_in = feature.get_shape()[-1].value 284 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) 285 | f_xyz = tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 286 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 287 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 288 | f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training) 289 | 290 | f_xyz = tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 291 | f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx) 292 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 293 | f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training) 294 | return f_pc_agg 295 | 296 | def relative_pos_encoding(self, xyz, neigh_idx): 297 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 298 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 299 | relative_xyz = xyz_tile - neighbor_xyz 300 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 301 | relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1) 302 | return relative_feature 303 | 304 | @staticmethod 305 | def random_sample(feature, pool_idx): 306 | """ 307 | :param feature: [B, N, d] input features matrix 308 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 309 | :return: pool_features = [B, N', d] pooled features matrix 310 | """ 311 | feature = tf.squeeze(feature, axis=2) 312 | num_neigh = tf.shape(pool_idx)[-1] 313 | d = feature.get_shape()[-1] 314 | batch_size = tf.shape(pool_idx)[0] 315 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 316 | pool_features = tf.batch_gather(feature, pool_idx) 317 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 318 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 319 | return pool_features 320 | 321 | @staticmethod 322 | def nearest_interpolation(feature, interp_idx): 323 | """ 324 | :param feature: [B, N, d] input features matrix 325 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 326 | :return: [B, up_num_points, d] interpolated features matrix 327 | """ 328 | feature = tf.squeeze(feature, axis=2) 329 | batch_size = tf.shape(interp_idx)[0] 330 | up_num_points = tf.shape(interp_idx)[1] 331 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 332 | interpolated_features = tf.batch_gather(feature, interp_idx) 333 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 334 | return interpolated_features 335 | 336 | @staticmethod 337 | def gather_neighbour(pc, neighbor_idx): 338 | # gather the coordinates or features of neighboring points 339 | batch_size = tf.shape(pc)[0] 340 | num_points = tf.shape(pc)[1] 341 | d = pc.get_shape()[2].value 342 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 343 | features = tf.batch_gather(pc, index_input) 344 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 345 | return features 346 | 347 | @staticmethod 348 | def att_pooling(feature_set, d_out, name, is_training): 349 | batch_size = tf.shape(feature_set)[0] 350 | num_points = tf.shape(feature_set)[1] 351 | num_neigh = tf.shape(feature_set)[2] 352 | d = feature_set.get_shape()[3].value 353 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 354 | att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc') 355 | att_scores = tf.nn.softmax(att_activation, axis=1) 356 | f_agg = f_reshaped * att_scores 357 | f_agg = tf.reduce_sum(f_agg, axis=1) 358 | f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d]) 359 | f_agg = tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 360 | return f_agg 361 | -------------------------------------------------------------------------------- /tf_util.py: -------------------------------------------------------------------------------- 1 | """ Wrapper functions for TensorFlow layers. 2 | 3 | Author: Charles R. Qi 4 | Date: November 2016 5 | """ 6 | 7 | import numpy as np 8 | import tensorflow as tf 9 | 10 | 11 | def _variable_on_cpu(name, shape, initializer, use_fp16=False): 12 | """Helper to create a Variable stored on CPU memory. 13 | Args: 14 | name: name of the variable 15 | shape: list of ints 16 | initializer: initializer for Variable 17 | Returns: 18 | Variable Tensor 19 | """ 20 | with tf.device('/cpu:0'): 21 | dtype = tf.float16 if use_fp16 else tf.float32 22 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 23 | return var 24 | 25 | 26 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True): 27 | """Helper to create an initialized Variable with weight decay. 28 | 29 | Note that the Variable is initialized with a truncated normal distribution. 30 | A weight decay is added only if one is specified. 31 | 32 | Args: 33 | name: name of the variable 34 | shape: list of ints 35 | stddev: standard deviation of a truncated Gaussian 36 | wd: add L2Loss weight decay multiplied by this float. If None, weight 37 | decay is not added for this Variable. 38 | use_xavier: bool, whether to use xavier initializer 39 | 40 | Returns: 41 | Variable Tensor 42 | """ 43 | if use_xavier: 44 | initializer = tf.contrib.layers.xavier_initializer() 45 | var = _variable_on_cpu(name, shape, initializer) 46 | else: 47 | # initializer = tf.truncated_normal_initializer(stddev=stddev) 48 | with tf.device('/cpu:0'): 49 | var = tf.truncated_normal(shape, stddev=np.sqrt(2 / shape[-1])) 50 | var = tf.round(var * tf.constant(1000, dtype=tf.float32)) / tf.constant(1000, dtype=tf.float32) 51 | var = tf.Variable(var, name='weights') 52 | if wd is not None: 53 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 54 | tf.add_to_collection('losses', weight_decay) 55 | return var 56 | 57 | 58 | def conv1d(inputs, 59 | num_output_channels, 60 | kernel_size, 61 | scope, 62 | stride=1, 63 | padding='SAME', 64 | use_xavier=True, 65 | stddev=1e-3, 66 | weight_decay=0.0, 67 | activation_fn=tf.nn.relu, 68 | bn=False, 69 | bn_decay=None, 70 | is_training=None): 71 | """ 1D convolution with non-linear operation. 72 | 73 | Args: 74 | inputs: 3-D tensor variable BxLxC 75 | num_output_channels: int 76 | kernel_size: int 77 | scope: string 78 | stride: int 79 | padding: 'SAME' or 'VALID' 80 | use_xavier: bool, use xavier_initializer if true 81 | stddev: float, stddev for truncated_normal init 82 | weight_decay: float 83 | activation_fn: function 84 | bn: bool, whether to use batch norm 85 | bn_decay: float or float tensor variable in [0,1] 86 | is_training: bool Tensor variable 87 | 88 | Returns: 89 | Variable tensor 90 | """ 91 | with tf.variable_scope(scope) as sc: 92 | num_in_channels = inputs.get_shape()[-1].value 93 | kernel_shape = [kernel_size, 94 | num_in_channels, num_output_channels] 95 | kernel = _variable_with_weight_decay('weights', 96 | shape=kernel_shape, 97 | use_xavier=use_xavier, 98 | stddev=stddev, 99 | wd=weight_decay) 100 | outputs = tf.nn.conv1d(inputs, kernel, 101 | stride=stride, 102 | padding=padding) 103 | biases = _variable_on_cpu('biases', [num_output_channels], 104 | tf.constant_initializer(0.0)) 105 | outputs = tf.nn.bias_add(outputs, biases) 106 | 107 | if bn: 108 | outputs = batch_norm_for_conv1d(outputs, is_training, 109 | bn_decay=bn_decay, scope='bn') 110 | if activation_fn is not None: 111 | outputs = activation_fn(outputs) 112 | return outputs 113 | 114 | 115 | def conv2d(inputs, 116 | num_output_channels, 117 | kernel_size, 118 | scope, 119 | stride=[1, 1], 120 | padding='SAME', 121 | bn=False, 122 | is_training=None, 123 | use_xavier=False, 124 | stddev=1e-3, 125 | weight_decay=0.0, 126 | activation_fn=tf.nn.relu, 127 | bn_decay=None): 128 | """ 2D convolution with non-linear operation. 129 | 130 | Args: 131 | inputs: 4-D tensor variable BxHxWxC 132 | num_output_channels: int 133 | kernel_size: a list of 2 ints 134 | scope: string 135 | stride: a list of 2 ints 136 | padding: 'SAME' or 'VALID' 137 | use_xavier: bool, use xavier_initializer if true 138 | stddev: float, stddev for truncated_normal init 139 | weight_decay: float 140 | activation_fn: function 141 | bn: bool, whether to use batch norm 142 | bn_decay: float or float tensor variable in [0,1] 143 | is_training: bool Tensor variable 144 | 145 | Returns: 146 | Variable tensor 147 | """ 148 | with tf.variable_scope(scope) as sc: 149 | kernel_h, kernel_w = kernel_size 150 | num_in_channels = inputs.get_shape()[-1].value 151 | kernel_shape = [kernel_h, kernel_w, 152 | num_in_channels, num_output_channels] 153 | kernel = _variable_with_weight_decay('weights', 154 | shape=kernel_shape, 155 | use_xavier=use_xavier, 156 | stddev=stddev, 157 | wd=weight_decay) 158 | stride_h, stride_w = stride 159 | outputs = tf.nn.conv2d(inputs, kernel, 160 | [1, stride_h, stride_w, 1], 161 | padding=padding) 162 | biases = _variable_on_cpu('biases', [num_output_channels], 163 | tf.constant_initializer(0.0)) 164 | outputs = tf.nn.bias_add(outputs, biases) 165 | 166 | if bn: 167 | outputs = tf.layers.batch_normalization(outputs, momentum=0.99, epsilon=1e-6, training=is_training) 168 | if activation_fn is not None: 169 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 170 | return outputs 171 | 172 | 173 | def conv2d_transpose(inputs, 174 | num_output_channels, 175 | kernel_size, 176 | scope, 177 | stride=[1, 1], 178 | padding='SAME', 179 | use_xavier=False, 180 | stddev=1e-3, 181 | weight_decay=0.0, 182 | activation_fn=tf.nn.relu, 183 | bn=False, 184 | bn_decay=None, 185 | is_training=None): 186 | """ 2D convolution transpose with non-linear operation. 187 | 188 | Args: 189 | inputs: 4-D tensor variable BxHxWxC 190 | num_output_channels: int 191 | kernel_size: a list of 2 ints 192 | scope: string 193 | stride: a list of 2 ints 194 | padding: 'SAME' or 'VALID' 195 | use_xavier: bool, use xavier_initializer if true 196 | stddev: float, stddev for truncated_normal init 197 | weight_decay: float 198 | activation_fn: function 199 | bn: bool, whether to use batch norm 200 | bn_decay: float or float tensor variable in [0,1] 201 | is_training: bool Tensor variable 202 | 203 | Returns: 204 | Variable tensor 205 | 206 | Note: conv2d(conv2d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a 207 | """ 208 | with tf.variable_scope(scope) as sc: 209 | kernel_h, kernel_w = kernel_size 210 | num_in_channels = inputs.get_shape()[-1].value 211 | kernel_shape = [kernel_h, kernel_w, 212 | num_output_channels, num_in_channels] # reversed to conv2d 213 | kernel = _variable_with_weight_decay('weights', 214 | shape=kernel_shape, 215 | use_xavier=use_xavier, 216 | stddev=stddev, 217 | wd=weight_decay) 218 | stride_h, stride_w = stride 219 | 220 | # from slim.convolution2d_transpose 221 | def get_deconv_dim(dim_size, stride_size, kernel_size, padding): 222 | dim_size *= stride_size 223 | 224 | if padding == 'VALID' and dim_size is not None: 225 | dim_size += max(kernel_size - stride_size, 0) 226 | return dim_size 227 | 228 | # caculate output shape 229 | batch_size = tf.shape(inputs)[0] 230 | height = tf.shape(inputs)[1] 231 | width = tf.shape(inputs)[2] 232 | out_height = get_deconv_dim(height, stride_h, kernel_h, padding) 233 | out_width = get_deconv_dim(width, stride_w, kernel_w, padding) 234 | output_shape = tf.stack([batch_size, out_height, out_width, num_output_channels], axis=0) 235 | 236 | outputs = tf.nn.conv2d_transpose(inputs, kernel, output_shape, 237 | [1, stride_h, stride_w, 1], 238 | padding=padding) 239 | biases = _variable_on_cpu('biases', [num_output_channels], 240 | tf.constant_initializer(0.0)) 241 | outputs = tf.nn.bias_add(outputs, biases) 242 | 243 | if bn: 244 | # outputs = batch_norm_for_conv2d(outputs, is_training, 245 | # bn_decay=bn_decay, scope='bn') 246 | outputs = tf.layers.batch_normalization(outputs, momentum=0.99, epsilon=1e-6, training=is_training) 247 | if activation_fn is not None: 248 | # outputs = activation_fn(outputs) 249 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 250 | return outputs 251 | 252 | 253 | def conv3d(inputs, 254 | num_output_channels, 255 | kernel_size, 256 | scope, 257 | stride=[1, 1, 1], 258 | padding='SAME', 259 | use_xavier=True, 260 | stddev=1e-3, 261 | weight_decay=0.0, 262 | activation_fn=tf.nn.relu, 263 | bn=False, 264 | bn_decay=None, 265 | is_training=None): 266 | """ 3D convolution with non-linear operation. 267 | 268 | Args: 269 | inputs: 5-D tensor variable BxDxHxWxC 270 | num_output_channels: int 271 | kernel_size: a list of 3 ints 272 | scope: string 273 | stride: a list of 3 ints 274 | padding: 'SAME' or 'VALID' 275 | use_xavier: bool, use xavier_initializer if true 276 | stddev: float, stddev for truncated_normal init 277 | weight_decay: float 278 | activation_fn: function 279 | bn: bool, whether to use batch norm 280 | bn_decay: float or float tensor variable in [0,1] 281 | is_training: bool Tensor variable 282 | 283 | Returns: 284 | Variable tensor 285 | """ 286 | with tf.variable_scope(scope) as sc: 287 | kernel_d, kernel_h, kernel_w = kernel_size 288 | num_in_channels = inputs.get_shape()[-1].value 289 | kernel_shape = [kernel_d, kernel_h, kernel_w, 290 | num_in_channels, num_output_channels] 291 | kernel = _variable_with_weight_decay('weights', 292 | shape=kernel_shape, 293 | use_xavier=use_xavier, 294 | stddev=stddev, 295 | wd=weight_decay) 296 | stride_d, stride_h, stride_w = stride 297 | outputs = tf.nn.conv3d(inputs, kernel, 298 | [1, stride_d, stride_h, stride_w, 1], 299 | padding=padding) 300 | biases = _variable_on_cpu('biases', [num_output_channels], 301 | tf.constant_initializer(0.0)) 302 | outputs = tf.nn.bias_add(outputs, biases) 303 | 304 | if bn: 305 | outputs = batch_norm_for_conv3d(outputs, is_training, 306 | bn_decay=bn_decay, scope='bn') 307 | 308 | if activation_fn is not None: 309 | outputs = activation_fn(outputs) 310 | return outputs 311 | 312 | 313 | def fully_connected(inputs, 314 | num_outputs, 315 | scope, 316 | use_xavier=True, 317 | stddev=1e-3, 318 | weight_decay=0.0, 319 | activation_fn=tf.nn.relu, 320 | bn=False, 321 | bn_decay=None, 322 | is_training=None): 323 | """ Fully connected layer with non-linear operation. 324 | 325 | Args: 326 | inputs: 2-D tensor BxN 327 | num_outputs: int 328 | 329 | Returns: 330 | Variable tensor of size B x num_outputs. 331 | """ 332 | with tf.variable_scope(scope) as sc: 333 | num_input_units = inputs.get_shape()[-1].value 334 | weights = _variable_with_weight_decay('weights', 335 | shape=[num_input_units, num_outputs], 336 | use_xavier=use_xavier, 337 | stddev=stddev, 338 | wd=weight_decay) 339 | outputs = tf.matmul(inputs, weights) 340 | biases = _variable_on_cpu('biases', [num_outputs], 341 | tf.constant_initializer(0.0)) 342 | outputs = tf.nn.bias_add(outputs, biases) 343 | 344 | if bn: 345 | outputs = batch_norm_for_fc(outputs, is_training, bn_decay, 'bn') 346 | 347 | if activation_fn is not None: 348 | # outputs = activation_fn(outputs) 349 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 350 | return outputs 351 | 352 | 353 | def max_pool2d(inputs, 354 | kernel_size, 355 | scope, 356 | stride=[2, 2], 357 | padding='VALID'): 358 | """ 2D max pooling. 359 | 360 | Args: 361 | inputs: 4-D tensor BxHxWxC 362 | kernel_size: a list of 2 ints 363 | stride: a list of 2 ints 364 | 365 | Returns: 366 | Variable tensor 367 | """ 368 | with tf.variable_scope(scope) as sc: 369 | kernel_h, kernel_w = kernel_size 370 | stride_h, stride_w = stride 371 | outputs = tf.nn.max_pool(inputs, 372 | ksize=[1, kernel_h, kernel_w, 1], 373 | strides=[1, stride_h, stride_w, 1], 374 | padding=padding, 375 | name=sc.name) 376 | return outputs 377 | 378 | 379 | def avg_pool2d(inputs, 380 | kernel_size, 381 | scope, 382 | stride=[2, 2], 383 | padding='VALID'): 384 | """ 2D avg pooling. 385 | 386 | Args: 387 | inputs: 4-D tensor BxHxWxC 388 | kernel_size: a list of 2 ints 389 | stride: a list of 2 ints 390 | 391 | Returns: 392 | Variable tensor 393 | """ 394 | with tf.variable_scope(scope) as sc: 395 | kernel_h, kernel_w = kernel_size 396 | stride_h, stride_w = stride 397 | outputs = tf.nn.avg_pool(inputs, 398 | ksize=[1, kernel_h, kernel_w, 1], 399 | strides=[1, stride_h, stride_w, 1], 400 | padding=padding, 401 | name=sc.name) 402 | return outputs 403 | 404 | 405 | def max_pool3d(inputs, 406 | kernel_size, 407 | scope, 408 | stride=[2, 2, 2], 409 | padding='VALID'): 410 | """ 3D max pooling. 411 | 412 | Args: 413 | inputs: 5-D tensor BxDxHxWxC 414 | kernel_size: a list of 3 ints 415 | stride: a list of 3 ints 416 | 417 | Returns: 418 | Variable tensor 419 | """ 420 | with tf.variable_scope(scope) as sc: 421 | kernel_d, kernel_h, kernel_w = kernel_size 422 | stride_d, stride_h, stride_w = stride 423 | outputs = tf.nn.max_pool3d(inputs, 424 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 425 | strides=[1, stride_d, stride_h, stride_w, 1], 426 | padding=padding, 427 | name=sc.name) 428 | return outputs 429 | 430 | 431 | def avg_pool3d(inputs, 432 | kernel_size, 433 | scope, 434 | stride=[2, 2, 2], 435 | padding='VALID'): 436 | """ 3D avg pooling. 437 | 438 | Args: 439 | inputs: 5-D tensor BxDxHxWxC 440 | kernel_size: a list of 3 ints 441 | stride: a list of 3 ints 442 | 443 | Returns: 444 | Variable tensor 445 | """ 446 | with tf.variable_scope(scope) as sc: 447 | kernel_d, kernel_h, kernel_w = kernel_size 448 | stride_d, stride_h, stride_w = stride 449 | outputs = tf.nn.avg_pool3d(inputs, 450 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 451 | strides=[1, stride_d, stride_h, stride_w, 1], 452 | padding=padding, 453 | name=sc.name) 454 | return outputs 455 | 456 | 457 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay): 458 | """ Batch normalization on convolutional maps and beyond... 459 | Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow 460 | 461 | Args: 462 | inputs: Tensor, k-D input ... x C could be BC or BHWC or BDHWC 463 | is_training: boolean tf.Varialbe, true indicates training phase 464 | scope: string, variable scope 465 | moments_dims: a list of ints, indicating dimensions for moments calculation 466 | bn_decay: float or float tensor variable, controling moving average weight 467 | Return: 468 | normed: batch-normalized maps 469 | """ 470 | with tf.variable_scope(scope) as sc: 471 | num_channels = inputs.get_shape()[-1].value 472 | beta = tf.Variable(tf.constant(0.0, shape=[num_channels]), 473 | name='beta', trainable=True) 474 | gamma = tf.Variable(tf.constant(1.0, shape=[num_channels]), 475 | name='gamma', trainable=True) 476 | batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments') 477 | decay = bn_decay if bn_decay is not None else 0.9 478 | ema = tf.train.ExponentialMovingAverage(decay=decay) 479 | # Operator that maintains moving averages of variables. 480 | ema_apply_op = tf.cond(is_training, 481 | lambda: ema.apply([batch_mean, batch_var]), 482 | lambda: tf.no_op()) 483 | 484 | # Update moving average and return current batch's avg and var. 485 | def mean_var_with_update(): 486 | with tf.control_dependencies([ema_apply_op]): 487 | return tf.identity(batch_mean), tf.identity(batch_var) 488 | 489 | # ema.average returns the Variable holding the average of var. 490 | mean, var = tf.cond(is_training, 491 | mean_var_with_update, 492 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 493 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 494 | return normed 495 | 496 | 497 | def batch_norm_for_fc(inputs, is_training, bn_decay, scope): 498 | """ Batch normalization on FC data. 499 | 500 | Args: 501 | inputs: Tensor, 2D BxC input 502 | is_training: boolean tf.Varialbe, true indicates training phase 503 | bn_decay: float or float tensor variable, controling moving average weight 504 | scope: string, variable scope 505 | Return: 506 | normed: batch-normalized maps 507 | """ 508 | return batch_norm_template(inputs, is_training, scope, [0, ], bn_decay) 509 | 510 | 511 | def batch_norm_for_conv1d(inputs, is_training, bn_decay, scope): 512 | """ Batch normalization on 1D convolutional maps. 513 | 514 | Args: 515 | inputs: Tensor, 3D BLC input maps 516 | is_training: boolean tf.Varialbe, true indicates training phase 517 | bn_decay: float or float tensor variable, controling moving average weight 518 | scope: string, variable scope 519 | Return: 520 | normed: batch-normalized maps 521 | """ 522 | return batch_norm_template(inputs, is_training, scope, [0, 1], bn_decay) 523 | 524 | 525 | def batch_norm_for_conv2d(inputs, is_training, bn_decay, scope): 526 | """ Batch normalization on 2D convolutional maps. 527 | 528 | Args: 529 | inputs: Tensor, 4D BHWC input maps 530 | is_training: boolean tf.Varialbe, true indicates training phase 531 | bn_decay: float or float tensor variable, controling moving average weight 532 | scope: string, variable scope 533 | Return: 534 | normed: batch-normalized maps 535 | """ 536 | return batch_norm_template(inputs, is_training, scope, [0, 1, 2], bn_decay) 537 | 538 | 539 | def batch_norm_for_conv3d(inputs, is_training, bn_decay, scope): 540 | """ Batch normalization on 3D convolutional maps. 541 | 542 | Args: 543 | inputs: Tensor, 5D BDHWC input maps 544 | is_training: boolean tf.Varialbe, true indicates training phase 545 | bn_decay: float or float tensor variable, controling moving average weight 546 | scope: string, variable scope 547 | Return: 548 | normed: batch-normalized maps 549 | """ 550 | return batch_norm_template(inputs, is_training, scope, [0, 1, 2, 3], bn_decay) 551 | 552 | 553 | def dropout(inputs, 554 | is_training, 555 | scope, 556 | keep_prob=0.5, 557 | noise_shape=None): 558 | """ Dropout layer. 559 | 560 | Args: 561 | inputs: tensor 562 | is_training: boolean tf.Variable 563 | scope: string 564 | keep_prob: float in [0,1] 565 | noise_shape: list of ints 566 | 567 | Returns: 568 | tensor variable 569 | """ 570 | with tf.variable_scope(scope) as sc: 571 | outputs = tf.cond(is_training, 572 | lambda: tf.nn.dropout(inputs, keep_prob, noise_shape), 573 | lambda: inputs) 574 | return outputs 575 | --------------------------------------------------------------------------------