├── .idea ├── .gitignore ├── misc.xml ├── inspectionProfiles │ └── profiles_settings.xml ├── modules.xml └── BAF-LAC.iml ├── imgs └── BAF-LAC.jpg ├── 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 ├── meta │ ├── class_names.txt │ └── anno_paths.txt ├── nearest_neighbors │ ├── test.py │ ├── setup.py │ ├── knn_.h │ ├── knn.pyx │ ├── knn_.cxx │ └── KDTreeTableAdaptor.h ├── 6_fold_cv.py ├── data_prepare_s3dis.py ├── data_prepare_semantickitti.py ├── data_prepare_semantic3d.py ├── download_semantic3d.sh └── semantic-kitti.yaml ├── compile_op.sh ├── helper_requirements.txt ├── jobs_6_fold_cv_s3dis.sh ├── jobs_test_semantickitti.sh ├── README.md ├── tester_Semantic3D.py ├── tester_S3DIS.py ├── helper_tf_util.py ├── tester_SemanticKITTI.py ├── helper_ply.py ├── LICENSE ├── main_SemanticKITTI.py ├── main_S3DIS.py ├── helper_tool.py ├── main_Semantic3D.py └── BAF_LAC.py /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | -------------------------------------------------------------------------------- /imgs/BAF-LAC.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xiangxu-0103/BAF-LAC/HEAD/imgs/BAF-LAC.jpg -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 ../../../ 8 | -------------------------------------------------------------------------------- /utils/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /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.1.2 9 | -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.idea/BAF-LAC.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | ext_modules = [Extension( 8 | "nearest_neighbors", 9 | sources=["knn.pyx", "knn_.cxx"], # source file(s) 10 | include_dirs=["./", numpy.get_include()], 11 | language="c++", 12 | extra_compile_args=["-std=c++11", "-fopenmp"], 13 | extra_link_args=["-std=c++11", '-fopenmp'], 14 | )] 15 | 16 | setup( 17 | name="KNN NanoFLANN", 18 | ext_modules=ext_modules, 19 | cmdclass={'build_ext': build_ext}, 20 | ) 21 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /jobs_6_fold_cv_s3dis.sh: -------------------------------------------------------------------------------- 1 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 1 2 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 1 3 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 2 4 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 2 5 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 3 6 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 3 7 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 4 8 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 4 9 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 5 10 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 5 11 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 6 12 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 6 13 | -------------------------------------------------------------------------------- /jobs_test_semantickitti.sh: -------------------------------------------------------------------------------- 1 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 11 2 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 12 3 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 13 4 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 14 5 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 15 6 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 16 7 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 17 8 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 18 9 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 19 10 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 20 11 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 21 12 | -------------------------------------------------------------------------------- /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); 28 | -------------------------------------------------------------------------------- /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 | } 68 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/6_fold_cv.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob, os, sys 3 | 4 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 5 | ROOT_DIR = os.path.dirname(BASE_DIR) 6 | sys.path.append(ROOT_DIR) 7 | from helper_ply import read_ply 8 | from helper_tool import Plot 9 | 10 | 11 | if __name__ == '__main__': 12 | base_dir = '/home/data/S3DIS/results' 13 | original_data_dir = '/home/data/S3DIS/original_ply' 14 | data_path = glob.glob(os.path.join(base_dir, '*.ply')) 15 | data_path = np.sort(data_path) 16 | 17 | test_total_correct = 0 18 | test_total_seen = 0 19 | gt_classes = [0 for _ in range(13)] 20 | positive_classes = [0 for _ in range(13)] 21 | true_positive_classes = [0 for _ in range(13)] 22 | visualization = False 23 | 24 | for file_name in data_path: 25 | pred_data = read_ply(file_name) 26 | pred = pred_data['pred'] 27 | original_data = read_ply(os.path.join(original_data_dir, file_name.split('/')[-1][:-4] + '.ply')) 28 | labels = original_data['class'] 29 | points = np.vstack((original_data['x'], original_data['y'], original_data['z'])).T 30 | 31 | ################## 32 | # Visualize data # 33 | ################## 34 | if visualization: 35 | colors = np.vstack((original_data['red'], original_data['green'], original_data['blue'])).T 36 | xyzrgb = np.concatenate([points, colors], axis=-1) 37 | Plot.draw_pc(xyzrgb) # visualize raw point clouds 38 | Plot.draw_pc_sem_ins(points, labels) # visualize ground-truth 39 | Plot.draw_pc_sem_ins(points, pred) # visualize prediction 40 | 41 | correct = np.sum(pred == labels) 42 | print(str(file_name.split('/')[-1][:-4]) + '_acc:' + str(correct / float(len(labels)))) 43 | test_total_correct += correct 44 | test_total_seen += len(labels) 45 | 46 | for j in range(len(labels)): 47 | gt_l = int(labels[j]) 48 | pred_l = int(pred[j]) 49 | gt_classes[gt_l] += 1 50 | positive_classes[pred_l] += 1 51 | true_positive_classes[gt_l] += int(gt_l == pred_l) 52 | 53 | iou_list = [] 54 | for n in range(13): 55 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 56 | iou_list.append(iou) 57 | mean_iou = sum(iou_list) / 13.0 58 | print('eval accuracy: {}'.format(test_total_correct / float(test_total_seen))) 59 | print('mean IoU: {}'.format(mean_iou)) 60 | print(iou_list) 61 | 62 | acc_list = [] 63 | for n in range(13): 64 | acc = true_positive_classes[n] / float(gt_classes[n]) 65 | acc_list.append(acc) 66 | mean_acc = sum(acc_list) / 13.0 67 | print('mAcc value is: {}'.format(mean_acc)) 68 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/data_prepare_s3dis.py: -------------------------------------------------------------------------------- 1 | from sklearn.neighbors import KDTree 2 | from os.path import join, exists, dirname, abspath 3 | import numpy as np 4 | import pandas as pd 5 | import os, sys, glob, pickle 6 | 7 | BASE_DIR = dirname(abspath(__file__)) 8 | ROOT_DIR = dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(ROOT_DIR) 11 | from helper_ply import write_ply 12 | from helper_tool import DataProcessing as DP 13 | 14 | dataset_path = '/home/data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version' 15 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 16 | anno_paths = [join(dataset_path, p) for p in anno_paths] 17 | 18 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 19 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 20 | 21 | sub_grid_size = 0.04 22 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 23 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 24 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 25 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 26 | out_format = '.ply' 27 | 28 | 29 | def convert_pc2ply(anno_path, save_path): 30 | """ 31 | Convert original dataset files to ply file (each line is XYZRGBL). 32 | We aggregated all the points from each instance in the room. 33 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 34 | :param save_path: path to save original point clouds (each line is XYZRGBL) 35 | :return: None 36 | """ 37 | data_list = [] 38 | 39 | for f in glob.glob(join(anno_path, '*.txt')): 40 | class_name = os.path.basename(f).split('_')[0] 41 | if class_name not in gt_class: # note: in some room there is 'staris' ckass.. 42 | class_name = 'clutter' 43 | pc = pd.read_csv(f, header=None, delim_whitespace=True).values 44 | labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 45 | data_list.append(np.concatenate([pc, labels], 1)) # Nx7 46 | 47 | pc_label = np.concatenate(data_list, 0) 48 | xyz_min = np.amin(pc_label, axis=0)[0:3] 49 | pc_label[:, 0:3] -= xyz_min 50 | 51 | xyz = pc_label[:, :3].astype(np.float32) 52 | colors = pc_label[:, 3:6].astype(np.uint8) 53 | labels = pc_label[:, 6].astype(np.uint8) 54 | write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 55 | 56 | # Save sub_cloud and KDTree file 57 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size) 58 | sub_colors = sub_colors / 255.0 59 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 60 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 61 | 62 | search_tree = KDTree(sub_xyz) 63 | kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl') 64 | with open(kd_tree_file, 'wb') as f: 65 | pickle.dump(search_tree, f) 66 | 67 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 68 | proj_idx = proj_idx.astype(np.int32) 69 | proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl') 70 | with open(proj_save, 'wb') as f: 71 | pickle.dump([proj_idx, labels], f) 72 | 73 | 74 | if __name__ == '__main__': 75 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 76 | for annotation_path in anno_paths: 77 | print(annotation_path) 78 | elements = str(annotation_path).split('/') 79 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 80 | convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name)) 81 | -------------------------------------------------------------------------------- /utils/data_prepare_semantickitti.py: -------------------------------------------------------------------------------- 1 | import pickle, yaml, os, sys 2 | import numpy as np 3 | from os.path import join, exists, dirname, abspath 4 | from sklearn.neighbors import KDTree 5 | 6 | BASE_DIR = dirname(abspath(__file__)) 7 | ROOT_DIR = dirname(BASE_DIR) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(ROOT_DIR) 10 | from helper_tool import DataProcessing as DP 11 | 12 | data_config = os.path.join(BASE_DIR, 'semantic-kitti.yaml') 13 | DATA = yaml.safe_load(open(data_config, 'r')) 14 | remap_dict = DATA['learning_map'] 15 | max_key = max(remap_dict.keys()) 16 | remap_lut = np.zeros((max_key + 100), dtype=np.int32) 17 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values()) 18 | 19 | grid_size = 0.06 20 | dataset_path = '/home/data/semantic_kitti/dataset/sequences' 21 | output_path = '/home/data/semantic_kitti/dataset/sequences' + '_' + str(grid_size) 22 | seq_list = np.sort(os.listdir(dataset_path)) 23 | 24 | for seq_id in seq_list: 25 | print('sequence' + seq_id + ' start') 26 | seq_path = join(dataset_path, seq_id) 27 | seq_path_out = join(output_path, seq_id) 28 | pc_path = join(seq_path, 'velodyne') 29 | pc_path_out = join(seq_path_out, 'velodyne') 30 | KDTree_path_out = join(seq_path_out, 'KDTree') 31 | os.makedirs(seq_path_out) if not exists(seq_path_out) else None 32 | os.makedirs(pc_path_out) if not exists(pc_path_out) else None 33 | os.makedirs(KDTree_path_out) if not exists(KDTree_path_out) else None 34 | 35 | if int(seq_id) < 11: 36 | label_path = join(seq_path, 'labels') 37 | label_path_out = join(seq_path_out, 'labels') 38 | os.makedirs(label_path_out) if not exists(label_path_out) else None 39 | scan_list = np.sort(os.listdir(pc_path)) 40 | for scan_id in scan_list: 41 | print(scan_id) 42 | points = DP.load_pc_kitti(join(pc_path, scan_id)) 43 | labels = DP.load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut) 44 | sub_points, sub_labels = DP.grid_sub_sampling(points, labels=labels, grid_size=grid_size) 45 | search_tree = KDTree(sub_points) 46 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl') 47 | np.save(join(pc_path_out, scan_id)[:-4], sub_points) 48 | np.save(join(label_path_out, scan_id)[:-4], sub_labels) 49 | with open(KDTree_save, 'wb') as f: 50 | pickle.dump(search_tree, f) 51 | if seq_id == '08': 52 | proj_path = join(seq_path_out, 'proj') 53 | os.makedirs(proj_path) if not exists(proj_path) else None 54 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False)) 55 | proj_inds = proj_inds.astype(np.int32) 56 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl') 57 | with open(proj_save, 'wb') as f: 58 | pickle.dump([proj_inds], f) 59 | else: 60 | proj_path = join(seq_path_out, 'proj') 61 | os.makedirs(proj_path) if not exists(proj_path) else None 62 | scan_list = np.sort(os.listdir(pc_path)) 63 | for scan_id in scan_list: 64 | print(scan_id) 65 | points = DP.load_pc_kitti(join(pc_path, scan_id)) 66 | sub_points = DP.grid_sub_sampling(points, grid_size=0.06) 67 | search_tree = KDTree(sub_points) 68 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False)) 69 | proj_inds = proj_inds.astype(np.int32) 70 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl') 71 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl') 72 | np.save(join(pc_path_out, scan_id)[:-4], sub_points) 73 | with open(KDTree_save, 'wb') as f: 74 | pickle.dump(search_tree, f) 75 | with open(proj_save, 'wb') as f: 76 | pickle.dump([proj_inds], f) 77 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Backward Attentive Fusing Network with Local Aggregation Classifier for 3D Point Cloud Semantic Segmentation 2 | 3 | This is the official implementation of **BAF-LAC** (TIP 2021), a novel point cloud semantic segmentation paradigm that introduces more context information. For technical details, please refer to: 4 | 5 | **Backward Attentive Fusing Network with Local Aggregation Classifier for 3D Point Cloud Semantic Segmentation**
6 | Hui Shuai, Xiang Xu, Qingshan Liu.
7 | **[[paper](https://ieeexplore.ieee.org/abstract/document/9410334)]**
8 | 9 | ![architecture](./imgs/BAF-LAC.jpg) 10 | 11 | ### (1) Setup 12 | 13 | This code has been tested with Python 3.6, Tensorflow 1.13.1, CUDA 10.0 on Ubuntu 16.04. 14 | 15 | - Clone the repository 16 | 17 | ```shell 18 | git clone https://github.com/Xiangxu-0103/BAF-LAC.git && cd BAF-LAC 19 | ``` 20 | 21 | - Setup python environment 22 | 23 | ```shell 24 | conda create -n baflac python=3.6.8 25 | conda activate baflac 26 | pip install tensorflow-gpu==1.13.1 27 | pip install -r helper_requirements.txt 28 | sh compile_op.sh 29 | ``` 30 | 31 | ### (2) S3DIS 32 | 33 | S3DIS dataset can be found here. Download the files named "Stanford3dDataset_v1.2_Aligned_Version.zip". Uncompress the folder and move it to `/home/data/S3DIS`. 34 | 35 | - Preparing the dataset: 36 | 37 | ```shell 38 | python utils/data_prepare_s3dis.py 39 | ``` 40 | 41 | - Start 6-fold cross validation: 42 | 43 | ```shell 44 | sh jobs_6_fold_cv_s3dis.sh 45 | ``` 46 | 47 | - Move all the generated results (*.ply) in `/test` folder to `/home/data/S3DIS/results`, calculate the final mean IoU results: 48 | 49 | ```shell 50 | python utils/6_fold_cv.py 51 | ``` 52 | 53 | ### (3) Semantic3D 54 | 55 | 7zip is required to uncompress the raw data in this dataset, to install p7zip: 56 | 57 | ```shell 58 | sudo apt-get install p7zip-full 59 | ``` 60 | 61 | - Download and extract the dataset. First, please specify the path of the dataset by changing the `BASE_DIR` in "download_semantic3d.sh" 62 | 63 | ```shell 64 | sh utils/download_semantic3d.sh 65 | ``` 66 | 67 | - Preparing the dataset: 68 | 69 | ```shell 70 | python utils/data_prepare_semantic3d.py 71 | ``` 72 | 73 | - Start training: 74 | 75 | ```shell 76 | python main_Semantic3D.py --mode train --gpu 0 77 | ``` 78 | 79 | - Evaluation: 80 | 81 | ```shell 82 | python main_Semantic3D.py --mode test --gpu 0 83 | ``` 84 | 85 | **Note:** 86 | 87 | - Preferably with more than 64G RAM to process this dataset due to the large volume of point cloud 88 | 89 | ### (4) SemanticKITTI 90 | 91 | SemanticKITTI dataset can be found here. Download the files related to semantic segmentation and extract everything into the same folder. Uncompress the folder and move it to `/home/data/semantic_kitti/dataset`. 92 | 93 | - Preparing the dataset 94 | 95 | ```shell 96 | python utils/data_prepare_semantickitti.py 97 | ``` 98 | 99 | - Start training: 100 | 101 | ```shell 102 | python main_SemanticKITTI.py --mode train --gpu 0 103 | ``` 104 | 105 | - Evaluation: 106 | 107 | ```shell 108 | sh jobs_test_semantickitti.sh 109 | ``` 110 | 111 | ### Citation 112 | If you find our work useful in your research, please consider citing: 113 | 114 | ``` 115 | @article{shuai2021backward, 116 | title={Backward Attentive Fusing Network With Local Aggregation Classifier for 3D Point Cloud Semantic Segmentation}, 117 | author={Shuai, Hui and Xu, Xiang and Liu, Qingshan}, 118 | journal={IEEE Transactions on Image Processing}, 119 | volume={30}, 120 | pages={4973--4984}, 121 | year={2021}, 122 | publisher={IEEE} 123 | } 124 | ``` 125 | 126 | ### Acknowledge 127 | 128 | - Our code refers to RandLA-Net. 129 | -------------------------------------------------------------------------------- /utils/data_prepare_semantic3d.py: -------------------------------------------------------------------------------- 1 | from sklearn.neighbors import KDTree 2 | from os.path import join, exists, dirname, abspath 3 | import numpy as np 4 | import os, glob, pickle 5 | import sys 6 | 7 | BASE_DIR = dirname(abspath(__file__)) 8 | ROOT_DIR = dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(ROOT_DIR) 11 | from helper_ply import write_ply 12 | from helper_tool import DataProcessing as DP 13 | 14 | grid_size = 0.06 15 | dataset_path = '/home/data/semantic3d/original_data' 16 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 17 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(grid_size)) 18 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 19 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 20 | 21 | for pc_path in glob.glob(join(dataset_path, '*.txt')): 22 | print(pc_path) 23 | file_name = pc_path.split('/')[-1][:-4] 24 | 25 | # Check if it has already calculated 26 | if exists(join(sub_pc_folder, file_name + '_KDTree.pkl')): 27 | continue 28 | 29 | pc = DP.load_pc_semantic3d(pc_path) 30 | # Check if label exists 31 | label_path = pc_path[:-4] + '.labels' 32 | if exists(label_path): 33 | labels = DP.load_label_semantic3d(label_path) 34 | full_ply_path = join(original_pc_folder, file_name + '.ply') 35 | 36 | # Subsample to save space 37 | sub_points, sub_colors, sub_labels = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), 38 | pc[:, 4:7].astype(np.uint8), labels, 0.01) 39 | sub_labels = np.squeeze(sub_labels) 40 | 41 | write_ply(full_ply_path, (sub_points, sub_colors, sub_labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 42 | 43 | # Save sub_cloud and KDTree file 44 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(sub_points, sub_colors, sub_labels, grid_size) 45 | sub_colors = sub_colors / 255.0 46 | sub_labels = np.squeeze(sub_labels) 47 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 48 | write_ply(sub_ply_file, (sub_xyz, sub_colors, sub_labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 49 | 50 | search_tree = KDTree(sub_xyz, leaf_size=50) 51 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 52 | with open(kd_tree_file, 'wb') as f: 53 | pickle.dump(search_tree, f) 54 | 55 | proj_idx = np.squeeze(search_tree.query(sub_points, return_distance=False)) 56 | proj_idx = proj_idx.astype(np.int32) 57 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 58 | with open(proj_save, 'wb') as f: 59 | pickle.dump([proj_idx, labels], f) 60 | 61 | else: 62 | full_ply_path = join(original_pc_folder, file_name + '.ply') 63 | write_ply(full_ply_path, (pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8)), 64 | ['x', 'y', 'z', 'red', 'green', 'blue']) 65 | 66 | # Save sub_cloud and KDTree file 67 | sub_xyz, sub_colors = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8), 68 | grid_size=grid_size) 69 | sub_colors = sub_colors / 255.0 70 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 71 | write_ply(sub_ply_file, [sub_xyz, sub_colors], ['x', 'y', 'z', 'red', 'green', 'blue']) 72 | labels = np.zeros(pc.shape[0], dtype=np.uint8) 73 | 74 | search_tree = KDTree(sub_xyz, leaf_size=50) 75 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 76 | with open(kd_tree_file, 'wb') as f: 77 | pickle.dump(search_tree, f) 78 | 79 | proj_idx = np.squeeze(search_tree.query(pc[:, :3].astype(np.float32), return_distance=False)) 80 | proj_idx = proj_idx.astype(np.int32) 81 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 82 | with open(proj_save, 'wb') as f: 83 | pickle.dump([proj_idx, labels], f) 84 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/download_semantic3d.sh: -------------------------------------------------------------------------------- 1 | BASE_DIR=${1-/home/data/semantic3d/original_data} 2 | 3 | # Training data 4 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station1_xyz_intensity_rgb.7z -P $BASE_DIR 5 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station3_xyz_intensity_rgb.7z -P $BASE_DIR 6 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station5_xyz_intensity_rgb.7z -P $BASE_DIR 7 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 8 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station2_xyz_intensity_rgb.7z -P $BASE_DIR/ 9 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station3_xyz_intensity_rgb.7z -P $BASE_DIR 10 | wget -c -N http://semantic3d.net/data/point-clouds/training1/neugasse_station1_xyz_intensity_rgb.7z -P $BASE_DIR 11 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station1_intensity_rgb.7z -P $BASE_DIR 12 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station2_intensity_rgb.7z -P $BASE_DIR 13 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station4_intensity_rgb.7z -P $BASE_DIR/ 14 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station5_intensity_rgb.7z -P $BASE_DIR 15 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station9_intensity_rgb.7z -P $BASE_DIR 16 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg28_station4_intensity_rgb.7z -P $BASE_DIR 17 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station1_xyz_intensity_rgb.7z -P $BASE_DIR 18 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station3_xyz_intensity_rgb.7z -P $BASE_DIR/ 19 | wget -c -N http://semantic3d.net/data/sem8_labels_training.7z -P $BASE_DIR 20 | 21 | 22 | # Test data 23 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/birdfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 24 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station1_intensity_rgb.7z -P $BASE_DIR 25 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station5_xyz_intensity_rgb.7z -P $BASE_DIR 26 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station1_intensity_rgb.7z -P $BASE_DIR 27 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station4_intensity_rgb.7z -P $BASE_DIR 28 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station7_intensity_rgb.7z -P $BASE_DIR 29 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station10_intensity_rgb.7z -P $BASE_DIR 30 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station3_intensity_rgb.7z -P $BASE_DIR 31 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station6_intensity_rgb.7z -P $BASE_DIR 32 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station8_intensity_rgb.7z -P $BASE_DIR 33 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station2_intensity_rgb.7z -P $BASE_DIR 34 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station5_xyz_intensity_rgb.7z -P $BASE_DIR 35 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station1_intensity_rgb.7z -P $BASE_DIR 36 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station3_intensity_rgb.7z -P $BASE_DIR 37 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station6_intensity_rgb.7z -P $BASE_DIR 38 | 39 | # reduced-8 40 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/MarketplaceFeldkirch_Station4_rgb_intensity-reduced.txt.7z -P $BASE_DIR 41 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/StGallenCathedral_station6_rgb_intensity-reduced.txt.7z -P $BASE_DIR 42 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg27_station10_rgb_intensity-reduced.txt.7z -P $BASE_DIR 43 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg28_Station2_rgb_intensity-reduced.txt.7z -P $BASE_DIR 44 | 45 | 46 | 47 | for entry in "$BASE_DIR"/* 48 | do 49 | 7z x "$entry" -o$(dirname "$entry") -y 50 | done 51 | 52 | mv $BASE_DIR/station1_xyz_intensity_rgb.txt $BASE_DIR/neugasse_station1_xyz_intensity_rgb.txt 53 | 54 | for entry in "$BASE_DIR"/*.7z 55 | do 56 | rm "$entry" 57 | done 58 | -------------------------------------------------------------------------------- /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 150 | -------------------------------------------------------------------------------- /tester_Semantic3D.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join 3 | from helper_ply import read_ply, write_ply 4 | import tensorflow as tf 5 | import numpy as np 6 | import time 7 | 8 | 9 | def log_string(out_str, log_out): 10 | log_out.write(out_str + '\n') 11 | log_out.flush() 12 | print(out_str) 13 | 14 | 15 | class ModelTester: 16 | def __init__(self, model, dataset, restore_snap=None): 17 | # Tensorflow Saver definition 18 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 19 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 20 | 21 | # Create a session for running Ops on the Graph. 22 | on_cpu = False 23 | if on_cpu: 24 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 25 | else: 26 | c_proto = tf.ConfigProto() 27 | c_proto.gpu_options.allow_growth = True 28 | self.sess = tf.Session(config=c_proto) 29 | self.sess.run(tf.global_variables_initializer()) 30 | 31 | if restore_snap is not None: 32 | self.saver.restore(self.sess, restore_snap) 33 | print('Model restored from ' + restore_snap) 34 | 35 | # Add a softmax operation for predictions 36 | self.prob_logits = tf.nn.softmax(model.logits) 37 | self.test_probs = [np.zeros((l.data.shape[0], model.config.num_classes), dtype=np.float16) 38 | for l in dataset.input_trees['test']] 39 | 40 | self.log_out = open('log_test_' + dataset.name + '.txt', 'a') 41 | 42 | def test(self, model, dataset, num_votes=100): 43 | 44 | # Smoothing parameter for votes 45 | test_smooth = 0.98 46 | 47 | # Initialise iterator with train data 48 | self.sess.run(dataset.test_init_op) 49 | 50 | # Test saving path 51 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 52 | test_path = join('test', saving_path.split('/')[-1]) 53 | makedirs(test_path) if not exists(test_path) else None 54 | makedirs(join(test_path, 'predictions')) if not exists(join(test_path, 'predictions')) else None 55 | makedirs(join(test_path, 'probs')) if not exists(join(test_path, 'probs')) else None 56 | 57 | ##################### 58 | # Network predictions 59 | ##################### 60 | 61 | step_id = 0 62 | epoch_id = 0 63 | last_min = -0.5 64 | 65 | while last_min < num_votes: 66 | 67 | try: 68 | ops = (self.prob_logits, 69 | model.labels, 70 | model.inputs['input_inds'], 71 | model.inputs['cloud_inds']) 72 | 73 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 74 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 75 | model.config.num_classes]) 76 | 77 | for j in range(np.shape(stacked_probs)[0]): 78 | probs = stacked_probs[j, :, :] 79 | inds = point_idx[j, :] 80 | c_i = cloud_idx[j][0] 81 | self.test_probs[c_i][inds] = test_smooth * self.test_probs[c_i][inds] + (1 - test_smooth) * probs 82 | step_id += 1 83 | log_string('Epoch {:3d}, step {:3d}. min possibility = {:.1f}'.format(epoch_id, step_id, np.min( 84 | dataset.min_possibility['test'])), self.log_out) 85 | 86 | except tf.errors.OutOfRangeError: 87 | 88 | # Save predicted cloud 89 | new_min = np.min(dataset.min_possibility['test']) 90 | log_string('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.log_out) 91 | 92 | if last_min + 4 < new_min: 93 | 94 | print('Saving clouds') 95 | 96 | # Update last min 97 | last_min = new_min 98 | 99 | # Project predictions 100 | print('\nReproject Vote #{:d}'.format(int(np.floor(new_min)))) 101 | t1 = time.time() 102 | files = dataset.test_files 103 | i_test = 0 104 | for i, file_path in enumerate(files): 105 | # Get file 106 | points = self.load_evaluation_points(file_path) 107 | points = points.astype(np.float16) 108 | 109 | # Reproject probs 110 | probs = np.zeros(shape=[np.shape(points)[0], 8], dtype=np.float16) 111 | proj_index = dataset.test_proj[i_test] 112 | 113 | probs = self.test_probs[i_test][proj_index, :] 114 | 115 | # Insert false columns for ignored labels 116 | probs2 = probs 117 | for l_ind, label_value in enumerate(dataset.label_values): 118 | if label_value in dataset.ignored_labels: 119 | probs2 = np.insert(probs2, l_ind, 0, axis=1) 120 | 121 | # Get the predicted labels 122 | preds = dataset.label_values[np.argmax(probs2, axis=1)].astype(np.uint8) 123 | 124 | # Save plys 125 | cloud_name = file_path.split('/')[-1] 126 | 127 | # Save acsii preds 128 | ascii_name = join(test_path, 'predictions', dataset.ascii_files[cloud_name]) 129 | np.savetxt(ascii_name, preds, fmt='%d') 130 | log_string(ascii_name + ' has saved', self.log_out) 131 | i_test += 1 132 | 133 | t2 = time.time() 134 | print('Done in {:.1f} s\n'.format(t2 - t1)) 135 | self.sess.close() 136 | return 137 | 138 | self.sess.run(dataset.test_init_op) 139 | epoch_id += 1 140 | step_id = 0 141 | continue 142 | return 143 | 144 | @staticmethod 145 | def load_evaluation_points(file_path): 146 | data = read_ply(file_path) 147 | return np.vstack((data['x'], data['y'], data['z'])).T 148 | -------------------------------------------------------------------------------- /utils/semantic-kitti.yaml: -------------------------------------------------------------------------------- 1 | # This file is covered by the LICENSE file in the root of this project. 2 | labels: 3 | 0 : "unlabeled" 4 | 1 : "outlier" 5 | 10: "car" 6 | 11: "bicycle" 7 | 13: "bus" 8 | 15: "motorcycle" 9 | 16: "on-rails" 10 | 18: "truck" 11 | 20: "other-vehicle" 12 | 30: "person" 13 | 31: "bicyclist" 14 | 32: "motorcyclist" 15 | 40: "road" 16 | 44: "parking" 17 | 48: "sidewalk" 18 | 49: "other-ground" 19 | 50: "building" 20 | 51: "fence" 21 | 52: "other-structure" 22 | 60: "lane-marking" 23 | 70: "vegetation" 24 | 71: "trunk" 25 | 72: "terrain" 26 | 80: "pole" 27 | 81: "traffic-sign" 28 | 99: "other-object" 29 | 252: "moving-car" 30 | 253: "moving-bicyclist" 31 | 254: "moving-person" 32 | 255: "moving-motorcyclist" 33 | 256: "moving-on-rails" 34 | 257: "moving-bus" 35 | 258: "moving-truck" 36 | 259: "moving-other-vehicle" 37 | color_map: # bgr 38 | 0 : [0, 0, 0] 39 | 1 : [0, 0, 255] 40 | 10: [245, 150, 100] 41 | 11: [245, 230, 100] 42 | 13: [250, 80, 100] 43 | 15: [150, 60, 30] 44 | 16: [255, 0, 0] 45 | 18: [180, 30, 80] 46 | 20: [255, 0, 0] 47 | 30: [30, 30, 255] 48 | 31: [200, 40, 255] 49 | 32: [90, 30, 150] 50 | 40: [255, 0, 255] 51 | 44: [255, 150, 255] 52 | 48: [75, 0, 75] 53 | 49: [75, 0, 175] 54 | 50: [0, 200, 255] 55 | 51: [50, 120, 255] 56 | 52: [0, 150, 255] 57 | 60: [170, 255, 150] 58 | 70: [0, 175, 0] 59 | 71: [0, 60, 135] 60 | 72: [80, 240, 150] 61 | 80: [150, 240, 255] 62 | 81: [0, 0, 255] 63 | 99: [255, 255, 50] 64 | 252: [245, 150, 100] 65 | 256: [255, 0, 0] 66 | 253: [200, 40, 255] 67 | 254: [30, 30, 255] 68 | 255: [90, 30, 150] 69 | 257: [250, 80, 100] 70 | 258: [180, 30, 80] 71 | 259: [255, 0, 0] 72 | content: # as a ratio with the total number of points 73 | 0: 0.018889854628292943 74 | 1: 0.0002937197336781505 75 | 10: 0.040818519255974316 76 | 11: 0.00016609538710764618 77 | 13: 2.7879693665067774e-05 78 | 15: 0.00039838616015114444 79 | 16: 0.0 80 | 18: 0.0020633612104619787 81 | 20: 0.0016218197275284021 82 | 30: 0.00017698551338515307 83 | 31: 1.1065903904919655e-08 84 | 32: 5.532951952459828e-09 85 | 40: 0.1987493871255525 86 | 44: 0.014717169549888214 87 | 48: 0.14392298360372 88 | 49: 0.0039048553037472045 89 | 50: 0.1326861944777486 90 | 51: 0.0723592229456223 91 | 52: 0.002395131480328884 92 | 60: 4.7084144280367186e-05 93 | 70: 0.26681502148037506 94 | 71: 0.006035012012626033 95 | 72: 0.07814222006271769 96 | 80: 0.002855498193863172 97 | 81: 0.0006155958086189918 98 | 99: 0.009923127583046915 99 | 252: 0.001789309418528068 100 | 253: 0.00012709999297008662 101 | 254: 0.00016059776092534436 102 | 255: 3.745553104802113e-05 103 | 256: 0.0 104 | 257: 0.00011351574470342043 105 | 258: 0.00010157861367183268 106 | 259: 4.3840131989471124e-05 107 | # classes that are indistinguishable from single scan or inconsistent in 108 | # ground truth are mapped to their closest equivalent 109 | learning_map: 110 | 0 : 0 # "unlabeled" 111 | 1 : 0 # "outlier" mapped to "unlabeled" --------------------------mapped 112 | 10: 1 # "car" 113 | 11: 2 # "bicycle" 114 | 13: 5 # "bus" mapped to "other-vehicle" --------------------------mapped 115 | 15: 3 # "motorcycle" 116 | 16: 5 # "on-rails" mapped to "other-vehicle" ---------------------mapped 117 | 18: 4 # "truck" 118 | 20: 5 # "other-vehicle" 119 | 30: 6 # "person" 120 | 31: 7 # "bicyclist" 121 | 32: 8 # "motorcyclist" 122 | 40: 9 # "road" 123 | 44: 10 # "parking" 124 | 48: 11 # "sidewalk" 125 | 49: 12 # "other-ground" 126 | 50: 13 # "building" 127 | 51: 14 # "fence" 128 | 52: 0 # "other-structure" mapped to "unlabeled" ------------------mapped 129 | 60: 9 # "lane-marking" to "road" ---------------------------------mapped 130 | 70: 15 # "vegetation" 131 | 71: 16 # "trunk" 132 | 72: 17 # "terrain" 133 | 80: 18 # "pole" 134 | 81: 19 # "traffic-sign" 135 | 99: 0 # "other-object" to "unlabeled" ----------------------------mapped 136 | 252: 1 # "moving-car" to "car" ------------------------------------mapped 137 | 253: 7 # "moving-bicyclist" to "bicyclist" ------------------------mapped 138 | 254: 6 # "moving-person" to "person" ------------------------------mapped 139 | 255: 8 # "moving-motorcyclist" to "motorcyclist" ------------------mapped 140 | 256: 5 # "moving-on-rails" mapped to "other-vehicle" --------------mapped 141 | 257: 5 # "moving-bus" mapped to "other-vehicle" -------------------mapped 142 | 258: 4 # "moving-truck" to "truck" --------------------------------mapped 143 | 259: 5 # "moving-other"-vehicle to "other-vehicle" ----------------mapped 144 | learning_map_inv: # inverse of previous map 145 | 0: 0 # "unlabeled", and others ignored 146 | 1: 10 # "car" 147 | 2: 11 # "bicycle" 148 | 3: 15 # "motorcycle" 149 | 4: 18 # "truck" 150 | 5: 20 # "other-vehicle" 151 | 6: 30 # "person" 152 | 7: 31 # "bicyclist" 153 | 8: 32 # "motorcyclist" 154 | 9: 40 # "road" 155 | 10: 44 # "parking" 156 | 11: 48 # "sidewalk" 157 | 12: 49 # "other-ground" 158 | 13: 50 # "building" 159 | 14: 51 # "fence" 160 | 15: 70 # "vegetation" 161 | 16: 71 # "trunk" 162 | 17: 72 # "terrain" 163 | 18: 80 # "pole" 164 | 19: 81 # "traffic-sign" 165 | learning_ignore: # Ignore classes 166 | 0: True # "unlabeled", and others ignored 167 | 1: False # "car" 168 | 2: False # "bicycle" 169 | 3: False # "motorcycle" 170 | 4: False # "truck" 171 | 5: False # "other-vehicle" 172 | 6: False # "person" 173 | 7: False # "bicyclist" 174 | 8: False # "motorcyclist" 175 | 9: False # "road" 176 | 10: False # "parking" 177 | 11: False # "sidewalk" 178 | 12: False # "other-ground" 179 | 13: False # "building" 180 | 14: False # "fence" 181 | 15: False # "vegetation" 182 | 16: False # "trunk" 183 | 17: False # "terrain" 184 | 18: False # "pole" 185 | 19: False # "traffic-sign" 186 | split: # sequence numbers 187 | train: 188 | - 0 189 | - 1 190 | - 2 191 | - 3 192 | - 4 193 | - 5 194 | - 6 195 | - 7 196 | - 9 197 | - 10 198 | valid: 199 | - 8 200 | test: 201 | - 11 202 | - 12 203 | - 13 204 | - 14 205 | - 15 206 | - 16 207 | - 17 208 | - 18 209 | - 19 210 | - 20 211 | - 21 212 | -------------------------------------------------------------------------------- /tester_S3DIS.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join 3 | from helper_ply import write_ply 4 | from sklearn.metrics import confusion_matrix 5 | from helper_tool import DataProcessing as DP 6 | import tensorflow as tf 7 | import numpy as np 8 | import time 9 | 10 | 11 | def log_out(out_str, log_f_out): 12 | log_f_out.write(out_str + '\n') 13 | log_f_out.flush() 14 | print(out_str) 15 | 16 | 17 | class ModelTester: 18 | def __init__(self, model, dataset, restore_snap=None): 19 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 20 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 21 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 22 | 23 | # Create a session for running Ops on the Graph. 24 | on_cpu = False 25 | if on_cpu: 26 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 27 | else: 28 | c_proto = tf.ConfigProto() 29 | c_proto.gpu_options.allow_growth = True 30 | self.sess = tf.Session(config=c_proto) 31 | self.sess.run(tf.global_variables_initializer()) 32 | 33 | # Load trained model 34 | if restore_snap is not None: 35 | self.saver.restore(self.sess, restore_snap) 36 | print('Model restored from ' + restore_snap) 37 | 38 | self.prob_logits = tf.nn.softmax(model.logits) 39 | 40 | # Initiate global prediction over all test clouds 41 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 42 | for l in dataset.input_labels['validation']] 43 | 44 | def test(self, model, dataset, num_votes=100): 45 | 46 | # Smoothing parameter for votes 47 | test_smooth = 0.95 48 | 49 | # Initialise iterator with validation/test data 50 | self.sess.run(dataset.val_init_op) 51 | 52 | # Number of points per class in validation set 53 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 54 | i = 0 55 | for label_val in dataset.label_values: 56 | if label_val not in dataset.ignored_labels: 57 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 58 | i += 1 59 | 60 | # Test saving path 61 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 62 | test_path = join('test', saving_path.split('/')[-1]) 63 | makedirs(test_path) if not exists(test_path) else None 64 | makedirs(join(test_path, 'val_preds')) if not exists(join(test_path, 'val_preds')) else None 65 | 66 | step_id = 0 67 | epoch_id = 0 68 | last_min = -0.5 69 | 70 | while last_min < num_votes: 71 | try: 72 | ops = (self.prob_logits, 73 | model.labels, 74 | model.inputs['input_inds'], 75 | model.inputs['cloud_inds'], 76 | ) 77 | 78 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 79 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 80 | acc = correct / float(np.prod(np.shape(stacked_labels))) 81 | print('step' + str(step_id) + ' acc: ' + str(acc)) 82 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 83 | model.config.num_classes]) 84 | 85 | for j in range(np.shape(stacked_probs)[0]): 86 | probs = stacked_probs[j, :, :] 87 | p_idx = point_idx[j, :] 88 | c_i = cloud_idx[j][0] 89 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 90 | step_id += 1 91 | 92 | except tf.errors.OutOfRangeError: 93 | 94 | new_min = np.min(dataset.min_possibility['validation']) 95 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 96 | 97 | if last_min + 1 < new_min: 98 | 99 | # Update last_min 100 | last_min += 1 101 | 102 | # Show vote results (On subcloud so it is not the good values here) 103 | log_out('\nConfusion on sub clouds', self.Log_file) 104 | confusion_list = [] 105 | 106 | num_val = len(dataset.input_labels['validation']) 107 | 108 | for i_test in range(num_val): 109 | probs = self.test_probs[i_test] 110 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 111 | labels = dataset.input_labels['validation'][i_test] 112 | 113 | # Confs 114 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 115 | 116 | # Regroup confusions 117 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 118 | 119 | # Rescale with the right number of point per class 120 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 121 | 122 | # Compute IoUs 123 | IoUs = DP.IoU_from_confusions(C) 124 | m_IoU = np.mean(IoUs) 125 | s = '{:5.2f} | '.format(100 * m_IoU) 126 | for IoU in IoUs: 127 | s += '{:5.2f} '.format(100 * IoU) 128 | log_out(s + '\n', self.Log_file) 129 | 130 | if int(np.ceil(new_min)) % 1 == 0: 131 | 132 | # Project predictions 133 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 134 | proj_probs_list = [] 135 | 136 | for i_val in range(num_val): 137 | # Reproject probs back to the evaluations points 138 | proj_idx = dataset.val_proj[i_val] 139 | probs = self.test_probs[i_val][proj_idx, :] 140 | proj_probs_list += [probs] 141 | 142 | # Show vote results 143 | log_out('Confusion on full clouds', self.Log_file) 144 | confusion_list = [] 145 | for i_test in range(num_val): 146 | # Get the predicted labels 147 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 148 | 149 | # Confusion 150 | labels = dataset.val_labels[i_test] 151 | acc = np.sum(preds == labels) / len(labels) 152 | log_out(dataset.input_names['validation'][i_test] + ' Acc: ' + str(acc), self.Log_file) 153 | 154 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 155 | name = dataset.input_names['validation'][i_test] + '.ply' 156 | write_ply(join(test_path, 'val_preds', name), [preds, labels], ['pred', 'label']) 157 | 158 | # Regroup confusions 159 | C = np.sum(np.stack(confusion_list), axis=0) 160 | 161 | IoUs = DP.IoU_from_confusions(C) 162 | m_IoU = np.mean(IoUs) 163 | s = '{:5.2f} | '.format(100 * m_IoU) 164 | for IoU in IoUs: 165 | s += '{:5.2f} '.format(100 * IoU) 166 | log_out('-' * len(s), self.Log_file) 167 | log_out(s, self.Log_file) 168 | log_out('-' * len(s) + '\n', self.Log_file) 169 | print('finished \n') 170 | self.sess.close() 171 | return 172 | 173 | self.sess.run(dataset.val_init_op) 174 | epoch_id += 1 175 | step_id = 0 176 | continue 177 | 178 | return 179 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /tester_SemanticKITTI.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join, isfile, dirname, abspath 3 | from helper_tool import DataProcessing as DP 4 | from sklearn.metrics import confusion_matrix 5 | import tensorflow as tf 6 | import numpy as np 7 | import yaml 8 | import pickle 9 | 10 | BASE_DIR = dirname(abspath(__file__)) 11 | 12 | data_config = join(BASE_DIR, 'utils', 'semantic-kitti.yaml') 13 | DATA = yaml.safe_load(open(data_config, 'r')) 14 | remap_dict = DATA['learning_map_inv'] 15 | 16 | # make lookup table for mapping 17 | max_key = max(remap_dict.keys()) 18 | remap_lut = np.zeros((max_key + 100), dtype=np.int32) 19 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values()) 20 | 21 | remap_dict_val = DATA['learning_map'] 22 | max_key = max(remap_dict_val.keys()) 23 | remap_lut_val = np.zeros((max_key + 100), dtype=np.int32) 24 | remap_lut_val[list(remap_dict_val.keys())] = list(remap_dict_val.values()) 25 | 26 | 27 | def log_out(out_str, f_out): 28 | f_out.write(out_str + '\n') 29 | f_out.flush() 30 | print(out_str) 31 | 32 | 33 | class ModelTester: 34 | def __init__(self, model, dataset, restore_snap=None): 35 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 36 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 37 | self.Log_file = open('log_test_' + dataset.name + '.txt', 'a') 38 | 39 | # Create a session for running Ops on the Graph. 40 | on_cpu = False 41 | if on_cpu: 42 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 43 | else: 44 | c_proto = tf.ConfigProto() 45 | c_proto.gpu_options.allow_growth = True 46 | self.sess = tf.Session(config=c_proto) 47 | self.sess.run(tf.global_variables_initializer()) 48 | 49 | # Name of the snapshot to restore to (None if you want to start from beginning) 50 | if restore_snap is not None: 51 | self.saver.restore(self.sess, restore_snap) 52 | print('Model restored from ' + restore_snap) 53 | 54 | self.prob_logits = tf.nn.softmax(model.logits) 55 | self.test_probs = 0 56 | self.idx = 0 57 | 58 | def test(self, model, dataset): 59 | 60 | # Initialise iterator with train data 61 | self.sess.run(dataset.test_init_op) 62 | self.test_probs = [np.zeros(shape=[len(l), model.config.num_classes], dtype=np.float16) 63 | for l in dataset.possibility] 64 | 65 | test_path = join('test', 'sequences') 66 | makedirs(test_path) if not exists(test_path) else None 67 | save_path = join(test_path, dataset.test_scan_number, 'predictions') 68 | makedirs(save_path) if not exists(save_path) else None 69 | test_smooth = 0.98 70 | epoch_ind = 0 71 | 72 | while True: 73 | try: 74 | ops = (self.prob_logits, 75 | model.labels, 76 | model.inputs['input_inds'], 77 | model.inputs['cloud_inds']) 78 | stacked_probs, labels, point_inds, cloud_inds = self.sess.run(ops, {model.is_training: False}) 79 | if self.idx % 10 == 0: 80 | print('step ' + str(self.idx)) 81 | self.idx += 1 82 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, 83 | model.config.num_points, 84 | model.config.num_classes]) 85 | for j in range(np.shape(stacked_probs)[0]): 86 | probs = stacked_probs[j, :, :] 87 | inds = point_inds[j, :] 88 | c_i = cloud_inds[j][0] 89 | self.test_probs[c_i][inds] = test_smooth * self.test_probs[c_i][inds] + (1 - test_smooth) * probs 90 | 91 | except tf.errors.OutOfRangeError: 92 | new_min = np.min(dataset.min_possibility) 93 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_ind, new_min), self.Log_file) 94 | if np.min(dataset.min_possibility) > 0.5: # 0.5 95 | log_out(' Min possibility = {:.1f}'.format(np.min(dataset.min_possibility)), self.Log_file) 96 | print('\nReproject Vote #{:d}'.format(int(np.floor(new_min)))) 97 | 98 | # For validation set 99 | num_classes = 19 100 | gt_classes = [0 for _ in range(num_classes)] 101 | positive_classes = [0 for _ in range(num_classes)] 102 | true_positive_classes = [0 for _ in range(num_classes)] 103 | val_total_correct = 0 104 | val_total_seen = 0 105 | 106 | for j in range(len(self.test_probs)): 107 | test_file_name = dataset.test_list[j] 108 | frame = test_file_name.split('/')[-1][:-4] 109 | proj_path = join(dataset.dataset_path, dataset.test_scan_number, 'proj') 110 | proj_file = join(proj_path, str(frame) + '_proj.pkl') 111 | if isfile(proj_file): 112 | with open(proj_file, 'rb') as f: 113 | proj_inds = pickle.load(f) 114 | probs = self.test_probs[j][proj_inds[0], :] 115 | pred = np.argmax(probs, 1) 116 | if dataset.test_scan_number == '08': 117 | label_path = join(dirname(dataset.dataset_path), 'sequences', dataset.test_scan_number, 118 | 'labels') 119 | label_file = join(label_path, str(frame) + '.label') 120 | labels = DP.load_label_kitti(label_file, remap_lut_val) 121 | invalid_idx = np.where(labels == 0)[0] 122 | labels_valid = np.delete(labels, invalid_idx) 123 | pred_valid = np.delete(pred, invalid_idx) 124 | labels_valid = labels_valid - 1 125 | correct = np.sum(pred_valid == labels_valid) 126 | val_total_correct += correct 127 | val_total_seen += len(labels_valid) 128 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, num_classes, 1)) 129 | gt_classes += np.sum(conf_matrix, axis=1) 130 | positive_classes += np.sum(conf_matrix, axis=0) 131 | true_positive_classes += np.diagonal(conf_matrix) 132 | else: 133 | store_path = join(test_path, dataset.test_scan_number, 'predictions', 134 | str(frame) + '.label') 135 | pred = pred + 1 136 | pred = pred.astype(np.uint32) 137 | upper_half = pred >> 16 # get upper half for instances 138 | lower_half = pred & 0xFFFF # get lower half for semantics 139 | lower_half = remap_lut[lower_half] # do the remapping of semantics 140 | pred = (upper_half << 16) + lower_half # reconstruct full label 141 | pred = pred.astype(np.uint32) 142 | pred.tofile(store_path) 143 | log_out(str(dataset.test_scan_number) + ' finished', self.Log_file) 144 | if dataset.test_scan_number == '08': 145 | iou_list = [] 146 | for n in range(0, num_classes, 1): 147 | iou = true_positive_classes[n] / float( 148 | gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 149 | iou_list.append(iou) 150 | mean_iou = sum(iou_list) / float(num_classes) 151 | 152 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 153 | log_out('mean IoU: {}'.format(mean_iou), self.Log_file) 154 | 155 | mean_iou = 100 * mean_iou 156 | print('Mean IoU = {:.1f}%'.format(mean_iou)) 157 | s = '{:5.2f} | '.format(mean_iou) 158 | for IoU in iou_list: 159 | s += '{:5.2f} '.format(100 * IoU) 160 | print('-' * len(s)) 161 | print(s) 162 | print('-' * len(s) + '\n') 163 | self.sess.close() 164 | return 165 | self.sess.run(dataset.test_init_op) 166 | epoch_ind += 1 167 | continue 168 | -------------------------------------------------------------------------------- /utils/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /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 | } 287 | -------------------------------------------------------------------------------- /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 | while b'end_header' not in line and line != b'': 91 | line = plyfile.readline() 92 | 93 | # Find point element 94 | if b'element vertex' in line: 95 | current_element = 'vertex' 96 | line = line.split() 97 | num_points = int(line[2]) 98 | 99 | elif b'element face' in line: 100 | current_element = 'face' 101 | line = line.split() 102 | num_faces = int(line[2]) 103 | 104 | elif b'property' in line: 105 | if current_element == 'vertex': 106 | line = line.split() 107 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 108 | elif current_element == 'vertex': 109 | if not line.startswith('property list uchar int'): 110 | raise ValueError('Unsupported faces property : ' + line) 111 | 112 | return num_points, num_faces, vertex_properties 113 | 114 | 115 | def read_ply(filename, triangular_mesh=False): 116 | """ 117 | Read ".ply" files 118 | 119 | Parameters 120 | ---------- 121 | filename : string 122 | the name of the file to read. 123 | 124 | Returns 125 | ------- 126 | result : array 127 | data stored in the file 128 | 129 | Examples 130 | -------- 131 | Store data in file 132 | 133 | >>> points = np.random.rand(5, 3) 134 | >>> values = np.random.randint(2, size=10) 135 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 136 | 137 | Read the file 138 | 139 | >>> data = read_ply('example.ply') 140 | >>> values = data['values'] 141 | array([0, 0, 1, 1, 0]) 142 | 143 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 144 | array([[ 0.466 0.595 0.324] 145 | [ 0.538 0.407 0.654] 146 | [ 0.850 0.018 0.988] 147 | [ 0.395 0.394 0.363] 148 | [ 0.873 0.996 0.092]]) 149 | 150 | """ 151 | 152 | with open(filename, 'rb') as plyfile: 153 | 154 | # Check if the file start with ply 155 | if b'ply' not in plyfile.readline(): 156 | raise ValueError('The file does not start whith the word ply') 157 | 158 | # get binary_little/big or ascii 159 | fmt = plyfile.readline().split()[1].decode() 160 | if fmt == "ascii": 161 | raise ValueError('The file is not binary') 162 | 163 | # get extension for building the numpy dtypes 164 | ext = valid_formats[fmt] 165 | 166 | # PointCloud reader vs mesh reader 167 | if triangular_mesh: 168 | 169 | # Parse header 170 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 171 | 172 | # Get point data 173 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 174 | 175 | # Get face data 176 | face_properties = [('k', ext + 'u1'), 177 | ('v1', ext + 'i4'), 178 | ('v2', ext + 'i4'), 179 | ('v3', ext + 'i4')] 180 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 181 | 182 | # Return vertex data and concatenated faces 183 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 184 | data = [vertex_data, faces] 185 | 186 | else: 187 | 188 | # Parse header 189 | num_points, properties = parse_header(plyfile, ext) 190 | 191 | # Get data 192 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 193 | 194 | return data 195 | 196 | 197 | def header_properties(field_list, field_names): 198 | # List of lines to write 199 | lines = [] 200 | 201 | # First line describing element vertex 202 | lines.append('element vertex %d' % field_list[0].shape[0]) 203 | 204 | # Properties lines 205 | i = 0 206 | for fields in field_list: 207 | for field in fields.T: 208 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 209 | i += 1 210 | 211 | return lines 212 | 213 | 214 | def write_ply(filename, field_list, field_names, triangular_faces=None): 215 | """ 216 | Write ".ply" files 217 | 218 | Parameters 219 | ---------- 220 | filename : string 221 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 222 | file name if it does no already have one. 223 | 224 | field_list : list, tuple, numpy array 225 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 226 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 227 | as one field. 228 | 229 | field_names : list 230 | the name of each fields as a list of strings. Has to be the same length as the number of 231 | fields. 232 | 233 | Examples 234 | -------- 235 | >>> points = np.random.rand(10, 3) 236 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 237 | 238 | >>> values = np.random.randint(2, size=10) 239 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 240 | 241 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 242 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', 'values'] 243 | >>> write_ply('example3.ply', [points, colors, values], field_names) 244 | 245 | """ 246 | 247 | # Format list input to the right form 248 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 249 | for i, field in enumerate(field_list): 250 | if field.ndim < 2: 251 | field_list[i] = field.reshape(-1, 1) 252 | if field.ndim > 2: 253 | print('fields have more than 2 dimensions') 254 | return False 255 | 256 | # check all fields have the same number of data 257 | n_points = [field.shape[0] for field in field_list] 258 | if not np.all(np.equal(n_points, n_points[0])): 259 | print('wrong field dimensions') 260 | return False 261 | 262 | # Check if field_names and field_list have same nb of column 263 | n_fields = np.sum([field.shape[1] for field in field_list]) 264 | if (n_fields != len(field_names)): 265 | print('wrong number of field names') 266 | return False 267 | 268 | # Add extension if not there 269 | if not filename.endswith('.ply'): 270 | filename += '.ply' 271 | 272 | # open in text mode to write the header 273 | with open(filename, 'w') as plyfile: 274 | 275 | # First magical word 276 | header = ['ply'] 277 | 278 | # Encoding format 279 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 280 | 281 | # Points properties description 282 | header.extend(header_properties(field_list, field_names)) 283 | 284 | # Add faces if needded 285 | if triangular_faces is not None: 286 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 287 | header.append('property list uchar int vertex_indices') 288 | 289 | # End of header 290 | header.append('end_header') 291 | 292 | # Write all lines 293 | for line in header: 294 | plyfile.write("%s\n" % line) 295 | 296 | # open in binary/append to use tofile 297 | with open(filename, 'ab') as plyfile: 298 | 299 | # Create a structured array 300 | i = 0 301 | type_list = [] 302 | for fields in field_list: 303 | for field in fields.T: 304 | type_list += [(field_names[i], field.dtype.str)] 305 | i += 1 306 | data = np.empty(field_list[0].shape[0], dtype=type_list) 307 | i = 0 308 | for fields in field_list: 309 | for field in fields.T: 310 | data[field_names[i]] = field 311 | i += 1 312 | 313 | data.tofile(plyfile) 314 | 315 | if triangular_faces is not None: 316 | triangular_faces = triangular_faces.astype(np.int32) 317 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 318 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 319 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 320 | data['0'] = triangular_faces[:, 0] 321 | data['1'] = triangular_faces[:, 1] 322 | data['2'] = triangular_faces[:, 2] 323 | data.tofile(plyfile) 324 | 325 | return True 326 | 327 | 328 | def describe_element(name, df): 329 | """ Takes the columns of the dataframe and builds a ply-like description 330 | 331 | Parameters 332 | ---------- 333 | name: str 334 | df: pandas DataFrame 335 | 336 | Returns 337 | ------- 338 | element: list[str] 339 | """ 340 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 341 | element = ['element ' + name + ' ' + str(len(df))] 342 | 343 | if name == 'face': 344 | element.append("property list uchar int points_indices") 345 | 346 | else: 347 | for i in range(len(df.columns)): 348 | # get first letter of dtype to infer format 349 | f = property_formats[str(df.dtypes[i])[0]] 350 | element.append('property ' + f + ' ' + df.columns.values[i]) 351 | 352 | return element 353 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /main_SemanticKITTI.py: -------------------------------------------------------------------------------- 1 | from helper_tool import DataProcessing as DP 2 | from helper_tool import ConfigSemanticKITTI as cfg 3 | from helper_tool import Plot 4 | from os.path import join 5 | from BAF_LAC import Network 6 | from tester_SemanticKITTI import ModelTester 7 | import tensorflow as tf 8 | import numpy as np 9 | import os, argparse, pickle 10 | 11 | 12 | class SemanticKITTI: 13 | def __init__(self, test_id): 14 | self.name = 'SemanticKITTI' 15 | self.dataset_path = '/home/data/semantic_kitti/dataset/sequences_0.06' 16 | self.label_to_names = {0: 'unlabeled', 17 | 1: 'car', 18 | 2: 'bicycle', 19 | 3: 'motorcycle', 20 | 4: 'truck', 21 | 5: 'other-vehicle', 22 | 6: 'person', 23 | 7: 'bicyclist', 24 | 8: 'motorcyclist', 25 | 9: 'road', 26 | 10: 'parking', 27 | 11: 'sidewalk', 28 | 12: 'other-ground', 29 | 13: 'building', 30 | 14: 'fence', 31 | 15: 'vegetation', 32 | 16: 'trunk', 33 | 17: 'terrain', 34 | 18: 'pole', 35 | 19: 'traffic-sign'} 36 | self.num_classes = len(self.label_to_names) 37 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 38 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 39 | self.ignored_labels = np.sort([0]) 40 | 41 | self.val_split = '08' 42 | 43 | self.seq_list = np.sort(os.listdir(self.dataset_path)) 44 | self.test_scan_number = str(test_id) 45 | self.train_list, self.val_list, self.test_list = DP.get_file_list(self.dataset_path, 46 | self.test_scan_number) 47 | self.train_list = DP.shuffle_list(self.train_list) 48 | self.val_list = DP.shuffle_list(self.val_list) 49 | 50 | self.possibility = [] 51 | self.min_possibility = [] 52 | 53 | # Generate the input data flow 54 | def get_batch_gen(self, split): 55 | if split == 'training': 56 | num_per_epoch = int(len(self.train_list) / cfg.batch_size) * cfg.batch_size 57 | path_list = self.train_list 58 | elif split == 'validation': 59 | num_per_epoch = int(len(self.val_list) / cfg.val_batch_size) * cfg.val_batch_size 60 | cfg.val_steps = int(len(self.val_list) / cfg.batch_size) 61 | path_list = self.val_list 62 | elif split == 'test': 63 | num_per_epoch = int(len(self.test_list) / cfg.val_batch_size) * cfg.val_batch_size * 4 64 | path_list = self.test_list 65 | for test_file_name in path_list: 66 | points = np.load(test_file_name) 67 | self.possibility += [np.random.rand(points.shape[0]) * 1e-3] 68 | self.min_possibility += [float(np.min(self.possibility[-1]))] 69 | 70 | def spatially_regular_gen(): 71 | # Generator loop 72 | for i in range(num_per_epoch): 73 | if split != 'test': 74 | cloud_ind = i 75 | pc_path = path_list[cloud_ind] 76 | pc, tree, labels = self.get_data(pc_path) 77 | # Crop a small point cloud 78 | pick_idx = np.random.choice(len(pc), 1) 79 | selected_pc, selected_labels, selected_idx = self.crop_pc(pc, labels, tree, pick_idx) 80 | else: 81 | cloud_ind = int(np.argmin(self.min_possibility)) 82 | pick_idx = np.argmin(self.possibility[cloud_ind]) 83 | pc_path = path_list[cloud_ind] 84 | pc, tree, labels = self.get_data(pc_path) 85 | selected_pc, selected_labels, selected_idx = self.crop_pc(pc, labels, tree, pick_idx) 86 | 87 | # Update the possibility of the selected pc 88 | dists = np.sum(np.square((selected_pc - pc[pick_idx]).astype(np.float32)), axis=1) 89 | delta = np.square(1 - dists / np.max(dists)) 90 | self.possibility[cloud_ind][selected_idx] += delta 91 | self.min_possibility[cloud_ind] = np.min(self.possibility[cloud_ind]) 92 | 93 | if True: 94 | yield (selected_pc.astype(np.float32), 95 | selected_labels.astype(np.int32), 96 | selected_idx.astype(np.int32), 97 | np.array([cloud_ind], dtype=np.int32)) 98 | 99 | gen_func = spatially_regular_gen 100 | gen_types = (tf.float32, tf.int32, tf.int32, tf.int32) 101 | gen_shapes = ([None, 3], [None], [None], [None]) 102 | 103 | return gen_func, gen_types, gen_shapes 104 | 105 | def get_data(self, file_path): 106 | seq_id = file_path.split('/')[-3] 107 | frame_id = file_path.split('/')[-1][:-4] 108 | kd_tree_path = join(self.dataset_path, seq_id, 'KDTree', frame_id + '.pkl') 109 | # Read pkl with search tree 110 | with open(kd_tree_path, 'rb') as f: 111 | search_tree = pickle.load(f) 112 | points = np.array(search_tree.data, copy=False) 113 | # Load labels 114 | if int(seq_id) >= 11: 115 | labels = np.zeros(np.shape(points)[0], dtype=np.uint8) 116 | else: 117 | label_path = join(self.dataset_path, seq_id, 'labels', frame_id + '.npy') 118 | labels = np.squeeze(np.load(label_path)) 119 | return points, search_tree, labels 120 | 121 | @staticmethod 122 | def crop_pc(points, labels, search_tree, pick_idx): 123 | # Crop a fixed size point cloud for training 124 | center_point = points[pick_idx, :].reshape(1, -1) 125 | select_idx = search_tree.query(center_point, k=cfg.num_points)[1][0] 126 | select_idx = DP.shuffle_idx(select_idx) 127 | select_points = points[select_idx] 128 | select_labels = labels[select_idx] 129 | return select_points, select_labels, select_idx 130 | 131 | @staticmethod 132 | def get_tf_mapping2(): 133 | 134 | def tf_map(batch_pc, batch_label, batch_pc_idx, batch_cloud_idx): 135 | features = batch_pc 136 | input_points = [] 137 | input_neighbors = [] 138 | input_pools = [] 139 | input_up_samples = [] 140 | 141 | for i in range(cfg.num_layers): 142 | neighbour_idx = tf.py_func(DP.knn_search, [batch_pc, batch_pc, cfg.k_n], tf.int32) 143 | sub_points = batch_pc[:, :tf.shape(batch_pc)[1] // cfg.sub_sampling_ratio[i], :] 144 | pool_i = neighbour_idx[:, :tf.shape(batch_pc)[1] // cfg.sub_sampling_ratio[i], :] 145 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_pc, 1], tf.int32) 146 | input_points.append(batch_pc) 147 | input_neighbors.append(neighbour_idx) 148 | input_pools.append(pool_i) 149 | input_up_samples.append(up_i) 150 | batch_pc = sub_points 151 | 152 | batch_graph_idx = tf.py_func(DP.knn_search, [input_points[0], input_points[0], cfg.k_n], tf.int32) 153 | input_list = input_points + input_neighbors + input_pools + input_up_samples 154 | input_list += [features, batch_label, batch_pc_idx, batch_cloud_idx, batch_graph_idx] 155 | 156 | return input_list 157 | 158 | return tf_map 159 | 160 | def init_input_pipeline(self): 161 | print('Initiating input pipelines') 162 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 163 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 164 | gen_function_val, _, _ = self.get_batch_gen('validation') 165 | gen_function_test, _, _ = self.get_batch_gen('test') 166 | 167 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 168 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 169 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 170 | 171 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 172 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 173 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 174 | 175 | map_func = self.get_tf_mapping2() 176 | 177 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 178 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 179 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 180 | 181 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 182 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 183 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 184 | 185 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 186 | self.flat_inputs = iter.get_next() 187 | self.train_init_op = iter.make_initializer(self.batch_train_data) 188 | self.val_init_op = iter.make_initializer(self.batch_val_data) 189 | self.test_init_op = iter.make_initializer(self.batch_test_data) 190 | 191 | 192 | if __name__ == '__main__': 193 | parser = argparse.ArgumentParser() 194 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 195 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 196 | parser.add_argument('--test_area', type=str, default='14', help='options: 08, 11,12,13,14,15,16,17,18,19,20,21') 197 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 198 | FLAGS = parser.parse_args() 199 | 200 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 201 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 202 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 203 | Mode = FLAGS.mode 204 | 205 | test_area = FLAGS.test_area 206 | dataset = SemanticKITTI(test_area) 207 | dataset.init_input_pipeline() 208 | 209 | if Mode == 'train': 210 | model = Network(dataset, cfg) 211 | model.train(dataset) 212 | elif Mode == 'test': 213 | cfg.saving = False 214 | model = Network(dataset, cfg) 215 | if FLAGS.model_path is not 'None': 216 | chosen_snap = FLAGS.model_path 217 | else: 218 | chosen_snapshot = -1 219 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 220 | chosen_folder = logs[-1] 221 | snap_path = join(chosen_folder, 'snapshots') 222 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 223 | chosen_step = np.sort(snap_steps)[-1] 224 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 225 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 226 | tester.test(model, dataset) 227 | else: 228 | ################## 229 | # Visualize data # 230 | ################## 231 | 232 | with tf.Session() as sess: 233 | sess.run(tf.global_variables_initializer()) 234 | sess.run(dataset.train_init_op) 235 | while True: 236 | flat_inputs = sess.run(dataset.flat_inputs) 237 | pc_xyz = flat_inputs[0] 238 | sub_pc_xyz = flat_inputs[1] 239 | labels = flat_inputs[17] 240 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 241 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 242 | -------------------------------------------------------------------------------- /main_S3DIS.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from BAF_LAC import Network 3 | from tester_S3DIS import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import ConfigS3DIS as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow as tf 9 | import numpy as np 10 | import time, pickle, argparse, glob, os 11 | 12 | 13 | class S3DIS: 14 | def __init__(self, test_area_idx): 15 | self.name = 'S3DIS' 16 | self.path = '/home/data/S3DIS' 17 | self.label_to_names = {0: 'ceiling', 18 | 1: 'floor', 19 | 2: 'wall', 20 | 3: 'beam', 21 | 4: 'column', 22 | 5: 'window', 23 | 6: 'door', 24 | 7: 'table', 25 | 8: 'chair', 26 | 9: 'sofa', 27 | 10: 'bookcase', 28 | 11: 'board', 29 | 12: 'clutter'} 30 | self.num_classes = len(self.label_to_names) 31 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 32 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 33 | self.ignored_labels = np.array([]) 34 | 35 | self.val_split = 'Area_' + str(test_area_idx) 36 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 37 | 38 | # Initiate containers 39 | self.val_proj = [] 40 | self.val_labels = [] 41 | self.possibility = {} 42 | self.min_possibility = {} 43 | self.input_trees = {'training': [], 'validation': []} 44 | self.input_colors = {'training': [], 'validation': []} 45 | self.input_labels = {'training': [], 'validation': []} 46 | self.input_names = {'training': [], 'validation': []} 47 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 48 | 49 | def load_sub_sampled_clouds(self, sub_grid_size): 50 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 51 | for i, file_path in enumerate(self.all_files): 52 | t0 = time.time() 53 | cloud_name = file_path.split('/')[-1][:-4] 54 | if self.val_split in cloud_name: 55 | cloud_split = 'validation' 56 | else: 57 | cloud_split = 'training' 58 | 59 | # Name of the input files 60 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 61 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 62 | 63 | data = read_ply(sub_ply_file) 64 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 65 | sub_labels = data['class'] 66 | 67 | # Read pkl with search tree 68 | with open(kd_tree_file, 'rb') as f: 69 | search_tree = pickle.load(f) 70 | 71 | self.input_trees[cloud_split] += [search_tree] 72 | self.input_colors[cloud_split] += [sub_colors] 73 | self.input_labels[cloud_split] += [sub_labels] 74 | self.input_names[cloud_split] += [cloud_name] 75 | 76 | size = sub_colors.shape[0] * 4 * 7 77 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 78 | 79 | print('\nPreparing reprojected indices for testing') 80 | 81 | # Get validation and test reprojected indices 82 | for i, file_path in enumerate(self.all_files): 83 | t0 = time.time() 84 | cloud_name = file_path.split('/')[-1][:-4] 85 | 86 | # Validation projection and labels 87 | if self.val_split in cloud_name: 88 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 89 | with open(proj_file, 'rb') as f: 90 | proj_idx, labels = pickle.load(f) 91 | self.val_proj += [proj_idx] 92 | self.val_labels += [labels] 93 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 94 | 95 | # Generate the input data flow 96 | def get_batch_gen(self, split): 97 | if split == 'training': 98 | num_per_epoch = cfg.train_steps * cfg.batch_size 99 | elif split == 'validation': 100 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 101 | 102 | self.possibility[split] = [] 103 | self.min_possibility[split] = [] 104 | # Random initialize 105 | for i, tree in enumerate(self.input_colors[split]): 106 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 107 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 108 | 109 | def spatially_regular_gen(): 110 | # Generator loop 111 | for i in range(num_per_epoch): 112 | 113 | # Choose the cloud with the lowest probability 114 | cloud_idx = int(np.argmin(self.min_possibility[split])) 115 | 116 | # Choose the point with the minimum of possibility in the cloud as query point 117 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 118 | 119 | # Get all points within the cloud from tree structure 120 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 121 | 122 | # Center point of input region 123 | center_point = points[point_ind, :].reshape(1, -1) 124 | 125 | # Add noise to the center point 126 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 127 | pick_point = center_point + noise.astype(center_point.dtype) 128 | 129 | # Check if the number of points in the selected cloud is less than the predicted num_points 130 | if len(points) < cfg.num_points: 131 | # Query all points within the cloud 132 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 133 | else: 134 | # Query the predefined number of points 135 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 136 | 137 | # Shuffle index 138 | queried_idx = DP.shuffle_idx(queried_idx) 139 | # Get corresponding points and colors based on the index 140 | queried_pc_xyz = points[queried_idx] 141 | queried_pc_xyz = queried_pc_xyz - pick_point 142 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 143 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 144 | 145 | # Update the possibility of the selected points 146 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 147 | delta = np.square(1 - dists / np.max(dists)) 148 | self.possibility[split][cloud_idx][queried_idx] += delta 149 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 150 | 151 | # Up_sampled with replacement 152 | if len(points) < cfg.num_points: 153 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 154 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 155 | 156 | if True: 157 | yield (queried_pc_xyz.astype(np.float32), 158 | queried_pc_colors.astype(np.float32), 159 | queried_pc_labels, 160 | queried_idx.astype(np.int32), 161 | np.array([cloud_idx], dtype=np.int32)) 162 | 163 | gen_func = spatially_regular_gen 164 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 165 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 166 | return gen_func, gen_types, gen_shapes 167 | 168 | @staticmethod 169 | def get_tf_mapping2(): 170 | # Collect flat inputs 171 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 172 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 173 | input_points = [] 174 | input_neighbors = [] 175 | input_pools = [] 176 | input_up_samples = [] 177 | 178 | for i in range(cfg.num_layers): 179 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 180 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 181 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 182 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 183 | input_points.append(batch_xyz) 184 | input_neighbors.append(neighbour_idx) 185 | input_pools.append(pool_i) 186 | input_up_samples.append(up_i) 187 | batch_xyz = sub_points 188 | 189 | batch_graph_idx = tf.py_func(DP.knn_search, [input_points[0], input_points[0], cfg.k_n], tf.int32) 190 | input_list = input_points + input_neighbors + input_pools + input_up_samples 191 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx, batch_graph_idx] 192 | 193 | return input_list 194 | 195 | return tf_map 196 | 197 | def init_input_pipeline(self): 198 | print('Initiating input pipelines') 199 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 200 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 201 | gen_function_val, _, _ = self.get_batch_gen('validation') 202 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 203 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 204 | 205 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 206 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 207 | map_func = self.get_tf_mapping2() 208 | 209 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 210 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 211 | 212 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 213 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 214 | 215 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 216 | self.flat_inputs = iter.get_next() 217 | self.train_init_op = iter.make_initializer(self.batch_train_data) 218 | self.val_init_op = iter.make_initializer(self.batch_val_data) 219 | 220 | 221 | if __name__ == '__main__': 222 | parser = argparse.ArgumentParser() 223 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 224 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 225 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 226 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 227 | FLAGS = parser.parse_args() 228 | 229 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 230 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 231 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 232 | Mode = FLAGS.mode 233 | 234 | test_area = FLAGS.test_area 235 | dataset = S3DIS(test_area) 236 | dataset.init_input_pipeline() 237 | 238 | if Mode == 'train': 239 | model = Network(dataset, cfg) 240 | model.train(dataset) 241 | elif Mode == 'test': 242 | cfg.saving = False 243 | model = Network(dataset, cfg) 244 | if FLAGS.model_path is not 'None': 245 | chosen_snap = FLAGS.model_path 246 | else: 247 | chosen_snapshot = -1 248 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 249 | chosen_folder = logs[-1] 250 | snap_path = join(chosen_folder, 'snapshots') 251 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 252 | chosen_step = np.sort(snap_steps)[-1] 253 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 254 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 255 | tester.test(model, dataset) 256 | else: 257 | ################## 258 | # Visualize data # 259 | ################## 260 | 261 | with tf.Session() as sess: 262 | sess.run(tf.global_variables_initializer()) 263 | sess.run(dataset.train_init_op) 264 | while True: 265 | flat_inputs = sess.run(dataset.flat_inputs) 266 | pc_xyz = flat_inputs[0] 267 | sub_pc_xyz = flat_inputs[1] 268 | labels = flat_inputs[21] 269 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 270 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 271 | -------------------------------------------------------------------------------- /helper_tool.py: -------------------------------------------------------------------------------- 1 | from open3d import linux as open3d 2 | from os.path import join 3 | import numpy as np 4 | import colorsys, random, os, sys 5 | import pandas as pd 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 ConfigSemanticKITTI: 19 | k_n = 16 # KNN 20 | num_layers = 4 # number of layers 21 | num_points = 4096 * 11 # number of input points 22 | num_classes = 19 # number of valid classes 23 | sub_grid_size = 0.06 # preprocess_parameter 24 | 25 | batch_size = 6 # batch_size during training 26 | val_batch_size = 20 # batch_size during validation and test 27 | train_steps = 1000 # number of steps per epochs 28 | val_steps = 200 # number of validation steps per epoch 29 | 30 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer 31 | d_out = [16, 32, 64, 128] # feature dimension 32 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256] 33 | 34 | noise_init = 3.5 # noise initial parameter 35 | max_epoch = 100 # maximum epoch during training 36 | learning_rate = 1e-2 # initial learning rate 37 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 38 | 39 | train_sum_dir = 'train_log' 40 | saving = True 41 | saving_path = None 42 | 43 | 44 | class ConfigS3DIS: 45 | k_n = 16 # KNN 46 | num_layers = 5 # number of layers 47 | num_points = 40960 # number of input points 48 | num_classes = 13 # number of valid classes 49 | sub_grid_size = 0.04 # preprocess_parameter 50 | 51 | batch_size = 6 # batch_size during training 52 | val_batch_size = 20 # batch_size during validation and test 53 | train_steps = 1000 # number of steps per epochs 54 | val_steps = 200 # number of validation steps per epoch 55 | 56 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 57 | d_out = [16, 32, 64, 128, 256] # feature dimension 58 | 59 | noise_init = 3.5 # noise initial parameter 60 | max_epoch = 100 # maximum epoch during training 61 | learning_rate = 1e-2 # initial learning rate 62 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 63 | 64 | train_sum_dir = 'train_log' 65 | saving = True 66 | saving_path = None 67 | 68 | 69 | class ConfigSemantic3D: 70 | k_n = 16 # KNN 71 | num_layers = 5 # number of layers 72 | num_points = 65536 # number of input points 73 | num_classes = 8 # number of valid classes 74 | sub_grid_size = 0.06 # preprocess_parameter 75 | 76 | batch_size = 4 # batch_size during training 77 | val_batch_size = 16 # batch_size during validation and test 78 | train_steps = 1000 # number of steps per epochs 79 | val_steps = 200 # number of validation steps per epoch 80 | 81 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 82 | d_out = [16, 32, 64, 128, 256] # feature dimension 83 | 84 | noise_init = 3.5 # noise initial parameter 85 | max_epoch = 100 # maximum epoch during training 86 | learning_rate = 1e-2 # initial learning rate 87 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 88 | 89 | train_sum_dir = 'train_log' 90 | saving = True 91 | saving_path = None 92 | 93 | augment_scale_anisotropic = True 94 | augment_symmetries = [True, False, False] 95 | augment_rotation = 'vertical' 96 | augment_scale_min = 0.8 97 | augment_scale_max = 1.2 98 | augment_noise = 0.001 99 | augment_occlusion = 'none' 100 | augment_color = 0.8 101 | 102 | 103 | class DataProcessing: 104 | @staticmethod 105 | def load_pc_semantic3d(filename): 106 | pc_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.float16) 107 | pc = pc_pd.values 108 | return pc 109 | 110 | @staticmethod 111 | def load_label_semantic3d(filename): 112 | label_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.uint8) 113 | cloud_labels = label_pd.values 114 | return cloud_labels 115 | 116 | @staticmethod 117 | def load_pc_kitti(pc_path): 118 | scan = np.fromfile(pc_path, dtype=np.float32) 119 | scan = scan.reshape((-1, 4)) 120 | points = scan[:, 0:3] # get xyz 121 | return points 122 | 123 | @staticmethod 124 | def load_label_kitti(label_path, remap_lut): 125 | label = np.fromfile(label_path, dtype=np.uint32) 126 | label = label.reshape((-1)) 127 | sem_label = label & 0xFFFF # semantic label in lower half 128 | inst_label = label >> 16 # instance id in upper half 129 | assert ((sem_label + (inst_label << 16) == label).all()) 130 | sem_label = remap_lut[sem_label] 131 | return sem_label.astype(np.int32) 132 | 133 | @staticmethod 134 | def get_file_list(dataset_path, test_scan_num): 135 | seq_list = np.sort(os.listdir(dataset_path)) 136 | 137 | train_file_list = [] 138 | test_file_list = [] 139 | val_file_list = [] 140 | for seq_id in seq_list: 141 | seq_path = join(dataset_path, seq_id) 142 | pc_path = join(seq_path, 'velodyne') 143 | if seq_id == '08': 144 | val_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 145 | if seq_id == test_scan_num: 146 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 147 | elif int(seq_id) >= 11 and seq_id == test_scan_num: 148 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 149 | elif seq_id in ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']: 150 | train_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 151 | 152 | train_file_list = np.concatenate(train_file_list, axis=0) 153 | val_file_list = np.concatenate(val_file_list, axis=0) 154 | test_file_list = np.concatenate(test_file_list, axis=0) 155 | return train_file_list, val_file_list, test_file_list 156 | 157 | @staticmethod 158 | def knn_search(support_pts, query_pts, k): 159 | """ 160 | :param support_pts: points you have, B*N1*3 161 | :param query_pts: points you want to know the neighbour index, B*N2*3 162 | :param k: Number of neighbours in knn search 163 | :return: neighbor_idx: neighboring points indexes, B*N2*k 164 | """ 165 | 166 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 167 | return neighbor_idx.astype(np.int32) 168 | 169 | @staticmethod 170 | def data_aug(xyz, color, labels, idx, num_out): 171 | num_in = len(xyz) 172 | dup = np.random.choice(num_in, num_out - num_in) 173 | xyz_dup = xyz[dup, ...] 174 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) 175 | color_dup = color[dup, ...] 176 | color_aug = np.concatenate([color, color_dup], 0) 177 | idx_dup = list(range(num_in)) + list(dup) 178 | idx_aug = idx[idx_dup] 179 | label_aug = labels[idx_dup] 180 | return xyz_aug, color_aug, idx_aug, label_aug 181 | 182 | @staticmethod 183 | def shuffle_idx(x): 184 | # Random shuffle the index 185 | idx = np.arange(len(x)) 186 | np.random.shuffle(idx) 187 | return x[idx] 188 | 189 | @staticmethod 190 | def shuffle_list(data_list): 191 | indices = np.arange(np.shape(data_list)[0]) 192 | np.random.shuffle(indices) 193 | data_list = data_list[indices] 194 | return data_list 195 | 196 | @staticmethod 197 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 198 | """ 199 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 200 | :param points: (N, 3) matrix of input points 201 | :param features: optional (N, d) matrix of features (floating number) 202 | :param labels: optional (N,) matrix of integer labels 203 | :param grid_size: parameter defining the size of grid voxels 204 | :param verbose: 1 to display 205 | :return: sub_sampled points, with features and/or labels depending of the input 206 | """ 207 | 208 | if (features is None) and (labels is None): 209 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 210 | elif labels is None: 211 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 212 | elif features is None: 213 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 214 | else: 215 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 216 | verbose=verbose) 217 | 218 | @staticmethod 219 | def IoU_from_confusions(confusions): 220 | """ 221 | Computes IoU from confusion matrices. 222 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 223 | the last axes. n_c = number of classes 224 | :return: ([..., n_c] np.float32) IoU score 225 | """ 226 | 227 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 228 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 229 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 230 | TP_plus_FN = np.sum(confusions, axis=-1) 231 | TP_plus_FP = np.sum(confusions, axis=-2) 232 | 233 | # Compute IoU 234 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 235 | 236 | # Compute mIoU with only the actual classes 237 | mask = TP_plus_FN < 1e-3 238 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 239 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 240 | 241 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 242 | IoU += mask * mIoU 243 | return IoU 244 | 245 | @staticmethod 246 | def get_class_weights(dataset_name): 247 | # Pre-calculate the number of points in each category 248 | num_per_class = [] 249 | if dataset_name is 'S3DIS': 250 | num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733, 251 | 650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32) 252 | elif dataset_name is 'Semantic3D': 253 | num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353], 254 | dtype=np.int32) 255 | elif dataset_name is 'SemanticKITTI': 256 | num_per_class = np.array([55437630, 320797, 541736, 2578735, 3274484, 552662, 184064, 78858, 257 | 240942562, 17294618, 170599734, 6369672, 230413074, 101130274, 476491114, 258 | 9833174, 129609852, 4506626, 1168181]) 259 | weight = num_per_class / float(sum(num_per_class)) 260 | ce_label_weight = 1 / (weight + 0.02) 261 | return np.expand_dims(ce_label_weight, axis=0) 262 | 263 | 264 | class Plot: 265 | @staticmethod 266 | def random_colors(N, bright=True, seed=0): 267 | brightness = 1.0 if bright else 0.7 268 | hsv = [(0.15 + i / float(N), 1, brightness) for i in range(N)] 269 | colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)) 270 | random.seed(seed) 271 | random.shuffle(colors) 272 | return colors 273 | 274 | @staticmethod 275 | def draw_pc(pc_xyzrgb): 276 | pc = open3d.PointCloud() 277 | pc.points = open3d.Vector3dVector(pc_xyzrgb[:, 0:3]) 278 | if pc_xyzrgb.shape[1] == 3: 279 | open3d.draw_geometries([pc]) 280 | return 0 281 | if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255 282 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.) 283 | else: 284 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6]) 285 | open3d.draw_geometries([pc]) 286 | return 0 287 | 288 | @staticmethod 289 | def draw_pc_sem_ins(pc_xyz, pc_sem_ins, plot_colors=None): 290 | """ 291 | pc_xyz: 3D coordinates of point clouds 292 | pc_sem_ins: semantic or instance labels 293 | plot_colors: custom color list 294 | """ 295 | if plot_colors is not None: 296 | ins_colors = plot_colors 297 | else: 298 | ins_colors = Plot.random_colors(len(np.unique(pc_sem_ins)) + 1, seed=2) 299 | 300 | ############################## 301 | sem_ins_labels = np.unique(pc_sem_ins) 302 | sem_ins_bbox = [] 303 | Y_colors = np.zeros((pc_sem_ins.shape[0], 3)) 304 | for id, semins in enumerate(sem_ins_labels): 305 | valid_ind = np.argwhere(pc_sem_ins == semins)[:, 0] 306 | if semins <= -1: 307 | tp = [0, 0, 0] 308 | else: 309 | if plot_colors is not None: 310 | tp = ins_colors[semins] 311 | else: 312 | tp = ins_colors[id] 313 | 314 | Y_colors[valid_ind] = tp 315 | 316 | ### bbox 317 | valid_xyz = pc_xyz[valid_ind] 318 | 319 | xmin = np.min(valid_xyz[:, 0]); 320 | xmax = np.max(valid_xyz[:, 0]) 321 | ymin = np.min(valid_xyz[:, 1]); 322 | ymax = np.max(valid_xyz[:, 1]) 323 | zmin = np.min(valid_xyz[:, 2]); 324 | zmax = np.max(valid_xyz[:, 2]) 325 | sem_ins_bbox.append( 326 | [[xmin, ymin, zmin], [xmax, ymax, zmax], [min(tp[0], 1.), min(tp[1], 1.), min(tp[2], 1.)]]) 327 | 328 | Y_semins = np.concatenate([pc_xyz[:, 0:3], Y_colors], axis=-1) 329 | Plot.draw_pc(Y_semins) 330 | return Y_semins 331 | -------------------------------------------------------------------------------- /main_Semantic3D.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists 2 | from BAF_LAC import Network 3 | from tester_Semantic3D import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import Plot 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import ConfigSemantic3D as cfg 8 | import tensorflow as tf 9 | import numpy as np 10 | import pickle, argparse, os 11 | 12 | 13 | class Semantic3D: 14 | def __init__(self): 15 | self.name = 'Semantic3D' 16 | self.path = '/home/data/semantic3d' 17 | self.label_to_names = {0: 'unlabeled', 18 | 1: 'man-made terrain', 19 | 2: 'natural terrain', 20 | 3: 'high vegetation', 21 | 4: 'low vegetation', 22 | 5: 'buildings', 23 | 6: 'hard scape', 24 | 7: 'scanning artefacts', 25 | 8: 'cars'} 26 | self.num_classes = len(self.label_to_names) 27 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 28 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 29 | self.ignored_labels = np.sort([0]) 30 | 31 | self.original_folder = join(self.path, 'original_data') 32 | self.full_pc_folder = join(self.path, 'original_ply') 33 | self.sub_pc_folder = join(self.path, 'input_{:.3f}'.format(cfg.sub_grid_size)) 34 | 35 | # Following KPConv to do the train-validation split 36 | self.all_splits = [0, 1, 4, 5, 3, 4, 3, 0, 1, 2, 3, 4, 2, 0, 5] 37 | self.val_split = 1 38 | 39 | # Initial training-validation-testing files 40 | self.train_files = [] 41 | self.val_files = [] 42 | self.test_files = [] 43 | cloud_names = [file_name[:-4] for file_name in os.listdir(self.original_folder) if file_name[-4:] == '.txt'] 44 | for pc_name in cloud_names: 45 | if exists(join(self.original_folder, pc_name + '.labels')): 46 | self.train_files.append(join(self.sub_pc_folder, pc_name + '.ply')) 47 | else: 48 | self.test_files.append(join(self.full_pc_folder, pc_name + '.ply')) 49 | 50 | self.train_files = np.sort(self.train_files) 51 | self.test_files = np.sort(self.test_files) 52 | 53 | for i, file_path in enumerate(self.train_files): 54 | if self.all_splits[i] == self.val_split: 55 | self.val_files.append(file_path) 56 | 57 | self.train_files = np.sort([x for x in self.train_files if x not in self.val_files]) 58 | 59 | # Initiate containers 60 | self.val_proj = [] 61 | self.val_labels = [] 62 | self.test_proj = [] 63 | self.test_labels = [] 64 | 65 | self.possibility = {} 66 | self.min_possibility = {} 67 | self.class_weight = {} 68 | self.input_trees = {'training': [], 'validation': [], 'test': []} 69 | self.input_colors = {'training': [], 'validation': [], 'test': []} 70 | self.input_labels = {'training': [], 'validation': []} 71 | 72 | # Ascii files dict for testing 73 | self.ascii_files = { 74 | 'MarketplaceFeldkirch_Station4_rgb_intensity-reduced.ply': 'marketsquarefeldkirch4-reduced.labels', 75 | 'sg27_station10_rgb_intensity-reduced.ply': 'sg27_10-reduced.labels', 76 | 'sg28_Station2_rgb_intensity-reduced.ply': 'sg28_2-reduced.labels', 77 | 'StGallenCathedral_station6_rgb_intensity-reduced.ply': 'stgallencathedral6-reduced.labels', 78 | 'birdfountain_station1_xyz_intensity_rgb.ply': 'birdfountain1.labels', 79 | 'castleblatten_station1_intensity_rgb.ply': 'castleblatten1.labels', 80 | 'castleblatten_station5_xyz_intensity_rgb.ply': 'castleblatten5.labels', 81 | 'marketplacefeldkirch_station1_intensity_rgb.ply': 'marketsquarefeldkirch1.labels', 82 | 'marketplacefeldkirch_station4_intensity_rgb.ply': 'marketsquarefeldkirch4.labels', 83 | 'marketplacefeldkirch_station7_intensity_rgb.ply': 'marketsquarefeldkirch7.labels', 84 | 'sg27_station10_intensity_rgb.ply': 'sg27_10.labels', 85 | 'sg27_station3_intensity_rgb.ply': 'sg27_3.labels', 86 | 'sg27_station6_intensity_rgb.ply': 'sg27_6.labels', 87 | 'sg27_station8_intensity_rgb.ply': 'sg27_8.labels', 88 | 'sg28_station2_intensity_rgb.ply': 'sg28_2.labels', 89 | 'sg28_station5_xyz_intensity_rgb.ply': 'sg28_5.labels', 90 | 'stgallencathedral_station1_intensity_rgb.ply': 'stgallencathedral1.labels', 91 | 'stgallencathedral_station3_intensity_rgb.ply': 'stgallencathedral3.labels', 92 | 'stgallencathedral_station6_intensity_rgb.ply': 'stgallencathedral6.labels'} 93 | 94 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 95 | 96 | def load_sub_sampled_clouds(self, sub_grid_size): 97 | 98 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 99 | files = np.hstack((self.train_files, self.val_files, self.test_files)) 100 | 101 | for i, file_path in enumerate(files): 102 | cloud_name = file_path.split('/')[-1][:-4] 103 | print('Load_pc_' + str(i) + ': ' + cloud_name) 104 | if file_path in self.val_files: 105 | cloud_split = 'validation' 106 | elif file_path in self.train_files: 107 | cloud_split = 'training' 108 | else: 109 | cloud_split = 'test' 110 | 111 | # Name of the input files 112 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 113 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 114 | 115 | # Read ply with data 116 | data = read_ply(sub_ply_file) 117 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 118 | if cloud_split == 'test': 119 | sub_labels = None 120 | else: 121 | sub_labels = data['class'] 122 | 123 | # Read pkl with search tree 124 | with open(kd_tree_file, 'rb') as f: 125 | search_tree = pickle.load(f) 126 | 127 | self.input_trees[cloud_split] += [search_tree] 128 | self.input_colors[cloud_split] += [sub_colors] 129 | if cloud_split in ['training', 'validation']: 130 | self.input_labels[cloud_split] += [sub_labels] 131 | 132 | # Get validation and test re_projection indices 133 | print('\nPreparing reprojection indices for validation and test') 134 | 135 | for i, file_path in enumerate(files): 136 | 137 | # Get cloud name and split 138 | cloud_name = file_path.split('/')[-1][:-4] 139 | 140 | # Validation projection and labels 141 | if file_path in self.val_files: 142 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 143 | with open(proj_file, 'rb') as f: 144 | proj_idx, labels = pickle.load(f) 145 | self.val_proj += [proj_idx] 146 | self.val_labels += [labels] 147 | 148 | # Test projection 149 | if file_path in self.test_files: 150 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 151 | with open(proj_file, 'rb') as f: 152 | proj_idx, labels = pickle.load(f) 153 | self.test_proj += [proj_idx] 154 | self.test_labels += [labels] 155 | print('finished') 156 | return 157 | 158 | # Generate the input data flow 159 | def get_batch_gen(self, split): 160 | if split == 'training': 161 | num_per_epoch = cfg.train_steps * cfg.batch_size 162 | elif split == 'validation': 163 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 164 | elif split == 'test': 165 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 166 | 167 | # Reset possibility 168 | self.possibility[split] = [] 169 | self.min_possibility[split] = [] 170 | self.class_weight[split] = [] 171 | 172 | # Random initialize 173 | for i, tree in enumerate(self.input_trees[split]): 174 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 175 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 176 | 177 | if split != 'test': 178 | _, num_class_total = np.unique(np.hstack(self.input_labels[split]), return_counts=True) 179 | self.class_weight[split] += [np.squeeze([num_class_total / np.sum(num_class_total)], axis=0)] 180 | 181 | def spatially_regular_gen(): 182 | 183 | # Generator loop 184 | for i in range(num_per_epoch): # num_per_epoch 185 | 186 | # Choose the cloud with the lowest probability 187 | cloud_idx = int(np.argmin(self.min_possibility[split])) 188 | 189 | # Choose the point with the minimum of possibility in the cloud as query point 190 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 191 | 192 | # Get all points within the cloud from tree structure 193 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 194 | 195 | # Center point of input region 196 | center_point = points[point_ind, :].reshape(1, -1) 197 | 198 | # Add noise to the center point 199 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 200 | pick_point = center_point + noise.astype(center_point.dtype) 201 | query_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 202 | 203 | # Shuffle index 204 | query_idx = DP.shuffle_idx(query_idx) 205 | 206 | # Get corresponding points and colors based on the index 207 | queried_pc_xyz = points[query_idx] 208 | queried_pc_xyz[:, 0:2] = queried_pc_xyz[:, 0:2] - pick_point[:, 0:2] 209 | queried_pc_colors = self.input_colors[split][cloud_idx][query_idx] 210 | if split == 'test': 211 | queried_pc_labels = np.zeros(queried_pc_xyz.shape[0]) 212 | queried_pt_weight = 1 213 | else: 214 | queried_pc_labels = self.input_labels[split][cloud_idx][query_idx] 215 | queried_pc_labels = np.array([self.label_to_idx[l] for l in queried_pc_labels]) 216 | queried_pt_weight = np.array([self.class_weight[split][0][n] for n in queried_pc_labels]) 217 | 218 | # Update the possibility of the selected points 219 | dists = np.sum(np.square((points[query_idx] - pick_point).astype(np.float32)), axis=1) 220 | delta = np.square(1 - dists / np.max(dists)) * queried_pt_weight 221 | self.possibility[split][cloud_idx][query_idx] += delta 222 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 223 | 224 | if True: 225 | yield (queried_pc_xyz, 226 | queried_pc_colors.astype(np.float32), 227 | queried_pc_labels, 228 | query_idx.astype(np.int32), 229 | np.array([cloud_idx], dtype=np.int32)) 230 | 231 | gen_func = spatially_regular_gen 232 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 233 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 234 | return gen_func, gen_types, gen_shapes 235 | 236 | def get_tf_mapping(self): 237 | # Collect flat inputs 238 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 239 | batch_features = tf.map_fn(self.tf_augment_input, [batch_xyz, batch_features], dtype=tf.float32) 240 | input_points = [] 241 | input_neighbors = [] 242 | input_pools = [] 243 | input_up_samples = [] 244 | 245 | for i in range(cfg.num_layers): 246 | neigh_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 247 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 248 | pool_i = neigh_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 249 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 250 | input_points.append(batch_xyz) 251 | input_neighbors.append(neigh_idx) 252 | input_pools.append(pool_i) 253 | input_up_samples.append(up_i) 254 | batch_xyz = sub_points 255 | 256 | batch_graph_idx = tf.py_func(DP.knn_search, [input_points[0], input_points[0], cfg.k_n], tf.int32) 257 | input_list = input_points + input_neighbors + input_pools + input_up_samples 258 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx, batch_graph_idx] 259 | 260 | return input_list 261 | 262 | return tf_map 263 | 264 | # Data augmentation 265 | @staticmethod 266 | def tf_augment_input(inputs): 267 | xyz = inputs[0] 268 | features = inputs[1] 269 | theta = tf.random_uniform((1,), minval=0, maxval=2 * np.pi) 270 | # Rotation matrices 271 | c, s = tf.cos(theta), tf.sin(theta) 272 | cs0 = tf.zeros_like(c) 273 | cs1 = tf.ones_like(c) 274 | R = tf.stack([c, -s, cs0, s, c, cs0, cs0, cs0, cs1], axis=1) 275 | stacked_rots = tf.reshape(R, (3, 3)) 276 | 277 | # Apply rotations 278 | transformed_xyz = tf.reshape(tf.matmul(xyz, stacked_rots), [-1, 3]) 279 | # Choose random scales for each example 280 | min_s = cfg.augment_scale_min 281 | max_s = cfg.augment_scale_max 282 | if cfg.augment_scale_anisotropic: 283 | s = tf.random_uniform((1, 3), minval=min_s, maxval=max_s) 284 | else: 285 | s = tf.random_uniform((1, 1), minval=min_s, maxval=max_s) 286 | 287 | symmetries = [] 288 | for i in range(3): 289 | if cfg.augment_symmetries[i]: 290 | symmetries.append(tf.round(tf.random_uniform((1, 1))) * 2 - 1) 291 | else: 292 | symmetries.append(tf.ones([1, 1], dtype=tf.float32)) 293 | s *= tf.concat(symmetries, 1) 294 | 295 | # Create N x 3 vector of scales to multiply with stacked points 296 | stacked_scales = tf.tile(s, [tf.shape(transformed_xyz)[0], 1]) 297 | 298 | # Apply scales 299 | transformed_xyz = transformed_xyz * stacked_scales 300 | 301 | noise = tf.random_normal(tf.shape(transformed_xyz), stddev=cfg.augment_noise) 302 | transformed_xyz = transformed_xyz + noise 303 | rgb = features[:, :3] 304 | stacked_features = tf.concat([transformed_xyz, rgb], axis=-1) 305 | return stacked_features 306 | 307 | def init_input_pipeline(self): 308 | print('Initiating input pipelines') 309 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 310 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 311 | gen_function_val, _, _ = self.get_batch_gen('validation') 312 | gen_function_test, _, _ = self.get_batch_gen('test') 313 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 314 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 315 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 316 | 317 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 318 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 319 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 320 | map_func = self.get_tf_mapping() 321 | 322 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 323 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 324 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 325 | 326 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 327 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 328 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 329 | 330 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 331 | self.flat_inputs = iter.get_next() 332 | self.train_init_op = iter.make_initializer(self.batch_train_data) 333 | self.val_init_op = iter.make_initializer(self.batch_val_data) 334 | self.test_init_op = iter.make_initializer(self.batch_test_data) 335 | 336 | 337 | if __name__ == '__main__': 338 | parser = argparse.ArgumentParser() 339 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 340 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, val') 341 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 342 | FLAGS = parser.parse_args() 343 | 344 | GPU_ID = FLAGS.gpu 345 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 346 | os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU_ID) 347 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 348 | 349 | Mode = FLAGS.mode 350 | dataset = Semantic3D() 351 | dataset.init_input_pipeline() 352 | 353 | if Mode == 'train': 354 | model = Network(dataset, cfg) 355 | model.train(dataset) 356 | elif Mode == 'test': 357 | cfg.saving = False 358 | model = Network(dataset, cfg) 359 | if FLAGS.model_path is not 'None': 360 | chosen_snap = FLAGS.model_path 361 | else: 362 | chosen_snapshot = -1 363 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 364 | chosen_folder = logs[-1] 365 | snap_path = join(chosen_folder, 'snapshots') 366 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 367 | chosen_step = np.sort(snap_steps)[-1] 368 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 369 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 370 | tester.test(model, dataset) 371 | 372 | else: 373 | ################## 374 | # Visualize data # 375 | ################## 376 | 377 | with tf.Session() as sess: 378 | sess.run(tf.global_variables_initializer()) 379 | sess.run(dataset.train_init_op) 380 | while True: 381 | flat_inputs = sess.run(dataset.flat_inputs) 382 | pc_xyz = flat_inputs[0] 383 | sub_pc_xyz = flat_inputs[1] 384 | labels = flat_inputs[21] 385 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 386 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 387 | -------------------------------------------------------------------------------- /BAF_LAC.py: -------------------------------------------------------------------------------- 1 | from os.path import exists, join 2 | from os import makedirs 3 | from sklearn.metrics import confusion_matrix 4 | from helper_tool import DataProcessing as DP 5 | import tensorflow as tf 6 | import numpy as np 7 | import helper_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_%H-%M-%S', time.gmtime()) 25 | else: 26 | self.saving_path = self.config.saving_path 27 | makedirs(self.saving_path) if not exists(self.saving_path) else None 28 | 29 | with tf.variable_scope('inputs'): 30 | self.inputs = dict() 31 | num_layers = self.config.num_layers 32 | self.inputs['xyz'] = flat_inputs[:num_layers] 33 | self.inputs['neigh_idx'] = flat_inputs[num_layers:2 * num_layers] 34 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 35 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 36 | self.inputs['features'] = flat_inputs[4 * num_layers] 37 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 38 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 39 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 40 | self.inputs['graph_idx'] = flat_inputs[4 * num_layers + 4] 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.class_weights = DP.get_class_weights(dataset.name) 50 | self.Log_file = open('log_train_' + dataset.name + str(dataset.val_split) + '.txt', 'a') 51 | 52 | with tf.variable_scope('layers'): 53 | self.logits = self.inference(self.inputs, self.is_training) 54 | 55 | ##################################################################### 56 | # Ignore the invalid point (unlabeled) when calculating the loss # 57 | ##################################################################### 58 | with tf.variable_scope('loss'): 59 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 60 | self.labels = tf.reshape(self.labels, [-1]) 61 | 62 | # Boolean mask of points that should be ignored 63 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 64 | for ign_label in self.config.ignored_label_inds: 65 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 66 | 67 | # Collect logits and labels that are not ignored 68 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 69 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 70 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 71 | 72 | # Reduce label values in the range of logit shape 73 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 74 | inserted_value = tf.zeros((1,), dtype=tf.int32) 75 | for ign_label in self.config.ignored_label_inds: 76 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 77 | valid_labels = tf.gather(reducing_list, valid_labels_init) 78 | 79 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 80 | 81 | with tf.variable_scope('optimizer'): 82 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 83 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 84 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 85 | 86 | with tf.variable_scope('results'): 87 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 88 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 89 | self.prob_logits = tf.nn.softmax(self.logits) 90 | 91 | tf.summary.scalar('learning_rate', self.learning_rate) 92 | tf.summary.scalar('loss', self.loss) 93 | tf.summary.scalar('accuracy', self.accuracy) 94 | 95 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 96 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 97 | c_proto = tf.ConfigProto() 98 | c_proto.gpu_options.allow_growth = True 99 | self.sess = tf.Session(config=c_proto) 100 | self.merged = tf.summary.merge_all() 101 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 102 | self.sess.run(tf.global_variables_initializer()) 103 | 104 | def inference(self, inputs, is_training): 105 | 106 | d_out = self.config.d_out 107 | feature = inputs['features'] 108 | feature = tf.layers.dense(feature, 8, activation=None, name='fc0') 109 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 110 | feature = tf.expand_dims(feature, axis=2) 111 | 112 | # ###########################Encoder############################ 113 | f_encoder_list = [] 114 | for i in range(self.config.num_layers): 115 | f_encoder_i = self.dilated_res_block(feature, inputs['xyz'][i], inputs['neigh_idx'][i], d_out[i], 116 | 'Encoder_layer_' + str(i), is_training) 117 | f_sampled_i = self.random_sample(f_encoder_i, inputs['sub_idx'][i]) 118 | feature = f_sampled_i 119 | f_encoder_list.append(f_encoder_i) 120 | # ###########################Encoder############################ 121 | 122 | feature = helper_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 | f_encoder_list.append(feature) 125 | 126 | # ###########################Decoder############################ 127 | f_decoder_list = [] 128 | for j in range(self.config.num_layers): 129 | # Backward attentive fusing 130 | f_high_encoder = helper_tf_util.conv2d(f_encoder_list[-j - 1], f_encoder_list[-j - 2].get_shape()[-1].value, 131 | [1, 1], 'Squeeze_layer_' + str(j), [1, 1], 'VALID', bn=True, 132 | is_training=is_training) 133 | f_high_encoder_interp = self.nearest_interpolation(f_high_encoder, inputs['interp_idx'][-j - 1]) 134 | attention_map = helper_tf_util.conv2d_transpose(f_high_encoder_interp, 135 | f_encoder_list[-j - 2].get_shape()[-1].value, [1, 1], 136 | 'Attention_layer_' + str(j), [1, 1], 'VALID', bn=False, 137 | is_training=is_training, activation_fn=None) 138 | attention_map = tf.nn.sigmoid(attention_map) 139 | 140 | f_fuse = f_encoder_list[-j - 2] * attention_map + f_encoder_list[-j - 2] 141 | 142 | f_enhance = helper_tf_util.conv2d(f_fuse, f_encoder_list[-j - 2].get_shape()[-1].value, [1, 1], 143 | 'Enhance_layer_' + str(j), [1, 1], 'VALID', bn=True, 144 | is_training=is_training) 145 | 146 | # Decoder 147 | f_interp_i = self.nearest_interpolation(feature, inputs['interp_idx'][-j - 1]) 148 | f_interp_i = helper_tf_util.conv2d_transpose(f_interp_i, f_encoder_list[-j - 2].get_shape()[-1].value, 149 | [1, 1], 'Interp_layer_' + str(j), [1, 1], 'VALID', bn=True, 150 | is_training=is_training) 151 | 152 | f_decoder_i = helper_tf_util.conv2d(f_interp_i + f_enhance, f_encoder_list[-j - 2].get_shape()[-1].value, 153 | [1, 1], 'Decoder_layer_' + str(j), [1, 1], 'VALID', bn=True, 154 | is_training=is_training) 155 | feature = f_decoder_i 156 | f_decoder_list.append(f_decoder_i) 157 | # ###########################Decoder############################ 158 | 159 | f_layer_fc1 = helper_tf_util.conv2d(f_decoder_list[-1], 32, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 160 | f_layer_fc2 = helper_tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 161 | f_neigh = self.gather_neighbour(tf.squeeze(f_layer_fc2, axis=2), inputs['graph_idx']) 162 | f_neigh = helper_tf_util.conv2d(f_neigh, 32, [1, 1], 'local_graph', [1, 1], 'VALID', True, is_training) 163 | f_neigh = tf.reduce_max(f_neigh, axis=2, keepdims=True) + f_layer_fc2 164 | f_layer_drop = helper_tf_util.dropout(f_neigh, keep_prob=0.5, is_training=is_training, scope='dp1') 165 | f_layer_fc3 = helper_tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, 166 | is_training, activation_fn=None) 167 | f_out = tf.squeeze(f_layer_fc3, [2]) 168 | return f_out 169 | 170 | def train(self, dataset): 171 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 172 | self.sess.run(dataset.train_init_op) 173 | while self.training_epoch < self.config.max_epoch: 174 | t_start = time.time() 175 | try: 176 | ops = [self.train_op, 177 | self.extra_update_ops, 178 | self.merged, 179 | self.loss, 180 | self.logits, 181 | self.labels, 182 | self.accuracy] 183 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 184 | self.train_writer.add_summary(summary, self.training_step) 185 | t_end = time.time() 186 | if self.training_step % 50 == 0: 187 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 188 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 189 | self.training_step += 1 190 | 191 | except tf.errors.OutOfRangeError: 192 | 193 | m_iou = self.evaluate(dataset) 194 | if m_iou > np.max(self.mIou_list): 195 | # Save the best model 196 | snapshot_directory = join(self.saving_path, 'snapshots') 197 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 198 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 199 | self.mIou_list.append(m_iou) 200 | log_out('Best m_IoU is: {:5.3f}'.format(max(self.mIou_list)), self.Log_file) 201 | 202 | self.training_epoch += 1 203 | self.sess.run(dataset.train_init_op) 204 | # Update learning rate 205 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, 206 | self.config.lr_decays[self.training_epoch])) 207 | self.sess.run(op) 208 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 209 | 210 | except tf.errors.InvalidArgumentError as e: 211 | 212 | print('Caught a NaN error: ') 213 | print(e.error_code) 214 | print(e.message) 215 | print(e.op) 216 | print(e.op.name) 217 | print([t.name for t in e.op.inputs]) 218 | print([t.name for t in e.op.outputs]) 219 | 220 | a = 1 / 0 221 | 222 | print('finished') 223 | self.sess.close() 224 | 225 | def evaluate(self, dataset): 226 | 227 | # Initialise iterator with validation data 228 | self.sess.run(dataset.val_init_op) 229 | 230 | gt_classes = [0 for _ in range(self.config.num_classes)] 231 | positive_classes = [0 for _ in range(self.config.num_classes)] 232 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 233 | val_total_correct = 0 234 | val_total_seen = 0 235 | 236 | for step_id in range(self.config.val_steps): 237 | if step_id % 50 == 0: 238 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 239 | try: 240 | ops = (self.prob_logits, self.labels, self.accuracy) 241 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 242 | pred = np.argmax(stacked_prob, 1) 243 | if not self.config.ignored_label_inds: 244 | pred_valid = pred 245 | labels_valid = labels 246 | else: 247 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 248 | labels_valid = np.delete(labels, invalid_idx) 249 | labels_valid = labels_valid - 1 250 | pred_valid = np.delete(pred, invalid_idx) 251 | 252 | correct = np.sum(pred_valid == labels_valid) 253 | val_total_correct += correct 254 | val_total_seen += len(labels_valid) 255 | 256 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 257 | gt_classes += np.sum(conf_matrix, axis=1) 258 | positive_classes += np.sum(conf_matrix, axis=0) 259 | true_positive_classes += np.diagonal(conf_matrix) 260 | 261 | except tf.errors.OutOfRangeError: 262 | break 263 | 264 | iou_list = [] 265 | for n in range(0, self.config.num_classes, 1): 266 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 267 | iou_list.append(iou) 268 | mean_iou = sum(iou_list) / float(self.config.num_classes) 269 | 270 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 271 | log_out('mean IoU: {}'.format(mean_iou), self.Log_file) 272 | 273 | mean_iou = 100 * mean_iou 274 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 275 | s = '{:5.2f} | '.format(mean_iou) 276 | for IoU in iou_list: 277 | s += '{:5.2f} '.format(100 * IoU) 278 | log_out('-' * len(s), self.Log_file) 279 | log_out(s, self.Log_file) 280 | log_out('-' * len(s), self.Log_file) 281 | return mean_iou 282 | 283 | def get_loss(self, logits, labels, pre_cal_weights): 284 | # Calculate the weighted cross entropy according to the inverse frequency 285 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 286 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 287 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 288 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels) 289 | weighted_losses = unweighted_losses * weights 290 | output_loss = tf.reduce_mean(weighted_losses) 291 | return output_loss 292 | 293 | def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training): 294 | f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 295 | f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 296 | f_pc = helper_tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, 297 | activation_fn=None) 298 | shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', 299 | activation_fn=None, bn=True, is_training=is_training) 300 | return tf.nn.leaky_relu(f_pc + shortcut) 301 | 302 | def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training): 303 | d_in = feature.get_shape()[-1].value 304 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) 305 | f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 306 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 307 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 308 | f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training) 309 | 310 | f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 311 | f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx) 312 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 313 | f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training) 314 | return f_pc_agg 315 | 316 | def relative_pos_encoding(self, xyz, neigh_idx): 317 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 318 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 319 | relative_xyz = xyz_tile - neighbor_xyz 320 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 321 | relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1) 322 | return relative_feature 323 | 324 | @staticmethod 325 | def random_sample(feature, pool_idx): 326 | """ 327 | :param feature: [B, N, d] input features matrix 328 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 329 | :return: pool_features = [B, N', d] pooled features matrix 330 | """ 331 | feature = tf.squeeze(feature, axis=2) 332 | num_neigh = tf.shape(pool_idx)[-1] 333 | d = feature.get_shape()[-1] 334 | batch_size = tf.shape(pool_idx)[0] 335 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 336 | pool_features = tf.batch_gather(feature, pool_idx) 337 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 338 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 339 | return pool_features 340 | 341 | @staticmethod 342 | def nearest_interpolation(feature, interp_idx): 343 | """ 344 | :param feature: [B, N, d] input features matrix 345 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 346 | :return: [B, up_num_points, d] interpolated features matrix 347 | """ 348 | feature = tf.squeeze(feature, axis=2) 349 | batch_size = tf.shape(interp_idx)[0] 350 | up_num_points = tf.shape(interp_idx)[1] 351 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 352 | interpolated_features = tf.batch_gather(feature, interp_idx) 353 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 354 | return interpolated_features 355 | 356 | @staticmethod 357 | def gather_neighbour(pc, neighbor_idx): 358 | # Gather the coordinates or features of neighboring points 359 | batch_size = tf.shape(pc)[0] 360 | num_points = tf.shape(pc)[1] 361 | d = pc.get_shape()[2].value 362 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 363 | features = tf.batch_gather(pc, index_input) 364 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 365 | return features 366 | 367 | @staticmethod 368 | def att_pooling(feature_set, d_out, name, is_training): 369 | batch_size = tf.shape(feature_set)[0] 370 | num_points = tf.shape(feature_set)[1] 371 | num_neigh = tf.shape(feature_set)[2] 372 | d = feature_set.get_shape()[3].value 373 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 374 | att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc') 375 | att_scores = tf.nn.softmax(att_activation, axis=1) 376 | f_agg = f_reshaped * att_scores 377 | f_agg = tf.reduce_sum(f_agg, axis=1) 378 | f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d]) 379 | f_agg = helper_tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 380 | return f_agg 381 | --------------------------------------------------------------------------------