├── 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 ├── process_toronto3d.py ├── 6_fold_cv.py ├── data_prepare_sensaturban.py ├── data_prepare_s3dis.py └── data_prepare_toronto3d.py ├── compile_op.sh ├── helper_requirements.txt ├── README.md ├── tester_SensatUrban.py ├── tester_S3DIS.py ├── tester_Toronto3D.py ├── helper_ply.py ├── main_S3DIS.py ├── main_S3DIS_ds.py ├── main_Toronto3D.py ├── main_Toronto3D_ds.py ├── main_SensatUrban.py ├── main_SensatUrban_ds.py ├── RandLANet_UNetp.py └── RandLANet.py /utils/cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /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.4 9 | 10 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lib.python.nearest_neighbors as nearest_neighbors 3 | import time 4 | 5 | batch_size = 16 6 | num_points = 81920 7 | K = 16 8 | pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) 9 | 10 | # nearest neighbours 11 | start = time.time() 12 | neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) 13 | print(time.time() - start) 14 | 15 | 16 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | from Cython.Distutils import build_ext 4 | import numpy 5 | 6 | 7 | 8 | ext_modules = [Extension( 9 | "nearest_neighbors", 10 | sources=["knn.pyx", "knn_.cxx",], # source file(s) 11 | include_dirs=["./", numpy.get_include()], 12 | language="c++", 13 | extra_compile_args = [ "-std=c++11", "-fopenmp",], 14 | extra_link_args=["-std=c++11", '-fopenmp'], 15 | )] 16 | 17 | setup( 18 | name = "KNN NanoFLANN", 19 | ext_modules = ext_modules, 20 | cmdclass = {'build_ext': build_ext}, 21 | ) 22 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Small but Mighty: Enhancing 3D Point Clouds Semantic Segmentation with U-Next Framework 2 | 3 | ### (1) Setup 4 | This code has been tested with Python 3.8, Tensorflow 2.4, CUDA 11.0 and cuDNN 8.0.5 on Ubuntu 20.04. 5 | 6 | - Setup python environment 7 | ``` 8 | conda create -n unext python=3.8 9 | source activate unext 10 | pip install -r helper_requirements.txt 11 | sh compile_op.sh 12 | ``` 13 | 14 | ### (2) Data Prepare 15 | ``` 16 | cd utils/ 17 | python data_prepare_$data_you_want_use$.py 18 | ``` 19 | If you want to use Toronto3D, run this before the previous command line 20 | ``` 21 | python process_toronto3d.py 22 | ``` 23 | 24 | ### (3) Train & Test 25 | If you don't want to use MDS, run this 26 | ``` 27 | python main_$data_you_want_use$.py --mode train & test 28 | ``` 29 | 30 | If you want to use MDS, run this 31 | ``` 32 | python main_$data_you_want_use$_ds.py --mode train & test 33 | ``` 34 | -------------------------------------------------------------------------------- /utils/process_toronto3d.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob, os 3 | 4 | from helper_ply import read_ply 5 | from helper_ply import write_ply 6 | 7 | data_dir = r'G:\dl_data\Toronto_3D' 8 | save_dir = r'G:\dl_data\new_toronto_3d' 9 | 10 | 11 | result_path = glob.glob(os.path.join(data_dir, '*.ply')) 12 | result_path = np.sort(result_path) 13 | 14 | for file_name in result_path: 15 | del_idx = [] 16 | orig_data = read_ply(file_name) 17 | points = np.vstack((orig_data['x'], orig_data['y'], orig_data['z'], orig_data['red'], orig_data['green'], 18 | orig_data['blue'], orig_data['scalar_Label'])).T 19 | points[:, 0] -= 627285 20 | points[:, 1] -= 4841948 21 | for i in range(points.shape[0]): 22 | if points[i, 6] == 0: 23 | del_idx.append(i) 24 | pass 25 | points = np.delete(points, del_idx, axis=0) 26 | points[:, 6] -= 1 27 | save_path = os.path.join(save_dir, file_name.split('\\')[-1][:-4] + '.ply') 28 | write_ply(save_path, points, ['x', 'y', 'z', 'red', 'green', 'blue', 'label']) 29 | -------------------------------------------------------------------------------- /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); -------------------------------------------------------------------------------- /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 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /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 | if __name__ == '__main__': 11 | base_dir = r'/media/zengziyin/ZZY/Code/LACV-Net/test/Toronto3D_8396/val_preds' 12 | original_data_dir = r'/media/zengziyin/ZZY/Data/Toronto3D/original_ply' 13 | data_path = glob.glob(os.path.join(base_dir, '*.ply')) 14 | data_path = np.sort(data_path) 15 | 16 | test_total_correct = 0 17 | test_total_seen = 0 18 | gt_classes = [0 for _ in range(13)] 19 | positive_classes = [0 for _ in range(13)] 20 | true_positive_classes = [0 for _ in range(13)] 21 | visualization = False 22 | 23 | for file_name in data_path: 24 | pred_data = read_ply(file_name) 25 | pred = pred_data['pred'] 26 | original_data = read_ply(os.path.join(original_data_dir, file_name.split('/')[-1][:-4] + '.ply')) 27 | labels = original_data['class'] 28 | points = np.vstack((original_data['x'], original_data['y'], original_data['z'])).T 29 | 30 | ################## 31 | # Visualize data # 32 | ################## 33 | if visualization: 34 | colors = np.vstack((original_data['red'], original_data['green'], original_data['blue'])).T 35 | xyzrgb = np.concatenate([points, colors], axis=-1) 36 | Plot.draw_pc(xyzrgb) # visualize raw point clouds 37 | Plot.draw_pc_sem_ins(points, labels) # visualize ground-truth 38 | Plot.draw_pc_sem_ins(points, pred) # visualize prediction 39 | 40 | correct = np.sum(pred == labels) 41 | print(str(file_name.split('/')[-1][:-4]) + '_acc:' + str(correct / float(len(labels)))) 42 | test_total_correct += correct 43 | test_total_seen += len(labels) 44 | 45 | for j in range(len(labels)): 46 | gt_l = int(labels[j]) 47 | pred_l = int(pred[j]) 48 | gt_classes[gt_l] += 1 49 | positive_classes[pred_l] += 1 50 | true_positive_classes[gt_l] += int(gt_l == pred_l) 51 | 52 | iou_list = [] 53 | for n in range(13): 54 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 55 | iou_list.append(iou) 56 | mean_iou = sum(iou_list) / 13.0 57 | print('eval accuracy: {}'.format(test_total_correct / float(test_total_seen))) 58 | print('mean IOU:{}'.format(mean_iou)) 59 | print(iou_list) 60 | 61 | acc_list = [] 62 | for n in range(13): 63 | acc = true_positive_classes[n] / float(gt_classes[n]) 64 | acc_list.append(acc) 65 | mean_acc = sum(acc_list) / 13.0 66 | print('mAcc value is :{}'.format(mean_acc)) 67 | -------------------------------------------------------------------------------- /utils/data_prepare_sensaturban.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname 2 | from sklearn.neighbors import KDTree 3 | from helper_tool import DataProcessing as DP 4 | from helper_ply import write_ply 5 | import numpy as np 6 | import os, pickle, argparse 7 | 8 | if __name__ == '__main__': 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--dataset_path', type=str, required=True, help='the number of GPUs to use [default: 0]') 11 | FLAGS = parser.parse_args() 12 | dataset_name = 'SensatUrban' 13 | dataset_path = FLAGS.dataset_path 14 | preparation_types = ['grid'] # Grid sampling & Random sampling 15 | grid_size = 0.04 16 | random_sample_ratio = 10 17 | train_files = np.sort([join(dataset_path, 'train', i) for i in os.listdir(join(dataset_path, 'train'))]) 18 | test_files = np.sort([join(dataset_path, 'test', i) for i in os.listdir(join(dataset_path, 'test'))]) 19 | files = np.sort(np.hstack((train_files, test_files))) 20 | 21 | for sample_type in preparation_types: 22 | for pc_path in files: 23 | cloud_name = pc_path.split('/')[-1][:-4] 24 | print('start to process:', cloud_name) 25 | 26 | # create output directory 27 | out_folder = join(dirname(dataset_path), sample_type + '_{:.3f}'.format(grid_size)) 28 | os.makedirs(out_folder) if not exists(out_folder) else None 29 | 30 | # check if it has already calculated 31 | if exists(join(out_folder, cloud_name + '_KDTree.pkl')): 32 | print(cloud_name, 'already exists, skipped') 33 | continue 34 | 35 | if pc_path in train_files: 36 | xyz, rgb, labels = DP.read_ply_data(pc_path, with_rgb=True) 37 | else: 38 | xyz, rgb = DP.read_ply_data(pc_path, with_rgb=True, with_label=False) 39 | labels = np.zeros(len(xyz), dtype=np.uint8) 40 | 41 | sub_ply_file = join(out_folder, cloud_name + '.ply') 42 | if sample_type == 'grid': 43 | sub_xyz, sub_rgb, sub_labels = DP.grid_sub_sampling(xyz, rgb, labels, grid_size) 44 | else: 45 | sub_xyz, sub_rgb, sub_labels = DP.random_sub_sampling(xyz, rgb, labels, random_sample_ratio) 46 | 47 | sub_rgb = sub_rgb / 255.0 48 | sub_labels = np.squeeze(sub_labels) 49 | write_ply(sub_ply_file, [sub_xyz, sub_rgb, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 50 | 51 | search_tree = KDTree(sub_xyz, leaf_size=50) 52 | kd_tree_file = join(out_folder, cloud_name + '_KDTree.pkl') 53 | with open(kd_tree_file, 'wb') as f: 54 | pickle.dump(search_tree, f) 55 | 56 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 57 | proj_idx = proj_idx.astype(np.int32) 58 | proj_save = join(out_folder, cloud_name + '_proj.pkl') 59 | with open(proj_save, 'wb') as f: 60 | pickle.dump([proj_idx, labels], f) 61 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /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 = '/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' class.. 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/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/data_prepare_toronto3d.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_ply import read_ply 13 | from helper_tool import DataProcessing as DP 14 | 15 | anno_paths = glob.glob(os.path.join('/media/cug210/data/ZZY/toronto_unprocess/new_toronto_3d', '*.ply')) 16 | 17 | gt_class = [x.rstrip() for x in open('/media/cug210/data/ZZY/toronto_unprocess/new_toronto_3d/toronto_3d_classes_8.txt')] 18 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 19 | 20 | sub_grid_size = 0.06 21 | original_pc_folder = join('/media/cug210/data/ZZY/toronto_process', 'original_ply') 22 | sub_pc_folder = join('/media/cug210/data/ZZY/toronto_process', 'input_{:.3f}'.format(sub_grid_size)) 23 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 24 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 25 | out_format = '.ply' 26 | 27 | 28 | def convert_pc2ply(anno_path, save_path): 29 | """ 30 | Convert original dataset files to ply file (each line is XYZRGBL). 31 | We aggregated all the points from each instance in the room. 32 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 33 | :param save_path: path to save original point clouds (each line is XYZRGBL) 34 | :return: None 35 | """ 36 | # data_list = [] 37 | # 38 | # for f in glob.glob(join(anno_path, '*.txt')): 39 | # class_name = os.path.basename(f).split('_')[0] 40 | # if class_name not in gt_class: # note: in some room there is 'staris' class.. 41 | # class_name = 'clutter' 42 | # pc = pd.read_csv(f, header=None, delim_whitespace=True).values 43 | # labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 44 | # data_list.append(np.concatenate([pc, labels], 1)) # Nx7 45 | # 46 | # pc_label = np.concatenate(data_list, 0) 47 | # xyz_min = np.amin(pc_label, axis=0)[0:3] 48 | # pc_label[:, 0:3] -= xyz_min 49 | 50 | process_data = read_ply(anno_path) 51 | xyz = np.vstack((process_data['x'], process_data['y'], process_data['z'])).T.astype(np.float32) 52 | colors = np.vstack((process_data['red'], process_data['green'], process_data['blue'])).T.astype(np.uint8) 53 | labels = np.expand_dims(process_data['label'], axis=-1).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[-1] 80 | print(out_file_name) 81 | convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name)) 82 | -------------------------------------------------------------------------------- /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 -------------------------------------------------------------------------------- /tester_SensatUrban.py: -------------------------------------------------------------------------------- 1 | from os import makedirs, system 2 | from os.path import exists, join, dirname, abspath 3 | from helper_ply import write_ply 4 | import tensorflow.compat.v1 as tf 5 | tf.disable_v2_behavior() 6 | import numpy as np 7 | import time 8 | 9 | 10 | def log_out(out_str, log_f_out): 11 | log_f_out.write(out_str + '\n') 12 | log_f_out.flush() 13 | print(out_str) 14 | 15 | 16 | class ModelTester: 17 | def __init__(self, model, dataset, restore_snap=None): 18 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 19 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 20 | self.Log_file = open('log_test_' + dataset.name + '.txt', 'a') 21 | 22 | # Create a session for running Ops on the Graph. 23 | on_cpu = False 24 | if on_cpu: 25 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 26 | else: 27 | c_proto = tf.ConfigProto() 28 | c_proto.gpu_options.allow_growth = True 29 | self.sess = tf.Session(config=c_proto) 30 | self.sess.run(tf.global_variables_initializer()) 31 | 32 | # Load trained model 33 | if restore_snap is not None: 34 | self.saver.restore(self.sess, restore_snap) 35 | print("Model restored from " + restore_snap) 36 | 37 | self.prob_logits = tf.nn.softmax(model.logits) 38 | 39 | # Initiate global prediction over all test clouds 40 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 41 | for l in dataset.input_labels['test']] 42 | 43 | def test(self, model, dataset, num_votes=100): 44 | 45 | # Smoothing parameter for votes 46 | test_smooth = 0.95 47 | 48 | # Initialise iterator with validation/test data 49 | self.sess.run(dataset.test_init_op) 50 | 51 | # Test saving path 52 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 53 | test_path = join('test', saving_path.split('/')[-1]) 54 | makedirs(test_path) if not exists(test_path) else None 55 | makedirs(join(test_path, 'test_preds')) if not exists(join(test_path, 'test_preds')) else None 56 | 57 | step_id = 0 58 | epoch_id = 0 59 | last_min = -0.5 60 | 61 | while last_min < num_votes: 62 | try: 63 | ops = (self.prob_logits, 64 | model.labels, 65 | model.inputs['input_inds'], 66 | model.inputs['cloud_inds']) 67 | 68 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 69 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 70 | model.config.num_classes]) 71 | 72 | for j in range(np.shape(stacked_probs)[0]): 73 | probs = stacked_probs[j, :, :] 74 | p_idx = point_idx[j, :] 75 | c_i = cloud_idx[j][0] 76 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 77 | step_id += 1 78 | 79 | except tf.errors.OutOfRangeError: 80 | 81 | new_min = np.min(dataset.min_possibility['test']) 82 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 83 | 84 | if last_min + 1 < new_min: 85 | 86 | # Update last_min 87 | last_min += 1 88 | 89 | # Show vote results (On subcloud so it is not the good values here) 90 | log_out('\nConfusion on sub clouds', self.Log_file) 91 | num_test = len(dataset.input_labels['test']) 92 | 93 | # Project predictions 94 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 95 | proj_probs_list = [] 96 | 97 | for i_test in range(num_test): 98 | # Reproject probs back to the evaluations points 99 | proj_idx = dataset.test_proj[i_test] 100 | probs = self.test_probs[i_test][proj_idx, :] 101 | proj_probs_list += [probs] 102 | 103 | # Show vote results 104 | log_out('Confusion on full clouds', self.Log_file) 105 | for i_test in range(num_test): 106 | # Get the predicted labels 107 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 108 | save_name = join(test_path, 'test_preds', dataset.input_names['test'][i_test] + '.label') 109 | preds = preds.astype(np.uint8) 110 | preds.tofile(save_name) 111 | 112 | name = dataset.input_names['test'][i_test] + '.ply' 113 | write_ply(join(test_path, 'test_preds', name), [preds], ['pred']) 114 | 115 | # creat submission files 116 | base_dir = dirname(abspath(__file__)) 117 | results_path = join(base_dir, test_path, 'test_preds') 118 | system('cd %s && zip -r %s/submission.zip *.label' % (results_path, results_path)) 119 | return 120 | 121 | self.sess.run(dataset.test_init_op) 122 | epoch_id += 1 123 | step_id = 0 124 | continue 125 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.cxx: -------------------------------------------------------------------------------- 1 | 2 | #include "knn_.h" 3 | #include "nanoflann.hpp" 4 | using namespace nanoflann; 5 | 6 | #include "KDTreeTableAdaptor.h" 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | using namespace std; 19 | 20 | 21 | 22 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 23 | const float* queries, const size_t nqueries, 24 | const size_t K, long* indices){ 25 | 26 | // create the kdtree 27 | typedef KDTreeTableAdaptor< float, float> KDTree; 28 | KDTree mat_index(npts, dim, points, 10); 29 | mat_index.index->buildIndex(); 30 | 31 | std::vector out_dists_sqr(K); 32 | std::vector out_ids(K); 33 | 34 | // iterate over the points 35 | for(size_t i=0; i resultSet(K); 38 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 39 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 40 | for(size_t j=0; j KDTree; 52 | KDTree mat_index(npts, dim, points, 10); 53 | mat_index.index->buildIndex(); 54 | 55 | 56 | // iterate over the points 57 | # pragma omp parallel for 58 | for(size_t i=0; i out_ids(K); 60 | std::vector out_dists_sqr(K); 61 | 62 | nanoflann::KNNResultSet resultSet(K); 63 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 64 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 65 | for(size_t j=0; j KDTree; 83 | KDTree mat_index(npts, dim, points, 10); 84 | 85 | mat_index.index->buildIndex(); 86 | 87 | std::vector out_dists_sqr(K); 88 | std::vector out_ids(K); 89 | 90 | // iterate over the points 91 | for(size_t i=0; i resultSet(K); 93 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 94 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 95 | for(size_t j=0; j KDTree; 116 | KDTree mat_index(npts, dim, points, 10); 117 | 118 | mat_index.index->buildIndex(); 119 | 120 | std::vector out_dists_sqr(K); 121 | std::vector out_ids(K); 122 | 123 | // iterate over the points 124 | for(size_t i=0; i resultSet(K); 126 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 127 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 128 | for(size_t j=0; j KDTree; 153 | KDTree tree(npts, dim, points, 10); 154 | tree.index->buildIndex(); 155 | 156 | vector used(npts, 0); 157 | int current_id = 0; 158 | for(size_t ptid=0; ptid possible_ids; 162 | while(possible_ids.size() == 0){ 163 | for(size_t i=0; i query(3); 178 | for(size_t i=0; i dists(K); 183 | std::vector ids(K); 184 | nanoflann::KNNResultSet resultSet(K); 185 | resultSet.init(&ids[0], &dists[0] ); 186 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 187 | 188 | for(size_t i=0; i KDTree; 221 | KDTree tree(npts, dim, points, 10); 222 | tree.index->buildIndex(); 223 | 224 | vector used(npts, 0); 225 | int current_id = 0; 226 | for(size_t ptid=0; ptid possible_ids; 230 | while(possible_ids.size() == 0){ 231 | for(size_t i=0; i query(3); 246 | for(size_t i=0; i dists(K); 251 | std::vector ids(K); 252 | nanoflann::KNNResultSet resultSet(K); 253 | resultSet.init(&ids[0], &dists[0] ); 254 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 255 | 256 | for(size_t i=0; i 34 | 35 | // ===== This example shows how to use nanoflann with these types of containers: ======= 36 | //typedef std::vector > my_vector_of_vectors_t; 37 | //typedef std::vector my_vector_of_vectors_t; // This requires #include 38 | // ===================================================================================== 39 | 40 | 41 | /** A simple vector-of-vectors adaptor for nanoflann, without duplicating the storage. 42 | * The i'th vector represents a point in the state space. 43 | * 44 | * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality for the points in the data set, allowing more compiler optimizations. 45 | * \tparam num_t The type of the point coordinates (typically, double or float). 46 | * \tparam Distance The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. 47 | * \tparam IndexType The type for indices in the KD-tree index (typically, size_t of int) 48 | */ 49 | // template 50 | // struct KDTreeVectorAdaptor 51 | // { 52 | // typedef KDTreeVectorAdaptor self_t; 53 | // typedef typename Distance::template traits::distance_t metric_t; 54 | // typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 55 | 56 | // index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 57 | // size_t dims; 58 | 59 | // /// Constructor: takes a const ref to the vector of vectors object with the data points 60 | // KDTreeVectorAdaptor(const size_t dims /* dimensionality */, const VectorType &mat, const int leaf_max_size = 10) : m_data(mat) 61 | // { 62 | // assert(mat.size() != 0); 63 | // this->dims= dims; 64 | // index = new index_t( static_cast(dims), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 65 | // index->buildIndex(); 66 | // } 67 | 68 | // ~KDTreeVectorAdaptor() { 69 | // delete index; 70 | // } 71 | 72 | // const VectorType &m_data; 73 | 74 | // /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 75 | // * Note that this is a short-cut method for index->findNeighbors(). 76 | // * The user can also call index->... methods as desired. 77 | // * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 78 | // */ 79 | // inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 80 | // { 81 | // nanoflann::KNNResultSet resultSet(num_closest); 82 | // resultSet.init(out_indices, out_distances_sq); 83 | // index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 84 | // } 85 | 86 | // /** @name Interface expected by KDTreeSingleIndexAdaptor 87 | // * @{ */ 88 | 89 | // const self_t & derived() const { 90 | // return *this; 91 | // } 92 | // self_t & derived() { 93 | // return *this; 94 | // } 95 | 96 | // // Must return the number of data points 97 | // inline size_t kdtree_get_point_count() const { 98 | // return m_data.size()/this->dims; 99 | // } 100 | 101 | // // Returns the dim'th component of the idx'th point in the class: 102 | // inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const { 103 | // return m_data[idx*this->dims + dim]; 104 | // } 105 | 106 | // // Optional bounding-box computation: return false to default to a standard bbox computation loop. 107 | // // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 108 | // // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 109 | // template 110 | // bool kdtree_get_bbox(BBOX & /*bb*/) const { 111 | // return false; 112 | // } 113 | 114 | // /** @} */ 115 | 116 | // }; // end of KDTreeVectorOfVectorsAdaptor 117 | 118 | 119 | 120 | 121 | template 122 | struct KDTreeTableAdaptor 123 | { 124 | typedef KDTreeTableAdaptor self_t; 125 | typedef typename Distance::template traits::distance_t metric_t; 126 | typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 127 | 128 | index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 129 | size_t dim; 130 | size_t npts; 131 | const TableType* m_data; 132 | 133 | /// Constructor: takes a const ref to the vector of vectors object with the data points 134 | KDTreeTableAdaptor(const size_t npts, const size_t dim, const TableType* mat, const int leaf_max_size = 10) : m_data(mat), dim(dim), npts(npts) 135 | { 136 | assert(npts != 0); 137 | index = new index_t( static_cast(dim), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 138 | index->buildIndex(); 139 | } 140 | 141 | ~KDTreeTableAdaptor() { 142 | delete index; 143 | } 144 | 145 | 146 | /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 147 | * Note that this is a short-cut method for index->findNeighbors(). 148 | * The user can also call index->... methods as desired. 149 | * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 150 | */ 151 | inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 152 | { 153 | nanoflann::KNNResultSet resultSet(num_closest); 154 | resultSet.init(out_indices, out_distances_sq); 155 | index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 156 | } 157 | 158 | /** @name Interface expected by KDTreeSingleIndexAdaptor 159 | * @{ */ 160 | 161 | const self_t & derived() const { 162 | return *this; 163 | } 164 | self_t & derived() { 165 | return *this; 166 | } 167 | 168 | // Must return the number of data points 169 | inline size_t kdtree_get_point_count() const { 170 | return this->npts; 171 | } 172 | 173 | // Returns the dim'th component of the idx'th point in the class: 174 | inline num_t kdtree_get_pt(const size_t pts_id, const size_t coord_id) const { 175 | return m_data[pts_id*this->dim + coord_id]; 176 | } 177 | 178 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 179 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 180 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 181 | template 182 | bool kdtree_get_bbox(BBOX & /*bb*/) const { 183 | return false; 184 | } 185 | 186 | /** @} */ 187 | 188 | }; // end of KDTreeVectorOfVectorsAdaptor 189 | 190 | -------------------------------------------------------------------------------- /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.compat.v1 as tf 7 | tf.disable_v2_behavior() 8 | import numpy as np 9 | import time 10 | 11 | 12 | def log_out(out_str, log_f_out): 13 | log_f_out.write(out_str + '\n') 14 | log_f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class ModelTester: 19 | def __init__(self, model, dataset, restore_snap=None): 20 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 21 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 22 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 23 | 24 | # Create a session for running Ops on the Graph. 25 | on_cpu = False 26 | if on_cpu: 27 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 28 | else: 29 | c_proto = tf.ConfigProto() 30 | c_proto.gpu_options.allow_growth = True 31 | self.sess = tf.Session(config=c_proto) 32 | self.sess.run(tf.global_variables_initializer()) 33 | 34 | # Load trained model 35 | if restore_snap is not None: 36 | self.saver.restore(self.sess, restore_snap) 37 | print("Model restored from " + restore_snap) 38 | 39 | self.prob_logits = tf.nn.softmax(model.logits) 40 | 41 | # Initiate global prediction over all test clouds 42 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 43 | for l in dataset.input_labels['validation']] 44 | 45 | def test(self, model, dataset, num_votes=100): 46 | 47 | # Smoothing parameter for votes 48 | test_smooth = 0.95 49 | 50 | # Initialise iterator with validation/test data 51 | self.sess.run(dataset.val_init_op) 52 | 53 | # Number of points per class in validation set 54 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 55 | i = 0 56 | for label_val in dataset.label_values: 57 | if label_val not in dataset.ignored_labels: 58 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 59 | i += 1 60 | 61 | # Test saving path 62 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 63 | test_path = join('test', saving_path.split('/')[-1]) 64 | makedirs(test_path) if not exists(test_path) else None 65 | makedirs(join(test_path, 'val_preds')) if not exists(join(test_path, 'val_preds')) else None 66 | 67 | step_id = 0 68 | epoch_id = 0 69 | last_min = -0.5 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 | if step_id == 0: 80 | time_start = time.time() # 记录开始时间 81 | print("-----strat-----") 82 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 83 | acc = correct / float(np.prod(np.shape(stacked_labels))) 84 | print('step' + str(step_id) + ' acc:' + str(acc)) 85 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 86 | model.config.num_classes]) 87 | 88 | for j in range(np.shape(stacked_probs)[0]): 89 | probs = stacked_probs[j, :, :] 90 | p_idx = point_idx[j, :] 91 | c_i = cloud_idx[j][0] 92 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 93 | step_id += 1 94 | 95 | except tf.errors.OutOfRangeError: 96 | time_end = time.time() # 记录结束时间 97 | time_sum = time_end - time_start 98 | print(time_sum) # 4000帧 99 | new_min = np.min(dataset.min_possibility['validation']) 100 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 101 | 102 | if last_min + 1 < new_min: 103 | 104 | # Update last_min 105 | last_min += 1 106 | 107 | # Show vote results (On subcloud so it is not the good values here) 108 | log_out('\nConfusion on sub clouds', self.Log_file) 109 | confusion_list = [] 110 | 111 | num_val = len(dataset.input_labels['validation']) 112 | 113 | for i_test in range(num_val): 114 | probs = self.test_probs[i_test] 115 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 116 | labels = dataset.input_labels['validation'][i_test] 117 | 118 | # Confs 119 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 120 | 121 | # Regroup confusions 122 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 123 | 124 | # Rescale with the right number of point per class 125 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 126 | 127 | # Compute IoUs 128 | IoUs = DP.IoU_from_confusions(C) 129 | m_IoU = np.mean(IoUs) 130 | s = '{:5.2f} | '.format(100 * m_IoU) 131 | for IoU in IoUs: 132 | s += '{:5.2f} '.format(100 * IoU) 133 | log_out(s + '\n', self.Log_file) 134 | 135 | if int(np.ceil(new_min)) % 1 == 0: 136 | 137 | # Project predictions 138 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 139 | proj_probs_list = [] 140 | 141 | for i_val in range(num_val): 142 | # Reproject probs back to the evaluations points 143 | proj_idx = dataset.val_proj[i_val] 144 | probs = self.test_probs[i_val][proj_idx, :] 145 | proj_probs_list += [probs] 146 | 147 | # Show vote results 148 | log_out('Confusion on full clouds', self.Log_file) 149 | confusion_list = [] 150 | start_time = time.time() 151 | for i_test in range(num_val): 152 | # Get the predicted labels 153 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 154 | 155 | # Confusion 156 | labels = dataset.val_labels[i_test] 157 | acc = np.sum(preds == labels) / len(labels) 158 | log_out(dataset.input_names['validation'][i_test] + ' Acc:' + str(acc), self.Log_file) 159 | 160 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 161 | name = dataset.input_names['validation'][i_test] + '.ply' 162 | write_ply(join(test_path, 'val_preds', name), [preds, labels], ['pred', 'label']) 163 | end_time = time.time() 164 | print("infer_time: ", end_time-start_time) 165 | # Regroup confusions 166 | C = np.sum(np.stack(confusion_list), axis=0) 167 | 168 | IoUs = DP.IoU_from_confusions(C) 169 | m_IoU = np.mean(IoUs) 170 | s = '{:5.2f} | '.format(100 * m_IoU) 171 | for IoU in IoUs: 172 | s += '{:5.2f} '.format(100 * IoU) 173 | log_out('-' * len(s), self.Log_file) 174 | log_out(s, self.Log_file) 175 | log_out('-' * len(s) + '\n', self.Log_file) 176 | print('finished \n') 177 | self.sess.close() 178 | return 179 | 180 | self.sess.run(dataset.val_init_op) 181 | epoch_id += 1 182 | step_id = 0 183 | continue 184 | 185 | return 186 | -------------------------------------------------------------------------------- /tester_Toronto3D.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.compat.v1 as tf 7 | tf.disable_v2_behavior() 8 | import numpy as np 9 | import time 10 | 11 | 12 | def log_out(out_str, log_f_out): 13 | log_f_out.write(out_str + '\n') 14 | log_f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class ModelTester: 19 | def __init__(self, model, dataset, restore_snap=None): 20 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 21 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 22 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 23 | 24 | # Create a session for running Ops on the Graph. 25 | on_cpu = False 26 | if on_cpu: 27 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 28 | else: 29 | c_proto = tf.ConfigProto() 30 | c_proto.gpu_options.allow_growth = True 31 | self.sess = tf.Session(config=c_proto) 32 | self.sess.run(tf.global_variables_initializer()) 33 | 34 | # Load trained model 35 | if restore_snap is not None: 36 | self.saver.restore(self.sess, restore_snap) 37 | print("Model restored from " + restore_snap) 38 | 39 | self.prob_logits = tf.nn.softmax(model.logits) 40 | 41 | # Initiate global prediction over all test clouds 42 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 43 | for l in dataset.input_labels['validation']] 44 | 45 | def test(self, model, dataset, num_votes=100): 46 | 47 | # Smoothing parameter for votes 48 | test_smooth = 0.95 49 | 50 | # Initialise iterator with validation/test data 51 | self.sess.run(dataset.val_init_op) 52 | 53 | # Number of points per class in validation set 54 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 55 | i = 0 56 | for label_val in dataset.label_values: 57 | if label_val not in dataset.ignored_labels: 58 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 59 | i += 1 60 | 61 | # Test saving path 62 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 63 | test_path = join('test', saving_path.split('/')[-1]) 64 | makedirs(test_path) if not exists(test_path) else None 65 | makedirs(join(test_path, 'val_preds')) if not exists(join(test_path, 'val_preds')) else None 66 | 67 | step_id = 0 68 | epoch_id = 0 69 | last_min = -0.5 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 | if step_id == 0: 80 | time_start = time.time() # 记录开始时间 81 | print("-----strat-----") 82 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 83 | acc = correct / float(np.prod(np.shape(stacked_labels))) 84 | print('step' + str(step_id) + ' acc:' + str(acc)) 85 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 86 | model.config.num_classes]) 87 | 88 | for j in range(np.shape(stacked_probs)[0]): 89 | probs = stacked_probs[j, :, :] 90 | p_idx = point_idx[j, :] 91 | c_i = cloud_idx[j][0] 92 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 93 | step_id += 1 94 | 95 | except tf.errors.OutOfRangeError: 96 | time_end = time.time() # 记录结束时间 97 | time_sum = time_end - time_start 98 | print(time_sum) # 4000帧 99 | new_min = np.min(dataset.min_possibility['validation']) 100 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 101 | 102 | if last_min + 1 < new_min: 103 | 104 | # Update last_min 105 | last_min += 1 106 | 107 | # Show vote results (On subcloud so it is not the good values here) 108 | log_out('\nConfusion on sub clouds', self.Log_file) 109 | confusion_list = [] 110 | 111 | num_val = len(dataset.input_labels['validation']) 112 | 113 | for i_test in range(num_val): 114 | probs = self.test_probs[i_test] 115 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 116 | labels = dataset.input_labels['validation'][i_test] 117 | 118 | # Confs 119 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 120 | 121 | # Regroup confusions 122 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 123 | 124 | # Rescale with the right number of point per class 125 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 126 | 127 | # Compute IoUs 128 | IoUs = DP.IoU_from_confusions(C) 129 | m_IoU = np.mean(IoUs) 130 | s = '{:5.2f} | '.format(100 * m_IoU) 131 | for IoU in IoUs: 132 | s += '{:5.2f} '.format(100 * IoU) 133 | log_out(s + '\n', self.Log_file) 134 | 135 | if int(np.ceil(new_min)) % 1 == 0: 136 | 137 | # Project predictions 138 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 139 | proj_probs_list = [] 140 | 141 | for i_val in range(num_val): 142 | # Reproject probs back to the evaluations points 143 | proj_idx = dataset.val_proj[i_val] 144 | probs = self.test_probs[i_val][proj_idx, :] 145 | proj_probs_list += [probs] 146 | 147 | # Show vote results 148 | log_out('Confusion on full clouds', self.Log_file) 149 | confusion_list = [] 150 | start_time = time.time() 151 | for i_test in range(num_val): 152 | # Get the predicted labels 153 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 154 | 155 | # Confusion 156 | labels = dataset.val_labels[i_test] 157 | acc = np.sum(preds == labels) / len(labels) 158 | log_out(dataset.input_names['validation'][i_test] + ' Acc:' + str(acc), self.Log_file) 159 | 160 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 161 | name = dataset.input_names['validation'][i_test] + '.ply' 162 | write_ply(join(test_path, 'val_preds', name), [preds, labels], ['pred', 'label']) 163 | end_time = time.time() 164 | print("infer_time: ", end_time-start_time) 165 | # Regroup confusions 166 | C = np.sum(np.stack(confusion_list), axis=0) 167 | 168 | IoUs = DP.IoU_from_confusions(C) 169 | m_IoU = np.mean(IoUs) 170 | s = '{:5.2f} | '.format(100 * m_IoU) 171 | for IoU in IoUs: 172 | s += '{:5.2f} '.format(100 * IoU) 173 | log_out('-' * len(s), self.Log_file) 174 | log_out(s, self.Log_file) 175 | log_out('-' * len(s) + '\n', self.Log_file) 176 | print('finished \n') 177 | self.sess.close() 178 | return 179 | 180 | self.sess.run(dataset.val_init_op) 181 | epoch_id += 1 182 | step_id = 0 183 | continue 184 | 185 | return 186 | -------------------------------------------------------------------------------- /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 | } -------------------------------------------------------------------------------- /helper_ply.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0===============================0 4 | # | PLY files reader/writer | 5 | # 0===============================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # function to read/write .ply files 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 10/02/2017 15 | # 16 | 17 | 18 | # ---------------------------------------------------------------------------------------------------------------------- 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Basic libs 26 | import numpy as np 27 | import sys 28 | 29 | 30 | # Define PLY types 31 | ply_dtypes = dict([ 32 | (b'int8', 'i1'), 33 | (b'char', 'i1'), 34 | (b'uint8', 'u1'), 35 | (b'uchar', 'u1'), 36 | (b'int16', 'i2'), 37 | (b'short', 'i2'), 38 | (b'uint16', 'u2'), 39 | (b'ushort', 'u2'), 40 | (b'int32', 'i4'), 41 | (b'int', 'i4'), 42 | (b'uint32', 'u4'), 43 | (b'uint', 'u4'), 44 | (b'float32', 'f4'), 45 | (b'float', 'f4'), 46 | (b'float64', 'f8'), 47 | (b'double', 'f8') 48 | ]) 49 | 50 | # Numpy reader format 51 | valid_formats = {'ascii': '', 'binary_big_endian': '>', 52 | 'binary_little_endian': '<'} 53 | 54 | 55 | # ---------------------------------------------------------------------------------------------------------------------- 56 | # 57 | # Functions 58 | # \***************/ 59 | # 60 | 61 | 62 | def parse_header(plyfile, ext): 63 | # Variables 64 | line = [] 65 | properties = [] 66 | num_points = None 67 | 68 | while b'end_header' not in line and line != b'': 69 | line = plyfile.readline() 70 | 71 | if b'element' in line: 72 | line = line.split() 73 | num_points = int(line[2]) 74 | 75 | elif b'property' in line: 76 | line = line.split() 77 | properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 78 | 79 | return num_points, properties 80 | 81 | 82 | def parse_mesh_header(plyfile, ext): 83 | # Variables 84 | line = [] 85 | vertex_properties = [] 86 | num_points = None 87 | num_faces = None 88 | current_element = None 89 | 90 | 91 | while b'end_header' not in line and line != b'': 92 | line = plyfile.readline() 93 | 94 | # Find point element 95 | if b'element vertex' in line: 96 | current_element = 'vertex' 97 | line = line.split() 98 | num_points = int(line[2]) 99 | 100 | elif b'element face' in line: 101 | current_element = 'face' 102 | line = line.split() 103 | num_faces = int(line[2]) 104 | 105 | elif b'property' in line: 106 | if current_element == 'vertex': 107 | line = line.split() 108 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 109 | elif current_element == 'vertex': 110 | if not line.startswith('property list uchar int'): 111 | raise ValueError('Unsupported faces property : ' + line) 112 | 113 | return num_points, num_faces, vertex_properties 114 | 115 | 116 | def read_ply(filename, triangular_mesh=False): 117 | """ 118 | Read ".ply" files 119 | 120 | Parameters 121 | ---------- 122 | filename : string 123 | the name of the file to read. 124 | 125 | Returns 126 | ------- 127 | result : array 128 | data stored in the file 129 | 130 | Examples 131 | -------- 132 | Store data in file 133 | 134 | >>> points = np.random.rand(5, 3) 135 | >>> values = np.random.randint(2, size=10) 136 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 137 | 138 | Read the file 139 | 140 | >>> data = read_ply('example.ply') 141 | >>> values = data['values'] 142 | array([0, 0, 1, 1, 0]) 143 | 144 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 145 | array([[ 0.466 0.595 0.324] 146 | [ 0.538 0.407 0.654] 147 | [ 0.850 0.018 0.988] 148 | [ 0.395 0.394 0.363] 149 | [ 0.873 0.996 0.092]]) 150 | 151 | """ 152 | 153 | with open(filename, 'rb') as plyfile: 154 | 155 | 156 | # Check if the file start with ply 157 | if b'ply' not in plyfile.readline(): 158 | raise ValueError('The file does not start whith the word ply') 159 | 160 | # get binary_little/big or ascii 161 | fmt = plyfile.readline().split()[1].decode() 162 | if fmt == "ascii": 163 | raise ValueError('The file is not binary') 164 | 165 | # get extension for building the numpy dtypes 166 | ext = valid_formats[fmt] 167 | 168 | # PointCloud reader vs mesh reader 169 | if triangular_mesh: 170 | 171 | # Parse header 172 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 173 | 174 | # Get point data 175 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 176 | 177 | # Get face data 178 | face_properties = [('k', ext + 'u1'), 179 | ('v1', ext + 'i4'), 180 | ('v2', ext + 'i4'), 181 | ('v3', ext + 'i4')] 182 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 183 | 184 | # Return vertex data and concatenated faces 185 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 186 | data = [vertex_data, faces] 187 | 188 | else: 189 | 190 | # Parse header 191 | num_points, properties = parse_header(plyfile, ext) 192 | 193 | # Get data 194 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 195 | 196 | return data 197 | 198 | 199 | def header_properties(field_list, field_names): 200 | 201 | # List of lines to write 202 | lines = [] 203 | 204 | # First line describing element vertex 205 | lines.append('element vertex %d' % field_list[0].shape[0]) 206 | 207 | # Properties lines 208 | i = 0 209 | for fields in field_list: 210 | for field in fields.T: 211 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 212 | i += 1 213 | 214 | return lines 215 | 216 | 217 | def write_ply(filename, field_list, field_names, triangular_faces=None): 218 | """ 219 | Write ".ply" files 220 | 221 | Parameters 222 | ---------- 223 | filename : string 224 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 225 | file name if it does no already have one. 226 | 227 | field_list : list, tuple, numpy array 228 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 229 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 230 | as one field. 231 | 232 | field_names : list 233 | the name of each fields as a list of strings. Has to be the same length as the number of 234 | fields. 235 | 236 | Examples 237 | -------- 238 | >>> points = np.random.rand(10, 3) 239 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 240 | 241 | >>> values = np.random.randint(2, size=10) 242 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 243 | 244 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 245 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', 'values'] 246 | >>> write_ply('example3.ply', [points, colors, values], field_names) 247 | 248 | """ 249 | 250 | # Format list input to the right form 251 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 252 | for i, field in enumerate(field_list): 253 | if field.ndim < 2: 254 | field_list[i] = field.reshape(-1, 1) 255 | if field.ndim > 2: 256 | print('fields have more than 2 dimensions') 257 | return False 258 | 259 | # check all fields have the same number of data 260 | n_points = [field.shape[0] for field in field_list] 261 | if not np.all(np.equal(n_points, n_points[0])): 262 | print('wrong field dimensions') 263 | return False 264 | 265 | # Check if field_names and field_list have same nb of column 266 | n_fields = np.sum([field.shape[1] for field in field_list]) 267 | if (n_fields != len(field_names)): 268 | print('wrong number of field names') 269 | return False 270 | 271 | # Add extension if not there 272 | if not filename.endswith('.ply'): 273 | filename += '.ply' 274 | 275 | # open in text mode to write the header 276 | with open(filename, 'w') as plyfile: 277 | 278 | # First magical word 279 | header = ['ply'] 280 | 281 | # Encoding format 282 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 283 | 284 | # Points properties description 285 | header.extend(header_properties(field_list, field_names)) 286 | 287 | # Add faces if needded 288 | if triangular_faces is not None: 289 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 290 | header.append('property list uchar int vertex_indices') 291 | 292 | # End of header 293 | header.append('end_header') 294 | 295 | # Write all lines 296 | for line in header: 297 | plyfile.write("%s\n" % line) 298 | 299 | # open in binary/append to use tofile 300 | with open(filename, 'ab') as plyfile: 301 | 302 | # Create a structured array 303 | i = 0 304 | type_list = [] 305 | for fields in field_list: 306 | for field in fields.T: 307 | type_list += [(field_names[i], field.dtype.str)] 308 | i += 1 309 | data = np.empty(field_list[0].shape[0], dtype=type_list) 310 | i = 0 311 | for fields in field_list: 312 | for field in fields.T: 313 | data[field_names[i]] = field 314 | i += 1 315 | 316 | data.tofile(plyfile) 317 | 318 | if triangular_faces is not None: 319 | triangular_faces = triangular_faces.astype(np.int32) 320 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 321 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 322 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 323 | data['0'] = triangular_faces[:, 0] 324 | data['1'] = triangular_faces[:, 1] 325 | data['2'] = triangular_faces[:, 2] 326 | data.tofile(plyfile) 327 | 328 | return True 329 | 330 | 331 | def describe_element(name, df): 332 | """ Takes the columns of the dataframe and builds a ply-like description 333 | 334 | Parameters 335 | ---------- 336 | name: str 337 | df: pandas DataFrame 338 | 339 | Returns 340 | ------- 341 | element: list[str] 342 | """ 343 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 344 | element = ['element ' + name + ' ' + str(len(df))] 345 | 346 | if name == 'face': 347 | element.append("property list uchar int points_indices") 348 | 349 | else: 350 | for i in range(len(df.columns)): 351 | # get first letter of dtype to infer format 352 | f = property_formats[str(df.dtypes[i])[0]] 353 | element.append('property ' + f + ' ' + df.columns.values[i]) 354 | 355 | return element 356 | 357 | -------------------------------------------------------------------------------- /main_S3DIS.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from RandLANet 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.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os 12 | 13 | 14 | class S3DIS: 15 | def __init__(self, test_area_idx): 16 | self.name = 'S3DIS' 17 | self.path = '/media/zengziyin/ZZY/Data/S3DIS' 18 | self.label_to_names = {0: 'ceiling', 19 | 1: 'floor', 20 | 2: 'wall', 21 | 3: 'beam', 22 | 4: 'column', 23 | 5: 'window', 24 | 6: 'door', 25 | 7: 'table', 26 | 8: 'chair', 27 | 9: 'sofa', 28 | 10: 'bookcase', 29 | 11: 'board', 30 | 12: 'clutter'} 31 | self.num_per_class = 0 32 | self.num_classes = len(self.label_to_names) 33 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 34 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 35 | self.ignored_labels = np.array([]) 36 | 37 | self.val_split = 'Area_' + str(test_area_idx) 38 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 39 | 40 | # Initiate containers 41 | self.val_proj = [] 42 | self.val_labels = [] 43 | self.possibility = {} 44 | self.min_possibility = {} 45 | self.input_trees = {'training': [], 'validation': []} 46 | self.input_colors = {'training': [], 'validation': []} 47 | self.input_labels = {'training': [], 'validation': []} 48 | self.input_names = {'training': [], 'validation': []} 49 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 50 | 51 | def load_sub_sampled_clouds(self, sub_grid_size): 52 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 53 | for i, file_path in enumerate(self.all_files): 54 | t0 = time.time() 55 | cloud_name = file_path.split('/')[-1][:-4] 56 | if self.val_split in cloud_name: 57 | cloud_split = 'validation' 58 | else: 59 | cloud_split = 'training' 60 | 61 | # Name of the input files 62 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 63 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 64 | 65 | data = read_ply(sub_ply_file) 66 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 67 | sub_labels = data['class'] 68 | 69 | # Read pkl with search tree 70 | with open(kd_tree_file, 'rb') as f: 71 | search_tree = pickle.load(f) 72 | 73 | self.input_trees[cloud_split] += [search_tree] 74 | self.input_colors[cloud_split] += [sub_colors] 75 | self.input_labels[cloud_split] += [sub_labels] 76 | self.input_names[cloud_split] += [cloud_name] 77 | 78 | size = sub_colors.shape[0] * 4 * 7 79 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 80 | 81 | print('\nPreparing reprojected indices for testing') 82 | 83 | # Get validation and test reprojected indices 84 | for i, file_path in enumerate(self.all_files): 85 | t0 = time.time() 86 | cloud_name = file_path.split('/')[-1][:-4] 87 | 88 | # Validation projection and labels 89 | if self.val_split in cloud_name: 90 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 91 | with open(proj_file, 'rb') as f: 92 | proj_idx, labels = pickle.load(f) 93 | self.val_proj += [proj_idx] 94 | self.val_labels += [labels] 95 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 96 | 97 | # Generate the input data flow 98 | def get_batch_gen(self, split): 99 | if split == 'training': 100 | num_per_epoch = cfg.train_steps * cfg.batch_size 101 | elif split == 'validation': 102 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 103 | 104 | self.possibility[split] = [] 105 | self.min_possibility[split] = [] 106 | # Random initialize 107 | for i, tree in enumerate(self.input_colors[split]): 108 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 109 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 110 | 111 | def spatially_regular_gen(): 112 | # Generator loop 113 | for i in range(num_per_epoch): 114 | 115 | # Choose the cloud with the lowest probability 116 | cloud_idx = int(np.argmin(self.min_possibility[split])) 117 | 118 | # choose the point with the minimum of possibility in the cloud as query point 119 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 120 | 121 | # Get all points within the cloud from tree structure 122 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 123 | 124 | # Center point of input region 125 | center_point = points[point_ind, :].reshape(1, -1) 126 | 127 | # Add noise to the center point 128 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 129 | pick_point = center_point + noise.astype(center_point.dtype) 130 | 131 | # Check if the number of points in the selected cloud is less than the predefined num_points 132 | if len(points) < cfg.num_points: 133 | # Query all points within the cloud 134 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 135 | else: 136 | # Query the predefined number of points 137 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 138 | 139 | # Shuffle index 140 | queried_idx = DP.shuffle_idx(queried_idx) 141 | # Get corresponding points and colors based on the index 142 | queried_pc_xyz = points[queried_idx] 143 | queried_pc_xyz = queried_pc_xyz - pick_point 144 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 145 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 146 | 147 | # Update the possibility of the selected points 148 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 149 | delta = np.square(1 - dists / np.max(dists)) 150 | self.possibility[split][cloud_idx][queried_idx] += delta 151 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 152 | 153 | # up_sampled with replacement 154 | if len(points) < cfg.num_points: 155 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 156 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 157 | 158 | if True: 159 | yield (queried_pc_xyz.astype(np.float32), 160 | queried_pc_colors.astype(np.float32), 161 | queried_pc_labels, 162 | queried_idx.astype(np.int32), 163 | np.array([cloud_idx], dtype=np.int32)) 164 | 165 | gen_func = spatially_regular_gen 166 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 167 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 168 | return gen_func, gen_types, gen_shapes 169 | 170 | @staticmethod 171 | def get_tf_mapping2(): 172 | # Collect flat inputs 173 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 174 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 175 | input_points = [] 176 | input_neighbors = [] 177 | input_pools = [] 178 | input_up_samples = [] 179 | 180 | for i in range(cfg.num_layers): 181 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 182 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 183 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 184 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 185 | input_points.append(batch_xyz) 186 | input_neighbors.append(neighbour_idx) 187 | input_pools.append(pool_i) 188 | input_up_samples.append(up_i) 189 | batch_xyz = sub_points 190 | 191 | input_list = input_points + input_neighbors + input_pools + input_up_samples 192 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 193 | 194 | return input_list 195 | 196 | return tf_map 197 | 198 | def init_input_pipeline(self): 199 | print('Initiating input pipelines') 200 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 201 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 202 | gen_function_val, _, _ = self.get_batch_gen('validation') 203 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 204 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 205 | 206 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 207 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 208 | map_func = self.get_tf_mapping2() 209 | 210 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 211 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 212 | 213 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 214 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 215 | 216 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 217 | self.flat_inputs = iter.get_next() 218 | self.train_init_op = iter.make_initializer(self.batch_train_data) 219 | self.val_init_op = iter.make_initializer(self.batch_val_data) 220 | 221 | 222 | if __name__ == '__main__': 223 | parser = argparse.ArgumentParser() 224 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 225 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 226 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 227 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 228 | FLAGS = parser.parse_args() 229 | 230 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 231 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 232 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 233 | Mode = FLAGS.mode 234 | 235 | test_area = FLAGS.test_area 236 | dataset = S3DIS(test_area) 237 | dataset.init_input_pipeline() 238 | 239 | if Mode == 'train': 240 | model = Network(dataset, cfg) 241 | model.train(dataset) 242 | elif Mode == 'test': 243 | cfg.saving = False 244 | model = Network(dataset, cfg) 245 | if FLAGS.model_path is not 'None': 246 | chosen_snap = FLAGS.model_path 247 | else: 248 | chosen_snapshot = -1 249 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 250 | chosen_folder = logs[-1] 251 | snap_path = join(chosen_folder, 'snapshots') 252 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 253 | chosen_step = np.sort(snap_steps)[-1] 254 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 255 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 256 | tester.test(model, dataset) 257 | else: 258 | ################## 259 | # Visualize data # 260 | ################## 261 | 262 | with tf.Session() as sess: 263 | sess.run(tf.global_variables_initializer()) 264 | sess.run(dataset.train_init_op) 265 | while True: 266 | flat_inputs = sess.run(dataset.flat_inputs) 267 | pc_xyz = flat_inputs[0] 268 | sub_pc_xyz = flat_inputs[1] 269 | labels = flat_inputs[21] 270 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 271 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 272 | -------------------------------------------------------------------------------- /main_S3DIS_ds.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from RandLANet_UNext_mds 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.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os 12 | 13 | 14 | class S3DIS: 15 | def __init__(self, test_area_idx): 16 | self.name = 'S3DIS' 17 | self.path = '/media/zengziyin/ZZY/Data/S3DIS' 18 | self.label_to_names = {0: 'ceiling', 19 | 1: 'floor', 20 | 2: 'wall', 21 | 3: 'beam', 22 | 4: 'column', 23 | 5: 'window', 24 | 6: 'door', 25 | 7: 'table', 26 | 8: 'chair', 27 | 9: 'sofa', 28 | 10: 'bookcase', 29 | 11: 'board', 30 | 12: 'clutter'} 31 | self.num_per_class = 0 32 | self.num_classes = len(self.label_to_names) 33 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 34 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 35 | self.ignored_labels = np.array([]) 36 | 37 | self.val_split = 'Area_' + str(test_area_idx) 38 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 39 | 40 | # Initiate containers 41 | self.val_proj = [] 42 | self.val_labels = [] 43 | self.possibility = {} 44 | self.min_possibility = {} 45 | self.input_trees = {'training': [], 'validation': []} 46 | self.input_colors = {'training': [], 'validation': []} 47 | self.input_labels = {'training': [], 'validation': []} 48 | self.input_names = {'training': [], 'validation': []} 49 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 50 | 51 | def load_sub_sampled_clouds(self, sub_grid_size): 52 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 53 | for i, file_path in enumerate(self.all_files): 54 | t0 = time.time() 55 | cloud_name = file_path.split('/')[-1][:-4] 56 | if self.val_split in cloud_name: 57 | cloud_split = 'validation' 58 | else: 59 | cloud_split = 'training' 60 | 61 | # Name of the input files 62 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 63 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 64 | 65 | data = read_ply(sub_ply_file) 66 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 67 | sub_labels = data['class'] 68 | 69 | # Read pkl with search tree 70 | with open(kd_tree_file, 'rb') as f: 71 | search_tree = pickle.load(f) 72 | 73 | self.input_trees[cloud_split] += [search_tree] 74 | self.input_colors[cloud_split] += [sub_colors] 75 | self.input_labels[cloud_split] += [sub_labels] 76 | self.input_names[cloud_split] += [cloud_name] 77 | 78 | size = sub_colors.shape[0] * 4 * 7 79 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 80 | 81 | print('\nPreparing reprojected indices for testing') 82 | 83 | # Get validation and test reprojected indices 84 | for i, file_path in enumerate(self.all_files): 85 | t0 = time.time() 86 | cloud_name = file_path.split('/')[-1][:-4] 87 | 88 | # Validation projection and labels 89 | if self.val_split in cloud_name: 90 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 91 | with open(proj_file, 'rb') as f: 92 | proj_idx, labels = pickle.load(f) 93 | self.val_proj += [proj_idx] 94 | self.val_labels += [labels] 95 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 96 | 97 | # Generate the input data flow 98 | def get_batch_gen(self, split): 99 | if split == 'training': 100 | num_per_epoch = cfg.train_steps * cfg.batch_size 101 | elif split == 'validation': 102 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 103 | 104 | self.possibility[split] = [] 105 | self.min_possibility[split] = [] 106 | # Random initialize 107 | for i, tree in enumerate(self.input_colors[split]): 108 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 109 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 110 | 111 | def spatially_regular_gen(): 112 | # Generator loop 113 | for i in range(num_per_epoch): 114 | 115 | # Choose the cloud with the lowest probability 116 | cloud_idx = int(np.argmin(self.min_possibility[split])) 117 | 118 | # choose the point with the minimum of possibility in the cloud as query point 119 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 120 | 121 | # Get all points within the cloud from tree structure 122 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 123 | 124 | # Center point of input region 125 | center_point = points[point_ind, :].reshape(1, -1) 126 | 127 | # Add noise to the center point 128 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 129 | pick_point = center_point + noise.astype(center_point.dtype) 130 | 131 | # Check if the number of points in the selected cloud is less than the predefined num_points 132 | if len(points) < cfg.num_points: 133 | # Query all points within the cloud 134 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 135 | else: 136 | # Query the predefined number of points 137 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 138 | 139 | # Shuffle index 140 | queried_idx = DP.shuffle_idx(queried_idx) 141 | # Get corresponding points and colors based on the index 142 | queried_pc_xyz = points[queried_idx] 143 | queried_pc_xyz = queried_pc_xyz - pick_point 144 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 145 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 146 | 147 | # Update the possibility of the selected points 148 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 149 | delta = np.square(1 - dists / np.max(dists)) 150 | self.possibility[split][cloud_idx][queried_idx] += delta 151 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 152 | 153 | # up_sampled with replacement 154 | if len(points) < cfg.num_points: 155 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 156 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 157 | 158 | if True: 159 | yield (queried_pc_xyz.astype(np.float32), 160 | queried_pc_colors.astype(np.float32), 161 | queried_pc_labels, 162 | queried_idx.astype(np.int32), 163 | np.array([cloud_idx], dtype=np.int32)) 164 | 165 | gen_func = spatially_regular_gen 166 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 167 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 168 | return gen_func, gen_types, gen_shapes 169 | 170 | @staticmethod 171 | def get_tf_mapping2(): 172 | # Collect flat inputs 173 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 174 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 175 | input_points = [] 176 | input_neighbors = [] 177 | input_pools = [] 178 | input_up_samples = [] 179 | input_labels = [] 180 | ori_batch_labels = batch_labels 181 | 182 | for i in range(cfg.num_layers): 183 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 184 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 185 | sub_labels = batch_labels[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i]] 186 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 187 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 188 | input_points.append(batch_xyz) 189 | input_labels.append(batch_labels) 190 | input_neighbors.append(neighbour_idx) 191 | input_pools.append(pool_i) 192 | input_up_samples.append(up_i) 193 | batch_xyz = sub_points 194 | batch_labels = sub_labels 195 | 196 | input_list = input_points + input_neighbors + input_pools + input_up_samples 197 | input_list += [batch_features, ori_batch_labels, batch_pc_idx, batch_cloud_idx] 198 | input_list += input_labels 199 | 200 | return input_list 201 | 202 | return tf_map 203 | 204 | def init_input_pipeline(self): 205 | print('Initiating input pipelines') 206 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 207 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 208 | gen_function_val, _, _ = self.get_batch_gen('validation') 209 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 210 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 211 | 212 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 213 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 214 | map_func = self.get_tf_mapping2() 215 | 216 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 217 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 218 | 219 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 220 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 221 | 222 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 223 | self.flat_inputs = iter.get_next() 224 | self.train_init_op = iter.make_initializer(self.batch_train_data) 225 | self.val_init_op = iter.make_initializer(self.batch_val_data) 226 | 227 | 228 | if __name__ == '__main__': 229 | parser = argparse.ArgumentParser() 230 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 231 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 232 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 233 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 234 | FLAGS = parser.parse_args() 235 | 236 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 237 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 238 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 239 | Mode = FLAGS.mode 240 | 241 | test_area = FLAGS.test_area 242 | dataset = S3DIS(test_area) 243 | dataset.init_input_pipeline() 244 | 245 | if Mode == 'train': 246 | model = Network(dataset, cfg) 247 | model.train(dataset) 248 | elif Mode == 'test': 249 | cfg.saving = False 250 | model = Network(dataset, cfg) 251 | if FLAGS.model_path is not 'None': 252 | chosen_snap = FLAGS.model_path 253 | else: 254 | chosen_snapshot = -1 255 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 256 | chosen_folder = logs[-1] 257 | snap_path = join(chosen_folder, 'snapshots') 258 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 259 | chosen_step = np.sort(snap_steps)[-1] 260 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 261 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 262 | tester.test(model, dataset) 263 | else: 264 | ################## 265 | # Visualize data # 266 | ################## 267 | 268 | with tf.Session() as sess: 269 | sess.run(tf.global_variables_initializer()) 270 | sess.run(dataset.train_init_op) 271 | while True: 272 | flat_inputs = sess.run(dataset.flat_inputs) 273 | pc_xyz = flat_inputs[0] 274 | sub_pc_xyz = flat_inputs[1] 275 | labels = flat_inputs[21] 276 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 277 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 278 | -------------------------------------------------------------------------------- /main_Toronto3D.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from RandLANet import Network 3 | from tester_Toronto3D import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import ConfigToronto3D as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os 12 | 13 | 14 | class Toronto3D: 15 | def __init__(self, test_area_idx): 16 | self.name = 'Toronto3D' 17 | self.path = '/media/zengziyin/ZZY/Data/Toronto3D' 18 | self.label_to_names = {0: 'road', 19 | 1: 'road_marking', 20 | 2: 'natural', 21 | 3: 'building', 22 | 4: 'utility_line', 23 | 5: 'pole', 24 | 6: 'car', 25 | 7: 'fence'} 26 | self.num_classes = len(self.label_to_names) 27 | self.num_per_class = np.zeros(self.num_classes) 28 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 29 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 30 | self.ignored_labels = np.array([]) 31 | 32 | self.val_split = 'L00' + str(test_area_idx) 33 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 34 | 35 | # Initiate containers 36 | self.val_proj = [] 37 | self.val_labels = [] 38 | self.possibility = {} 39 | self.min_possibility = {} 40 | self.input_trees = {'training': [], 'validation': []} 41 | self.input_colors = {'training': [], 'validation': []} 42 | self.input_labels = {'training': [], 'validation': []} 43 | self.input_names = {'training': [], 'validation': []} 44 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 45 | 46 | def load_sub_sampled_clouds(self, sub_grid_size): 47 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 48 | for i, file_path in enumerate(self.all_files): 49 | t0 = time.time() 50 | cloud_name = file_path.split('/')[-1][:-4] 51 | if self.val_split in cloud_name: 52 | cloud_split = 'validation' 53 | print(self.val_split, cloud_name) 54 | else: 55 | cloud_split = 'training' 56 | 57 | # Name of the input files 58 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 59 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 60 | 61 | data = read_ply(sub_ply_file) 62 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 63 | sub_labels = data['class'] 64 | 65 | self.num_per_class += DP.get_num_class_from_label(sub_labels, self.num_classes) 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 predefined 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 | input_list = input_points + input_neighbors + input_pools + input_up_samples 190 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 191 | 192 | return input_list 193 | 194 | return tf_map 195 | 196 | # @staticmethod 197 | # def get_tf_mapping2(): 198 | # # Collect flat inputs 199 | # def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 200 | # batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 201 | # input_list = [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 202 | # 203 | # return input_list 204 | 205 | return tf_map 206 | 207 | def init_input_pipeline(self): 208 | print('Initiating input pipelines') 209 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 210 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 211 | gen_function_val, _, _ = self.get_batch_gen('validation') 212 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 213 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 214 | 215 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 216 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 217 | map_func = self.get_tf_mapping2() 218 | 219 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 220 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 221 | 222 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 223 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 224 | 225 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 226 | self.flat_inputs = iter.get_next() 227 | self.train_init_op = iter.make_initializer(self.batch_train_data) 228 | self.val_init_op = iter.make_initializer(self.batch_val_data) 229 | 230 | 231 | if __name__ == '__main__': 232 | parser = argparse.ArgumentParser() 233 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 234 | parser.add_argument('--test_area', type=int, default=2, help='Which area to use for test, option: 1-6 [default: 5]') 235 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 236 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 237 | FLAGS = parser.parse_args() 238 | 239 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 240 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 241 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 242 | Mode = FLAGS.mode 243 | 244 | test_area = FLAGS.test_area 245 | dataset = Toronto3D(test_area) 246 | dataset.init_input_pipeline() 247 | 248 | if Mode == 'train': 249 | model = Network(dataset, cfg) 250 | model.train(dataset) 251 | elif Mode == 'test': 252 | cfg.saving = False 253 | model = Network(dataset, cfg) 254 | if FLAGS.model_path is not 'None': 255 | chosen_snap = FLAGS.model_path 256 | else: 257 | chosen_snapshot = -1 258 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 259 | chosen_folder = logs[-1] 260 | snap_path = join(chosen_folder, 'snapshots') 261 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 262 | chosen_step = np.sort(snap_steps)[-1] 263 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 264 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 265 | tester.test(model, dataset) 266 | else: 267 | ################## 268 | # Visualize data # 269 | ################## 270 | 271 | with tf.Session() as sess: 272 | sess.run(tf.global_variables_initializer()) 273 | sess.run(dataset.train_init_op) 274 | while True: 275 | flat_inputs = sess.run(dataset.flat_inputs) 276 | pc_xyz = flat_inputs[0] 277 | sub_pc_xyz = flat_inputs[1] 278 | labels = flat_inputs[21] 279 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 280 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 281 | -------------------------------------------------------------------------------- /main_Toronto3D_ds.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from RandLANet_UNext_mds import Network 3 | from tester_Toronto3D import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import ConfigToronto3D as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os 12 | 13 | 14 | class Toronto3D: 15 | def __init__(self, test_area_idx): 16 | self.name = 'Toronto3D' 17 | self.path = '/media/zengziyin/ZZY/Data/Toronto3D' 18 | self.label_to_names = {0: 'road', 19 | 1: 'road_marking', 20 | 2: 'natural', 21 | 3: 'building', 22 | 4: 'utility_line', 23 | 5: 'pole', 24 | 6: 'car', 25 | 7: 'fence'} 26 | self.num_classes = len(self.label_to_names) 27 | self.num_per_class = np.zeros(self.num_classes) 28 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 29 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 30 | self.ignored_labels = np.array([]) 31 | 32 | self.val_split = 'L00' + str(test_area_idx) 33 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 34 | 35 | # Initiate containers 36 | self.val_proj = [] 37 | self.val_labels = [] 38 | self.possibility = {} 39 | self.min_possibility = {} 40 | self.input_trees = {'training': [], 'validation': []} 41 | self.input_colors = {'training': [], 'validation': []} 42 | self.input_labels = {'training': [], 'validation': []} 43 | self.input_names = {'training': [], 'validation': []} 44 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 45 | 46 | def load_sub_sampled_clouds(self, sub_grid_size): 47 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 48 | for i, file_path in enumerate(self.all_files): 49 | t0 = time.time() 50 | cloud_name = file_path.split('/')[-1][:-4] 51 | if self.val_split in cloud_name: 52 | cloud_split = 'validation' 53 | print(self.val_split, cloud_name) 54 | else: 55 | cloud_split = 'training' 56 | 57 | # Name of the input files 58 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 59 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 60 | 61 | data = read_ply(sub_ply_file) 62 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 63 | sub_labels = data['class'] 64 | 65 | self.num_per_class += DP.get_num_class_from_label(sub_labels, self.num_classes) 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 predefined 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 | input_labels = [] 178 | ori_batch_labels = batch_labels 179 | 180 | for i in range(cfg.num_layers): 181 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 182 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 183 | sub_labels = batch_labels[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i]] 184 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 185 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 186 | input_points.append(batch_xyz) 187 | input_labels.append(batch_labels) 188 | input_neighbors.append(neighbour_idx) 189 | input_pools.append(pool_i) 190 | input_up_samples.append(up_i) 191 | batch_xyz = sub_points 192 | batch_labels = sub_labels 193 | 194 | input_list = input_points + input_neighbors + input_pools + input_up_samples 195 | input_list += [batch_features, ori_batch_labels, batch_pc_idx, batch_cloud_idx] 196 | input_list += input_labels 197 | 198 | return input_list 199 | 200 | return tf_map 201 | 202 | # @staticmethod 203 | # def get_tf_mapping2(): 204 | # # Collect flat inputs 205 | # def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 206 | # batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 207 | # input_list = [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 208 | # 209 | # return input_list 210 | 211 | return tf_map 212 | 213 | def init_input_pipeline(self): 214 | print('Initiating input pipelines') 215 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 216 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 217 | gen_function_val, _, _ = self.get_batch_gen('validation') 218 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 219 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 220 | 221 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 222 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 223 | map_func = self.get_tf_mapping2() 224 | 225 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 226 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 227 | 228 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 229 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 230 | 231 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 232 | self.flat_inputs = iter.get_next() 233 | self.train_init_op = iter.make_initializer(self.batch_train_data) 234 | self.val_init_op = iter.make_initializer(self.batch_val_data) 235 | 236 | 237 | if __name__ == '__main__': 238 | parser = argparse.ArgumentParser() 239 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 240 | parser.add_argument('--test_area', type=int, default=2, help='Which area to use for test, option: 1-6 [default: 5]') 241 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 242 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 243 | FLAGS = parser.parse_args() 244 | 245 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 246 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 247 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 248 | Mode = FLAGS.mode 249 | 250 | test_area = FLAGS.test_area 251 | dataset = Toronto3D(test_area) 252 | dataset.init_input_pipeline() 253 | 254 | if Mode == 'train': 255 | model = Network(dataset, cfg) 256 | model.train(dataset) 257 | elif Mode == 'test': 258 | cfg.saving = False 259 | model = Network(dataset, cfg) 260 | if FLAGS.model_path is not 'None': 261 | chosen_snap = FLAGS.model_path 262 | else: 263 | chosen_snapshot = -1 264 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 265 | chosen_folder = logs[-1] 266 | snap_path = join(chosen_folder, 'snapshots') 267 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 268 | chosen_step = np.sort(snap_steps)[-1] 269 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 270 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 271 | tester.test(model, dataset) 272 | else: 273 | ################## 274 | # Visualize data # 275 | ################## 276 | 277 | with tf.Session() as sess: 278 | sess.run(tf.global_variables_initializer()) 279 | sess.run(dataset.train_init_op) 280 | while True: 281 | flat_inputs = sess.run(dataset.flat_inputs) 282 | pc_xyz = flat_inputs[0] 283 | sub_pc_xyz = flat_inputs[1] 284 | labels = flat_inputs[21] 285 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 286 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 287 | -------------------------------------------------------------------------------- /main_SensatUrban.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname, abspath 2 | from RandLANet import Network 3 | from tester_SensatUrban import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import ConfigSensatUrban as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os, shutil 12 | 13 | 14 | class SensatUrban: 15 | def __init__(self): 16 | self.name = 'SensatUrban' 17 | root_path = '/media/zengziyin/ZZY/Data/SensatUrban' # path to the dataset 18 | self.path = root_path 19 | self.label_to_names = {0: 'Ground', 1: 'High Vegetation', 2: 'Buildings', 3: 'Walls', 20 | 4: 'Bridge', 5: 'Parking', 6: 'Rail', 7: 'traffic Roads', 8: 'Street Furniture', 21 | 9: 'Cars', 10: 'Footpath', 11: 'Bikes', 12: 'Water'} 22 | self.num_classes = len(self.label_to_names) 23 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 24 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 25 | self.ignored_labels = np.array([]) 26 | 27 | self.all_files = np.sort(glob.glob(join(self.path, 'original_block_ply', '*.ply'))) 28 | self.val_file_name = ['birmingham_block_1', 29 | 'birmingham_block_5', 30 | 'cambridge_block_10', 31 | 'cambridge_block_7'] 32 | self.test_file_name = ['birmingham_block_2', 'birmingham_block_8', 33 | 'cambridge_block_15', 'cambridge_block_22', 34 | 'cambridge_block_16', 'cambridge_block_27'] 35 | self.use_val = True # whether use validation set or not 36 | 37 | # initialize 38 | self.num_per_class = np.zeros(self.num_classes) 39 | self.val_proj = [] 40 | self.val_labels = [] 41 | self.test_proj = [] 42 | self.test_labels = [] 43 | self.possibility = {} 44 | self.min_possibility = {} 45 | self.input_trees = {'training': [], 'validation': [], 'test': []} 46 | self.input_colors = {'training': [], 'validation': [], 'test': []} 47 | self.input_labels = {'training': [], 'validation': [], 'test': []} 48 | self.input_names = {'training': [], 'validation': [], 'test': []} 49 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 50 | for ignore_label in self.ignored_labels: 51 | self.num_per_class = np.delete(self.num_per_class, ignore_label) 52 | 53 | def load_sub_sampled_clouds(self, sub_grid_size): 54 | tree_path = join(self.path, 'grid_{:.3f}'.format(sub_grid_size)) 55 | print(tree_path) 56 | 57 | for i, file_path in enumerate(self.all_files): 58 | t0 = time.time() 59 | cloud_name = file_path.split('/')[-1][:-4] 60 | if cloud_name in self.test_file_name: 61 | cloud_split = 'test' 62 | elif cloud_name in self.val_file_name: 63 | cloud_split = 'validation' 64 | else: 65 | cloud_split = 'training' 66 | 67 | # Name of the input files 68 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 69 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 70 | 71 | data = read_ply(sub_ply_file) 72 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 73 | sub_labels = data['class'] 74 | 75 | # compute num_per_class in training set 76 | if cloud_split == 'training': 77 | self.num_per_class += DP.get_num_class_from_label(sub_labels, self.num_classes) 78 | 79 | # Read pkl with search tree 80 | with open(kd_tree_file, 'rb') as f: 81 | search_tree = pickle.load(f) 82 | 83 | self.input_trees[cloud_split] += [search_tree] 84 | self.input_colors[cloud_split] += [sub_colors] 85 | self.input_labels[cloud_split] += [sub_labels] 86 | self.input_names[cloud_split] += [cloud_name] 87 | 88 | size = sub_colors.shape[0] * 4 * 7 89 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 90 | 91 | print('\nPreparing reprojected indices for testing') 92 | 93 | # Get validation and test reprojected indices 94 | 95 | for i, file_path in enumerate(self.all_files): 96 | t0 = time.time() 97 | cloud_name = file_path.split('/')[-1][:-4] 98 | 99 | # val projection and labels 100 | if cloud_name in self.val_file_name: 101 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 102 | with open(proj_file, 'rb') as f: 103 | proj_idx, labels = pickle.load(f) 104 | self.val_proj += [proj_idx] 105 | self.val_labels += [labels] 106 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 107 | 108 | # test projection and labels 109 | if cloud_name in self.test_file_name: 110 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 111 | with open(proj_file, 'rb') as f: 112 | proj_idx, labels = pickle.load(f) 113 | self.test_proj += [proj_idx] 114 | self.test_labels += [labels] 115 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 116 | 117 | def get_batch_gen(self, split): 118 | if split == 'training': 119 | num_per_epoch = cfg.train_steps * cfg.batch_size 120 | elif split == 'validation': 121 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 122 | else: 123 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 124 | 125 | # Reset possibility 126 | self.possibility[split] = [] 127 | self.min_possibility[split] = [] 128 | for i, tree in enumerate(self.input_colors[split]): 129 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 130 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 131 | 132 | def spatially_regular_gen(): 133 | # Generator loop 134 | for i in range(num_per_epoch): # num_per_epoch 135 | 136 | # Choose a random cloud 137 | cloud_idx = int(np.argmin(self.min_possibility[split])) 138 | 139 | # choose the point with the minimum of possibility as query point 140 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 141 | 142 | # Get points from tree structure 143 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 144 | 145 | # Center point of input region 146 | center_point = points[point_ind, :].reshape(1, -1) 147 | 148 | # Add noise to the center point 149 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 150 | pick_point = center_point + noise.astype(center_point.dtype) 151 | 152 | if len(points) < cfg.num_points: 153 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 154 | else: 155 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 156 | 157 | queried_idx = DP.shuffle_idx(queried_idx) 158 | # Collect points and colors 159 | queried_pc_xyz = points[queried_idx] 160 | queried_pc_xyz = queried_pc_xyz - pick_point 161 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 162 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 163 | 164 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 165 | delta = np.square(1 - dists / np.max(dists)) 166 | self.possibility[split][cloud_idx][queried_idx] += delta 167 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 168 | 169 | if len(points) < cfg.num_points: 170 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 171 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 172 | 173 | if True: 174 | yield (queried_pc_xyz.astype(np.float32), 175 | queried_pc_colors.astype(np.float32), 176 | queried_pc_labels, 177 | queried_idx.astype(np.int32), 178 | np.array([cloud_idx], dtype=np.int32)) 179 | 180 | gen_func = spatially_regular_gen 181 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 182 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 183 | return gen_func, gen_types, gen_shapes 184 | 185 | @staticmethod 186 | def get_tf_mapping2(): 187 | # Collect flat inputs 188 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 189 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 190 | input_points = [] 191 | input_neighbors = [] 192 | input_pools = [] 193 | input_up_samples = [] 194 | 195 | for i in range(cfg.num_layers): 196 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 197 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 198 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 199 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 200 | input_points.append(batch_xyz) 201 | input_neighbors.append(neighbour_idx) 202 | input_pools.append(pool_i) 203 | input_up_samples.append(up_i) 204 | batch_xyz = sub_points 205 | 206 | input_list = input_points + input_neighbors + input_pools + input_up_samples 207 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 208 | 209 | return input_list 210 | 211 | return tf_map 212 | 213 | def init_input_pipeline(self): 214 | print('Initiating input pipelines') 215 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 216 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 217 | gen_function_val, _, _ = self.get_batch_gen('validation') 218 | gen_function_test, _, _ = self.get_batch_gen('test') 219 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 220 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 221 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 222 | 223 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 224 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 225 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 226 | map_func = self.get_tf_mapping2() 227 | 228 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 229 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 230 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 231 | 232 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 233 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 234 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 235 | 236 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 237 | self.flat_inputs = iter.get_next() 238 | self.train_init_op = iter.make_initializer(self.batch_train_data) 239 | self.val_init_op = iter.make_initializer(self.batch_val_data) 240 | self.test_init_op = iter.make_initializer(self.batch_test_data) 241 | 242 | 243 | if __name__ == '__main__': 244 | parser = argparse.ArgumentParser() 245 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 246 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 247 | FLAGS = parser.parse_args() 248 | 249 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 250 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 251 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 252 | Mode = FLAGS.mode 253 | 254 | # shutil.rmtree('__pycache__') if exists('__pycache__') else None 255 | # if Mode == 'train': 256 | # shutil.rmtree('results') if exists('results') else None 257 | # shutil.rmtree('train_log') if exists('train_log') else None 258 | # for f in os.listdir(dirname(abspath(__file__))): 259 | # if f.startswith('log_'): 260 | # os.remove(f) 261 | 262 | dataset = SensatUrban() 263 | dataset.init_input_pipeline() 264 | 265 | if Mode == 'train': 266 | model = Network(dataset, cfg) 267 | model.train(dataset) 268 | elif Mode == 'test': 269 | cfg.saving = False 270 | model = Network(dataset, cfg) 271 | chosen_snapshot = -1 272 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 273 | chosen_folder = logs[-1] 274 | snap_path = join(chosen_folder, 'snapshots') 275 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 276 | chosen_step = np.sort(snap_steps)[-1] 277 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 278 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 279 | tester.test(model, dataset) 280 | shutil.rmtree('train_log') if exists('train_log') else None 281 | 282 | else: 283 | 284 | with tf.Session() as sess: 285 | sess.run(tf.global_variables_initializer()) 286 | sess.run(dataset.train_init_op) 287 | while True: 288 | data_list = sess.run(dataset.flat_inputs) 289 | xyz = data_list[0] 290 | sub_xyz = data_list[1] 291 | label = data_list[21] 292 | Plot.draw_pc_sem_ins(xyz[0, :, :], label[0, :]) -------------------------------------------------------------------------------- /main_SensatUrban_ds.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists, dirname, abspath 2 | from RandLANet_UNext_mds import Network 3 | from tester_SensatUrban import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import ConfigSensatUrban as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow.compat.v1 as tf 9 | tf.disable_v2_behavior() 10 | import numpy as np 11 | import time, pickle, argparse, glob, os, shutil 12 | 13 | 14 | class SensatUrban: 15 | def __init__(self): 16 | self.name = 'SensatUrban' 17 | root_path = '/media/zengziyin/ZZY/Data/ours_process' # path to the dataset 18 | self.path = root_path 19 | self.label_to_names = {0: 'Ground', 1: 'High Vegetation', 2: 'Buildings', 3: 'Walls', 20 | 4: 'Bridge', 5: 'Parking', 6: 'Rail', 7: 'traffic Roads', 8: 'Street Furniture', 21 | 9: 'Cars', 10: 'Footpath', 11: 'Bikes', 12: 'Water'} 22 | self.num_classes = len(self.label_to_names) 23 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 24 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 25 | self.ignored_labels = np.array([]) 26 | 27 | self.all_files = np.sort(glob.glob(join(self.path, 'original_block_ply', '*.ply'))) 28 | self.val_file_name = ['birmingham_block_1', 29 | 'birmingham_block_5', 30 | 'cambridge_block_10', 31 | 'cambridge_block_7'] 32 | self.test_file_name = ['birmingham_block_2', 'birmingham_block_8', 33 | 'cambridge_block_15', 'cambridge_block_22', 34 | 'cambridge_block_16', 'cambridge_block_27'] 35 | self.use_val = True # whether use validation set or not 36 | 37 | # initialize 38 | self.num_per_class = np.zeros(self.num_classes) 39 | self.val_proj = [] 40 | self.val_labels = [] 41 | self.test_proj = [] 42 | self.test_labels = [] 43 | self.possibility = {} 44 | self.min_possibility = {} 45 | self.input_trees = {'training': [], 'validation': [], 'test': []} 46 | self.input_colors = {'training': [], 'validation': [], 'test': []} 47 | self.input_labels = {'training': [], 'validation': [], 'test': []} 48 | self.input_names = {'training': [], 'validation': [], 'test': []} 49 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 50 | for ignore_label in self.ignored_labels: 51 | self.num_per_class = np.delete(self.num_per_class, ignore_label) 52 | 53 | def load_sub_sampled_clouds(self, sub_grid_size): 54 | tree_path = join(self.path, 'grid_{:.3f}'.format(sub_grid_size)) 55 | print(tree_path) 56 | 57 | for i, file_path in enumerate(self.all_files): 58 | t0 = time.time() 59 | cloud_name = file_path.split('/')[-1][:-4] 60 | print(cloud_name) 61 | if cloud_name in self.test_file_name: 62 | cloud_split = 'test' 63 | elif cloud_name in self.val_file_name: 64 | cloud_split = 'validation' 65 | else: 66 | cloud_split = 'training' 67 | 68 | # Name of the input files 69 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 70 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 71 | 72 | data = read_ply(sub_ply_file) 73 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 74 | sub_labels = data['class'] 75 | 76 | # compute num_per_class in training set 77 | if cloud_split == 'training': 78 | self.num_per_class += DP.get_num_class_from_label(sub_labels, self.num_classes) 79 | 80 | # Read pkl with search tree 81 | with open(kd_tree_file, 'rb') as f: 82 | search_tree = pickle.load(f) 83 | 84 | self.input_trees[cloud_split] += [search_tree] 85 | self.input_colors[cloud_split] += [sub_colors] 86 | self.input_labels[cloud_split] += [sub_labels] 87 | self.input_names[cloud_split] += [cloud_name] 88 | 89 | size = sub_colors.shape[0] * 4 * 7 90 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 91 | 92 | print('\nPreparing reprojected indices for testing') 93 | 94 | # Get validation and test reprojected indices 95 | 96 | for i, file_path in enumerate(self.all_files): 97 | t0 = time.time() 98 | cloud_name = file_path.split('/')[-1][:-4] 99 | 100 | # val projection and labels 101 | if cloud_name in self.val_file_name: 102 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 103 | with open(proj_file, 'rb') as f: 104 | proj_idx, labels = pickle.load(f) 105 | self.val_proj += [proj_idx] 106 | self.val_labels += [labels] 107 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 108 | 109 | # test projection and labels 110 | if cloud_name in self.test_file_name: 111 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 112 | with open(proj_file, 'rb') as f: 113 | proj_idx, labels = pickle.load(f) 114 | self.test_proj += [proj_idx] 115 | self.test_labels += [labels] 116 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 117 | 118 | def get_batch_gen(self, split): 119 | if split == 'training': 120 | num_per_epoch = cfg.train_steps * cfg.batch_size 121 | elif split == 'validation': 122 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 123 | else: 124 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 125 | 126 | # Reset possibility 127 | self.possibility[split] = [] 128 | self.min_possibility[split] = [] 129 | for i, tree in enumerate(self.input_colors[split]): 130 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 131 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 132 | 133 | def spatially_regular_gen(): 134 | # Generator loop 135 | for i in range(num_per_epoch): # num_per_epoch 136 | 137 | # Choose a random cloud 138 | cloud_idx = int(np.argmin(self.min_possibility[split])) 139 | 140 | # choose the point with the minimum of possibility as query point 141 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 142 | 143 | # Get points from tree structure 144 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 145 | 146 | # Center point of input region 147 | center_point = points[point_ind, :].reshape(1, -1) 148 | 149 | # Add noise to the center point 150 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 151 | pick_point = center_point + noise.astype(center_point.dtype) 152 | 153 | if len(points) < cfg.num_points: 154 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 155 | else: 156 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 157 | 158 | queried_idx = DP.shuffle_idx(queried_idx) 159 | # Collect points and colors 160 | queried_pc_xyz = points[queried_idx] 161 | queried_pc_xyz = queried_pc_xyz - pick_point 162 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 163 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 164 | 165 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 166 | delta = np.square(1 - dists / np.max(dists)) 167 | self.possibility[split][cloud_idx][queried_idx] += delta 168 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 169 | 170 | if len(points) < cfg.num_points: 171 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 172 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 173 | 174 | if True: 175 | yield (queried_pc_xyz.astype(np.float32), 176 | queried_pc_colors.astype(np.float32), 177 | queried_pc_labels, 178 | queried_idx.astype(np.int32), 179 | np.array([cloud_idx], dtype=np.int32)) 180 | 181 | gen_func = spatially_regular_gen 182 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 183 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 184 | return gen_func, gen_types, gen_shapes 185 | 186 | @staticmethod 187 | def get_tf_mapping2(): 188 | # Collect flat inputs 189 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 190 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 191 | input_points = [] 192 | input_neighbors = [] 193 | input_pools = [] 194 | input_up_samples = [] 195 | input_labels = [] 196 | ori_batch_labels = batch_labels 197 | 198 | for i in range(cfg.num_layers): 199 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 200 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 201 | sub_labels = batch_labels[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i]] 202 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 203 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 204 | input_points.append(batch_xyz) 205 | input_labels.append(batch_labels) 206 | input_neighbors.append(neighbour_idx) 207 | input_pools.append(pool_i) 208 | input_up_samples.append(up_i) 209 | batch_xyz = sub_points 210 | batch_labels = sub_labels 211 | 212 | input_list = input_points + input_neighbors + input_pools + input_up_samples 213 | input_list += [batch_features, ori_batch_labels, batch_pc_idx, batch_cloud_idx] 214 | input_list += input_labels 215 | 216 | return input_list 217 | 218 | return tf_map 219 | 220 | def init_input_pipeline(self): 221 | print('Initiating input pipelines') 222 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 223 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 224 | gen_function_val, _, _ = self.get_batch_gen('validation') 225 | gen_function_test, _, _ = self.get_batch_gen('test') 226 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 227 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 228 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 229 | 230 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 231 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 232 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 233 | map_func = self.get_tf_mapping2() 234 | 235 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 236 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 237 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 238 | 239 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 240 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 241 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 242 | 243 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 244 | self.flat_inputs = iter.get_next() 245 | self.train_init_op = iter.make_initializer(self.batch_train_data) 246 | self.val_init_op = iter.make_initializer(self.batch_val_data) 247 | self.test_init_op = iter.make_initializer(self.batch_test_data) 248 | 249 | 250 | if __name__ == '__main__': 251 | parser = argparse.ArgumentParser() 252 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 253 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 254 | FLAGS = parser.parse_args() 255 | 256 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 257 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 258 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 259 | Mode = FLAGS.mode 260 | 261 | # shutil.rmtree('__pycache__') if exists('__pycache__') else None 262 | # if Mode == 'train': 263 | # shutil.rmtree('results') if exists('results') else None 264 | # shutil.rmtree('train_log') if exists('train_log') else None 265 | # for f in os.listdir(dirname(abspath(__file__))): 266 | # if f.startswith('log_'): 267 | # os.remove(f) 268 | 269 | dataset = SensatUrban() 270 | dataset.init_input_pipeline() 271 | 272 | if Mode == 'train': 273 | model = Network(dataset, cfg) 274 | model.train(dataset) 275 | elif Mode == 'test': 276 | cfg.saving = False 277 | model = Network(dataset, cfg) 278 | chosen_snapshot = -1 279 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 280 | chosen_folder = logs[-1] 281 | snap_path = join(chosen_folder, 'snapshots') 282 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 283 | chosen_step = np.sort(snap_steps)[-1] 284 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 285 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 286 | tester.test(model, dataset) 287 | shutil.rmtree('train_log') if exists('train_log') else None 288 | 289 | else: 290 | 291 | with tf.Session() as sess: 292 | sess.run(tf.global_variables_initializer()) 293 | sess.run(dataset.train_init_op) 294 | while True: 295 | data_list = sess.run(dataset.flat_inputs) 296 | xyz = data_list[0] 297 | sub_xyz = data_list[1] 298 | label = data_list[21] 299 | Plot.draw_pc_sem_ins(xyz[0, :, :], label[0, :]) 300 | -------------------------------------------------------------------------------- /RandLANet_UNetp.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.compat.v1 as tf 6 | tf.disable_v2_behavior() 7 | import numpy as np 8 | import helper_tf_util 9 | import time 10 | 11 | 12 | def log_out(out_str, f_out): 13 | f_out.write(out_str + '\n') 14 | f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class Network: 19 | def __init__(self, dataset, config): 20 | flat_inputs = dataset.flat_inputs 21 | self.config = config 22 | # Path of the result folder 23 | if self.config.saving: 24 | if self.config.saving_path is None: 25 | self.saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 26 | else: 27 | self.saving_path = self.config.saving_path 28 | makedirs(self.saving_path) if not exists(self.saving_path) else None 29 | 30 | with tf.variable_scope('inputs'): 31 | self.inputs = dict() 32 | num_layers = self.config.num_layers 33 | self.inputs['xyz'] = flat_inputs[:num_layers] 34 | self.inputs['neigh_idx'] = flat_inputs[num_layers: 2 * num_layers] 35 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 36 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 37 | self.inputs['features'] = flat_inputs[4 * num_layers] 38 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 39 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 40 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 41 | 42 | self.labels = self.inputs['labels'] 43 | self.is_training = tf.placeholder(tf.bool, shape=()) 44 | self.training_step = 1 45 | self.training_epoch = 0 46 | self.correct_prediction = 0 47 | self.accuracy = 0 48 | self.mIou_list = [0] 49 | self.loss_type = 'sqrt' # wce, lovas 50 | 51 | self.class_weights = DP.get_class_weights(dataset.name, dataset.num_per_class, self.loss_type) 52 | self.Log_file = open('log_train_' + dataset.name + '.txt', 'a') 53 | 54 | with tf.variable_scope('layers'): 55 | self.logits = self.inference(self.inputs, self.is_training) 56 | 57 | ##################################################################### 58 | # Ignore the invalid point (unlabeled) when calculating the loss # 59 | ##################################################################### 60 | with tf.variable_scope('loss'): 61 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 62 | self.labels = tf.reshape(self.labels, [-1]) 63 | 64 | # Boolean mask of points that should be ignored 65 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 66 | for ign_label in self.config.ignored_label_inds: 67 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 68 | 69 | # Collect logits and labels that are not ignored 70 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 71 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 72 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 73 | 74 | # Reduce label values in the range of logit shape 75 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 76 | inserted_value = tf.zeros((1,), dtype=tf.int32) 77 | for ign_label in self.config.ignored_label_inds: 78 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 79 | valid_labels = tf.gather(reducing_list, valid_labels_init) 80 | 81 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 82 | 83 | with tf.variable_scope('optimizer'): 84 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 85 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 86 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 87 | 88 | with tf.variable_scope('results'): 89 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 90 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 91 | self.prob_logits = tf.nn.softmax(self.logits) 92 | 93 | tf.summary.scalar('learning_rate', self.learning_rate) 94 | tf.summary.scalar('loss', self.loss) 95 | tf.summary.scalar('accuracy', self.accuracy) 96 | 97 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 98 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 99 | c_proto = tf.ConfigProto() 100 | c_proto.gpu_options.allow_growth = True 101 | self.sess = tf.Session(config=c_proto) 102 | self.merged = tf.summary.merge_all() 103 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 104 | self.sess.run(tf.global_variables_initializer()) 105 | 106 | def inference(self, inputs, is_training): 107 | 108 | d_out = self.config.d_out 109 | feature = inputs['features'] 110 | feature = tf.layers.dense(feature, 16, activation=None, name='fc0') 111 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 112 | feature = tf.expand_dims(feature, axis=2) 113 | 114 | jin = [] 115 | feature_list = [] 116 | for i in range(self.config.num_layers): 117 | if i == 0: 118 | feature_list = [] 119 | for j in range(self.config.num_layers - i): 120 | # encoding[(i, j-1)] 121 | f_encoder_i_j = self.dilated_res_block(feature, inputs['xyz'][j], inputs['neigh_idx'][j], d_out[j], 'Encoder_layer_' + str(i) + str(j), is_training) 122 | f_sampled_i_j = self.random_sample(f_encoder_i_j, inputs['sub_idx'][j]) 123 | feature_list.append(f_encoder_i_j) 124 | feature = f_sampled_i_j 125 | jin.append(feature_list) 126 | if i != 0: 127 | feature_list = [] 128 | for j in range(self.config.num_layers - i): 129 | # decoding[(i-1, j); (i-1, j+1)] 130 | f_1 = jin[i-1][j] 131 | f_2 = self.nearest_interpolation(jin[i-1][j+1], inputs['interp_idx'][j]) 132 | feature = tf.concat([f_1, f_2], axis=3) 133 | f_encoder_i_j = self.dilated_res_block(feature, inputs['xyz'][j], inputs['neigh_idx'][j], d_out[j], 'Encoder_layer_' + str(i) + str(j), is_training) 134 | f_sampled_i_j = self.random_sample(f_encoder_i_j, inputs['sub_idx'][j]) 135 | feature_list.append(f_encoder_i_j) 136 | feature = f_sampled_i_j 137 | jin.append(feature_list) 138 | 139 | f_layer_fc1 = helper_tf_util.conv2d(jin[-1][-1], 64, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 140 | f_layer_fc2 = helper_tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 141 | f_layer_drop = helper_tf_util.dropout(f_layer_fc2, keep_prob=0.5, is_training=is_training, scope='dp1') 142 | f_layer_fc3 = helper_tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, is_training, activation_fn=None) 143 | f_out = tf.squeeze(f_layer_fc3, [2]) 144 | return f_out 145 | 146 | def train(self, dataset): 147 | flops = tf.profiler.profile(self.sess.graph, options=tf.profiler.ProfileOptionBuilder.float_operation()) 148 | print("Params: ", str(np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]))) 149 | print("FLOPs: ", flops.total_float_ops) 150 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 151 | self.sess.run(dataset.train_init_op) 152 | while self.training_epoch < self.config.max_epoch: 153 | t_start = time.time() 154 | try: 155 | ops = [self.train_op, 156 | self.extra_update_ops, 157 | self.merged, 158 | self.loss, 159 | self.logits, 160 | self.labels, 161 | self.accuracy] 162 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 163 | self.train_writer.add_summary(summary, self.training_step) 164 | t_end = time.time() 165 | if self.training_step % 50 == 0: 166 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 167 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 168 | self.training_step += 1 169 | 170 | except tf.errors.OutOfRangeError: 171 | 172 | m_iou = self.evaluate(dataset) 173 | if m_iou > np.max(self.mIou_list): 174 | # Save the best model 175 | snapshot_directory = join(self.saving_path, 'snapshots') 176 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 177 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 178 | self.mIou_list.append(m_iou) 179 | log_out('Best m_IoU is: {:5.3f}'.format(max(self.mIou_list)), self.Log_file) 180 | 181 | self.training_epoch += 1 182 | self.sess.run(dataset.train_init_op) 183 | # Update learning rate 184 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, 185 | self.config.lr_decays[self.training_epoch])) 186 | self.sess.run(op) 187 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 188 | 189 | except tf.errors.InvalidArgumentError as e: 190 | 191 | print('Caught a NaN error :') 192 | print(e.error_code) 193 | print(e.message) 194 | print(e.op) 195 | print(e.op.name) 196 | print([t.name for t in e.op.inputs]) 197 | print([t.name for t in e.op.outputs]) 198 | 199 | a = 1 / 0 200 | 201 | print('finished') 202 | self.sess.close() 203 | 204 | def evaluate(self, dataset): 205 | 206 | # Initialise iterator with validation data 207 | self.sess.run(dataset.val_init_op) 208 | 209 | gt_classes = [0 for _ in range(self.config.num_classes)] 210 | positive_classes = [0 for _ in range(self.config.num_classes)] 211 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 212 | val_total_correct = 0 213 | val_total_seen = 0 214 | 215 | for step_id in range(self.config.val_steps): 216 | if step_id % 50 == 0: 217 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 218 | try: 219 | ops = (self.prob_logits, self.labels, self.accuracy) 220 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 221 | pred = np.argmax(stacked_prob, 1) 222 | if not self.config.ignored_label_inds: 223 | pred_valid = pred 224 | labels_valid = labels 225 | else: 226 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 227 | labels_valid = np.delete(labels, invalid_idx) 228 | labels_valid = labels_valid - 1 229 | pred_valid = np.delete(pred, invalid_idx) 230 | 231 | correct = np.sum(pred_valid == labels_valid) 232 | val_total_correct += correct 233 | val_total_seen += len(labels_valid) 234 | 235 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 236 | gt_classes += np.sum(conf_matrix, axis=1) 237 | positive_classes += np.sum(conf_matrix, axis=0) 238 | true_positive_classes += np.diagonal(conf_matrix) 239 | 240 | except tf.errors.OutOfRangeError: 241 | break 242 | 243 | iou_list = [] 244 | for n in range(0, self.config.num_classes, 1): 245 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 246 | iou_list.append(iou) 247 | mean_iou = sum(iou_list) / float(self.config.num_classes) 248 | 249 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 250 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 251 | 252 | mean_iou = 100 * mean_iou 253 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 254 | s = '{:5.2f} | '.format(mean_iou) 255 | for IoU in iou_list: 256 | s += '{:5.2f} '.format(100 * IoU) 257 | log_out('-' * len(s), self.Log_file) 258 | log_out(s, self.Log_file) 259 | log_out('-' * len(s) + '\n', self.Log_file) 260 | return mean_iou 261 | 262 | def get_loss(self, logits, labels, pre_cal_weights): 263 | # calculate the weighted cross entropy according to the inverse frequency 264 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 265 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 266 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 267 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels) 268 | weighted_losses = unweighted_losses * weights 269 | output_loss = tf.reduce_mean(weighted_losses) 270 | return output_loss 271 | 272 | def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training): 273 | f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 274 | f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 275 | f_pc = helper_tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, 276 | activation_fn=None) 277 | shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', 278 | activation_fn=None, bn=True, is_training=is_training) 279 | return tf.nn.leaky_relu(f_pc + shortcut) 280 | 281 | def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training): 282 | d_in = feature.get_shape()[-1].value 283 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) 284 | f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 285 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 286 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 287 | f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training) 288 | 289 | f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 290 | f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx) 291 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 292 | f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training) 293 | return f_pc_agg 294 | 295 | def relative_pos_encoding(self, xyz, neigh_idx): 296 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 297 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 298 | relative_xyz = xyz_tile - neighbor_xyz 299 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 300 | relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1) 301 | return relative_feature 302 | 303 | @staticmethod 304 | def random_sample(feature, pool_idx): 305 | """ 306 | :param feature: [B, N, d] input features matrix 307 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 308 | :return: pool_features = [B, N', d] pooled features matrix 309 | """ 310 | feature = tf.squeeze(feature, axis=2) 311 | num_neigh = tf.shape(pool_idx)[-1] 312 | d = feature.get_shape()[-1] 313 | batch_size = tf.shape(pool_idx)[0] 314 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 315 | pool_features = tf.batch_gather(feature, pool_idx) 316 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 317 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 318 | return pool_features 319 | 320 | @staticmethod 321 | def nearest_interpolation(feature, interp_idx): 322 | """ 323 | :param feature: [B, N, d] input features matrix 324 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 325 | :return: [B, up_num_points, d] interpolated features matrix 326 | """ 327 | feature = tf.squeeze(feature, axis=2) 328 | batch_size = tf.shape(interp_idx)[0] 329 | up_num_points = tf.shape(interp_idx)[1] 330 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 331 | interpolated_features = tf.batch_gather(feature, interp_idx) 332 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 333 | return interpolated_features 334 | 335 | @staticmethod 336 | def gather_neighbour(pc, neighbor_idx): 337 | # gather the coordinates or features of neighboring points 338 | batch_size = tf.shape(pc)[0] 339 | num_points = tf.shape(pc)[1] 340 | d = pc.get_shape()[2].value 341 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 342 | features = tf.batch_gather(pc, index_input) 343 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 344 | return features 345 | 346 | @staticmethod 347 | def att_pooling(feature_set, d_out, name, is_training): 348 | batch_size = tf.shape(feature_set)[0] 349 | num_points = tf.shape(feature_set)[1] 350 | num_neigh = tf.shape(feature_set)[2] 351 | d = feature_set.get_shape()[3].value 352 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 353 | att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc') 354 | att_scores = tf.nn.softmax(att_activation, axis=1) 355 | f_agg = f_reshaped * att_scores 356 | f_agg = tf.reduce_sum(f_agg, axis=1) 357 | f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d]) 358 | f_agg = helper_tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 359 | return f_agg -------------------------------------------------------------------------------- /RandLANet.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.compat.v1 as tf 6 | tf.disable_v2_behavior() 7 | import numpy as np 8 | import helper_tf_util 9 | import time 10 | 11 | 12 | def log_out(out_str, f_out): 13 | f_out.write(out_str + '\n') 14 | f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class Network: 19 | def __init__(self, dataset, config): 20 | flat_inputs = dataset.flat_inputs 21 | self.config = config 22 | # Path of the result folder 23 | if self.config.saving: 24 | if self.config.saving_path is None: 25 | self.saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 26 | else: 27 | self.saving_path = self.config.saving_path 28 | makedirs(self.saving_path) if not exists(self.saving_path) else None 29 | 30 | with tf.variable_scope('inputs'): 31 | self.inputs = dict() 32 | num_layers = self.config.num_layers 33 | self.inputs['xyz'] = flat_inputs[:num_layers] 34 | self.inputs['neigh_idx'] = flat_inputs[num_layers: 2 * num_layers] 35 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 36 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 37 | self.inputs['features'] = flat_inputs[4 * num_layers] 38 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 39 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 40 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 41 | 42 | self.labels = self.inputs['labels'] 43 | self.is_training = tf.placeholder(tf.bool, shape=()) 44 | self.training_step = 1 45 | self.training_epoch = 0 46 | self.correct_prediction = 0 47 | self.accuracy = 0 48 | self.mIou_list = [0] 49 | self.loss_type = 'sqrt' # wce, lovas 50 | 51 | self.class_weights = DP.get_class_weights(dataset.name, dataset.num_per_class, self.loss_type) 52 | self.Log_file = open('log_train_' + dataset.name + '.txt', 'a') 53 | 54 | with tf.variable_scope('layers'): 55 | self.logits = self.inference(self.inputs, self.is_training) 56 | 57 | ##################################################################### 58 | # Ignore the invalid point (unlabeled) when calculating the loss # 59 | ##################################################################### 60 | with tf.variable_scope('loss'): 61 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 62 | self.labels = tf.reshape(self.labels, [-1]) 63 | 64 | # Boolean mask of points that should be ignored 65 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 66 | for ign_label in self.config.ignored_label_inds: 67 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 68 | 69 | # Collect logits and labels that are not ignored 70 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 71 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 72 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 73 | 74 | # Reduce label values in the range of logit shape 75 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 76 | inserted_value = tf.zeros((1,), dtype=tf.int32) 77 | for ign_label in self.config.ignored_label_inds: 78 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 79 | valid_labels = tf.gather(reducing_list, valid_labels_init) 80 | 81 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 82 | 83 | with tf.variable_scope('optimizer'): 84 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 85 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 86 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 87 | 88 | with tf.variable_scope('results'): 89 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 90 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 91 | self.prob_logits = tf.nn.softmax(self.logits) 92 | 93 | tf.summary.scalar('learning_rate', self.learning_rate) 94 | tf.summary.scalar('loss', self.loss) 95 | tf.summary.scalar('accuracy', self.accuracy) 96 | 97 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 98 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 99 | c_proto = tf.ConfigProto() 100 | c_proto.gpu_options.allow_growth = True 101 | self.sess = tf.Session(config=c_proto) 102 | self.merged = tf.summary.merge_all() 103 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 104 | self.sess.run(tf.global_variables_initializer()) 105 | 106 | def inference(self, inputs, is_training): 107 | 108 | d_out = self.config.d_out 109 | feature = inputs['features'] 110 | feature = tf.layers.dense(feature, 8, activation=None, name='fc0') 111 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 112 | feature = tf.expand_dims(feature, axis=2) 113 | 114 | # ###########################Encoder############################ 115 | f_encoder_list = [] 116 | for i in range(self.config.num_layers): 117 | f_encoder_i = self.dilated_res_block(feature, inputs['xyz'][i], inputs['neigh_idx'][i], d_out[i], 'Encoder_layer_' + str(i), is_training) 118 | f_sampled_i = self.random_sample(f_encoder_i, inputs['sub_idx'][i]) 119 | feature = f_sampled_i 120 | if i == 0: 121 | f_encoder_list.append(f_encoder_i) 122 | f_encoder_list.append(f_sampled_i) 123 | # ###########################Encoder############################ 124 | 125 | feature = helper_tf_util.conv2d(f_encoder_list[-1], f_encoder_list[-1].get_shape()[3].value, [1, 1], 'decoder_0', [1, 1], 'VALID', True, is_training) 126 | 127 | # ###########################Decoder############################ 128 | f_decoder_list = [] 129 | for j in range(self.config.num_layers): 130 | f_interp_i = self.nearest_interpolation(feature, inputs['interp_idx'][-j - 1]) 131 | f_decoder_i = helper_tf_util.conv2d_transpose(tf.concat([f_encoder_list[-j - 2], f_interp_i], axis=3), 132 | f_encoder_list[-j - 2].get_shape()[-1].value, [1, 1], 133 | 'Decoder_layer_' + str(j), [1, 1], 'VALID', bn=True, 134 | is_training=is_training) 135 | feature = f_decoder_i 136 | f_decoder_list.append(f_decoder_i) 137 | # ###########################Decoder############################ 138 | 139 | f_layer_fc1 = helper_tf_util.conv2d(f_decoder_list[-1], 64, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 140 | f_layer_fc2 = helper_tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 141 | f_layer_drop = helper_tf_util.dropout(f_layer_fc2, keep_prob=0.5, is_training=is_training, scope='dp1') 142 | f_layer_fc3 = helper_tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, 143 | is_training, activation_fn=None) 144 | f_out = tf.squeeze(f_layer_fc3, [2]) 145 | return f_out 146 | 147 | def train(self, dataset): 148 | flops = tf.profiler.profile(self.sess.graph, options=tf.profiler.ProfileOptionBuilder.float_operation()) 149 | print("Params: ", str(np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]))) 150 | print("FLOPs: ", flops.total_float_ops) 151 | 152 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 153 | self.sess.run(dataset.train_init_op) 154 | while self.training_epoch < self.config.max_epoch: 155 | t_start = time.time() 156 | try: 157 | ops = [self.train_op, 158 | self.extra_update_ops, 159 | self.merged, 160 | self.loss, 161 | self.logits, 162 | self.labels, 163 | self.accuracy] 164 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 165 | self.train_writer.add_summary(summary, self.training_step) 166 | t_end = time.time() 167 | if self.training_step % 50 == 0: 168 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 169 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 170 | self.training_step += 1 171 | 172 | except tf.errors.OutOfRangeError: 173 | 174 | m_iou = self.evaluate(dataset) 175 | if m_iou > np.max(self.mIou_list): 176 | # Save the best model 177 | snapshot_directory = join(self.saving_path, 'snapshots') 178 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 179 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 180 | self.mIou_list.append(m_iou) 181 | log_out('Best m_IoU is: {:5.3f}'.format(max(self.mIou_list)), self.Log_file) 182 | 183 | self.training_epoch += 1 184 | self.sess.run(dataset.train_init_op) 185 | # Update learning rate 186 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, 187 | self.config.lr_decays[self.training_epoch])) 188 | self.sess.run(op) 189 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 190 | 191 | except tf.errors.InvalidArgumentError as e: 192 | 193 | print('Caught a NaN error :') 194 | print(e.error_code) 195 | print(e.message) 196 | print(e.op) 197 | print(e.op.name) 198 | print([t.name for t in e.op.inputs]) 199 | print([t.name for t in e.op.outputs]) 200 | 201 | a = 1 / 0 202 | 203 | print('finished') 204 | self.sess.close() 205 | 206 | def evaluate(self, dataset): 207 | 208 | # Initialise iterator with validation data 209 | self.sess.run(dataset.val_init_op) 210 | 211 | gt_classes = [0 for _ in range(self.config.num_classes)] 212 | positive_classes = [0 for _ in range(self.config.num_classes)] 213 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 214 | val_total_correct = 0 215 | val_total_seen = 0 216 | 217 | for step_id in range(self.config.val_steps): 218 | if step_id % 50 == 0: 219 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 220 | try: 221 | ops = (self.prob_logits, self.labels, self.accuracy) 222 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 223 | pred = np.argmax(stacked_prob, 1) 224 | if not self.config.ignored_label_inds: 225 | pred_valid = pred 226 | labels_valid = labels 227 | else: 228 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 229 | labels_valid = np.delete(labels, invalid_idx) 230 | labels_valid = labels_valid - 1 231 | pred_valid = np.delete(pred, invalid_idx) 232 | 233 | correct = np.sum(pred_valid == labels_valid) 234 | val_total_correct += correct 235 | val_total_seen += len(labels_valid) 236 | 237 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 238 | gt_classes += np.sum(conf_matrix, axis=1) 239 | positive_classes += np.sum(conf_matrix, axis=0) 240 | true_positive_classes += np.diagonal(conf_matrix) 241 | 242 | except tf.errors.OutOfRangeError: 243 | break 244 | 245 | iou_list = [] 246 | for n in range(0, self.config.num_classes, 1): 247 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 248 | iou_list.append(iou) 249 | mean_iou = sum(iou_list) / float(self.config.num_classes) 250 | 251 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 252 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 253 | 254 | mean_iou = 100 * mean_iou 255 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 256 | s = '{:5.2f} | '.format(mean_iou) 257 | for IoU in iou_list: 258 | s += '{:5.2f} '.format(100 * IoU) 259 | log_out('-' * len(s), self.Log_file) 260 | log_out(s, self.Log_file) 261 | log_out('-' * len(s) + '\n', self.Log_file) 262 | return mean_iou 263 | 264 | def get_loss(self, logits, labels, pre_cal_weights): 265 | # calculate the weighted cross entropy according to the inverse frequency 266 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 267 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 268 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 269 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels) 270 | weighted_losses = unweighted_losses * weights 271 | output_loss = tf.reduce_mean(weighted_losses) 272 | return output_loss 273 | 274 | def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training): 275 | f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 276 | f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 277 | f_pc = helper_tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, 278 | activation_fn=None) 279 | shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', 280 | activation_fn=None, bn=True, is_training=is_training) 281 | return tf.nn.leaky_relu(f_pc + shortcut) 282 | 283 | def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training): 284 | d_in = feature.get_shape()[-1].value 285 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) 286 | f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 287 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 288 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 289 | f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training) 290 | 291 | f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 292 | f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx) 293 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 294 | f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training) 295 | return f_pc_agg 296 | 297 | def relative_pos_encoding(self, xyz, neigh_idx): 298 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 299 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 300 | relative_xyz = xyz_tile - neighbor_xyz 301 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 302 | relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1) 303 | return relative_feature 304 | 305 | @staticmethod 306 | def random_sample(feature, pool_idx): 307 | """ 308 | :param feature: [B, N, d] input features matrix 309 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 310 | :return: pool_features = [B, N', d] pooled features matrix 311 | """ 312 | feature = tf.squeeze(feature, axis=2) 313 | num_neigh = tf.shape(pool_idx)[-1] 314 | d = feature.get_shape()[-1] 315 | batch_size = tf.shape(pool_idx)[0] 316 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 317 | pool_features = tf.batch_gather(feature, pool_idx) 318 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 319 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 320 | return pool_features 321 | 322 | @staticmethod 323 | def nearest_interpolation(feature, interp_idx): 324 | """ 325 | :param feature: [B, N, d] input features matrix 326 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 327 | :return: [B, up_num_points, d] interpolated features matrix 328 | """ 329 | feature = tf.squeeze(feature, axis=2) 330 | batch_size = tf.shape(interp_idx)[0] 331 | up_num_points = tf.shape(interp_idx)[1] 332 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 333 | interpolated_features = tf.batch_gather(feature, interp_idx) 334 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 335 | return interpolated_features 336 | 337 | @staticmethod 338 | def gather_neighbour(pc, neighbor_idx): 339 | # gather the coordinates or features of neighboring points 340 | batch_size = tf.shape(pc)[0] 341 | num_points = tf.shape(pc)[1] 342 | d = pc.get_shape()[2].value 343 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 344 | features = tf.batch_gather(pc, index_input) 345 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 346 | return features 347 | 348 | @staticmethod 349 | def att_pooling(feature_set, d_out, name, is_training): 350 | batch_size = tf.shape(feature_set)[0] 351 | num_points = tf.shape(feature_set)[1] 352 | num_neigh = tf.shape(feature_set)[2] 353 | d = feature_set.get_shape()[3].value 354 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 355 | att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc') 356 | att_scores = tf.nn.softmax(att_activation, axis=1) 357 | f_agg = f_reshaped * att_scores 358 | f_agg = tf.reduce_sum(f_agg, axis=1) 359 | f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d]) 360 | f_agg = helper_tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 361 | return f_agg --------------------------------------------------------------------------------