├── .gitignore ├── README.md ├── cpp_wrappers ├── compile_wrappers.sh ├── cpp_neighbors │ ├── build.bat │ ├── build │ │ ├── lib.linux-x86_64-3.8 │ │ │ └── radius_neighbors.cpython-38-x86_64-linux-gnu.so │ │ ├── temp.linux-x86_64-3.8 │ │ │ ├── cpp_wrappers │ │ │ │ └── cpp_utils │ │ │ │ │ └── cloud │ │ │ │ │ └── cloud.o │ │ │ ├── neighbors │ │ │ │ └── neighbors.o │ │ │ └── wrapper.o │ │ └── temp.linux-x86_64-3.9 │ │ │ ├── cpp_wrappers │ │ │ └── cpp_utils │ │ │ │ └── cloud │ │ │ │ └── cloud.o │ │ │ ├── neighbors │ │ │ └── neighbors.o │ │ │ └── wrapper.o │ ├── neighbors │ │ ├── neighbors.cpp │ │ └── neighbors.h │ ├── radius_neighbors.cpython-38-x86_64-linux-gnu.so │ ├── radius_neighbors.cpython-39-x86_64-linux-gnu.so │ ├── setup.py │ └── wrapper.cpp ├── cpp_subsampling │ ├── build.bat │ ├── build │ │ ├── lib.linux-x86_64-3.8 │ │ │ └── grid_subsampling.cpython-38-x86_64-linux-gnu.so │ │ ├── temp.linux-x86_64-3.8 │ │ │ ├── cpp_wrappers │ │ │ │ └── cpp_utils │ │ │ │ │ └── cloud │ │ │ │ │ └── cloud.o │ │ │ ├── grid_subsampling │ │ │ │ └── grid_subsampling.o │ │ │ └── wrapper.o │ │ └── temp.linux-x86_64-3.9 │ │ │ ├── cpp_wrappers │ │ │ └── cpp_utils │ │ │ │ └── cloud │ │ │ │ └── cloud.o │ │ │ ├── grid_subsampling │ │ │ └── grid_subsampling.o │ │ │ └── wrapper.o │ ├── grid_subsampling.cpython-38-x86_64-linux-gnu.so │ ├── grid_subsampling.cpython-39-x86_64-linux-gnu.so │ ├── grid_subsampling │ │ ├── grid_subsampling.cpp │ │ └── grid_subsampling.h │ ├── setup.py │ └── wrapper.cpp └── cpp_utils │ ├── cloud │ ├── cloud.cpp │ └── cloud.h │ └── nanoflann │ └── nanoflann.hpp ├── data_utils ├── kpconv_loader.py ├── modelnet40_loader.py └── shapenet_loader.py ├── datasets └── ModelNet40.py ├── misc ├── layers.py ├── ops.py ├── pointconv_utils.py └── utils.py ├── networks ├── cls │ ├── blocks.py │ ├── data.txt │ ├── dgcnn.py │ ├── kernel_points.py │ ├── kernels │ │ └── dispositions │ │ │ └── k_015_center_3D.ply │ ├── kpconv.py │ ├── pointcnn.py │ ├── pointconv.py │ ├── pointnet.py │ └── pointnet2.py └── seg │ ├── dgcnn_partseg.py │ ├── pointcnn_partseg.py │ ├── pointconv_partseg.py │ ├── pointnet2_partseg.py │ └── pointnet_partseg.py ├── run_cls.sh ├── run_partseg.sh ├── train_cls.py └── train_partseg.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | data 3 | data.zip 4 | *.output 5 | *.o 6 | *.so 7 | *.tar -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 计图点云库 2 | 3 | ## 已经实现的模型 4 | 5 | | Model | Classification | Segmentation | 6 | | ----------- | -------------- | ------------ | 7 | | PointNet | √ | √ | 8 | | PointNet ++ | √ | √ | 9 | | PointCNN | √ | √ | 10 | | DGCNN | √ | √ | 11 | | PointConv | √ | √ | 12 | | KPConv | √ | | 13 | 14 | ## 使用方法 15 | 16 | ### 安装依赖 17 | 18 | ```bash 19 | sudo apt install python3.7-dev libomp-dev 20 | sudo python3.7 -m pip install git+https://github.com/Jittor/jittor.git 21 | python3.7 -m pip install sklearn lmdb msgpack_numpy 22 | ``` 23 | ### 安装点云库 24 | 25 | ```bash 26 | git clone https://github.com/Jittor/PointCloudLib.git # 将库下载的本地 27 | # 您需要将 ModelNet40 和 ShapeNet 数据集下载到 data_util/data/ 里面 28 | ModelNet40 数据集链接 ： https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip 29 | ShapeNet 数据集链接 ： https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip 30 | 31 | sh run_cls.sh # 点云分类的训练和测试（以PointNet为例） 32 | sh run_seg.sh # 点云分割的训练和测试（以PointNet为例） 33 | 34 | # 对于kpconv，需要额外执行脚本再开始训练 35 | cd cpp_wrappers 36 | bash compile_wrappers.sh 37 | cd .. 38 | python train_cls.py --model kpconv # kpconv训练 39 | python train_cls.py --model kpconv --eval # 修改train_cls.py中的chkp_path指定模型进行测试 40 | ``` 41 | 42 | ## 所依赖的库 43 | 44 | ```bash 45 | Python 3.7 46 | Jittor 47 | Numpy 48 | sklearn 49 | lmdb 50 | msgpack_numpy 51 | ... 52 | ``` 53 | 54 | 55 | 56 | ## 实验结果 57 | 58 | ### 分类训练效果测试 59 | 60 | | Model | Input | overall accuracy | 61 | | ----------- | ----------------- | ---------------- | 62 | | PointNet | 1024 xyz | 87.2 | 63 | | PointNet ++ | 4096 xyz + normal | 92.3 | 64 | | PointCNN | 1024 xyz | 92.6 | 65 | | DGCNN | 1024 xyz | 92.9 | 66 | | PointConv | 1024 xyz + normal | 92.4 | 67 | | KPConv | xyz + neighbors + pools + lengths + features | 92.5 | 68 | 69 | ### 分类训练时间测试 70 | 71 | | Model | Speed up ratio (Compare with Pytorch) | 72 | | :---------- | ------------------------------------- | 73 | | PointNet | 1.22 | 74 | | PointNet ++ | 2.72 | 75 | | PointCNN | 2.41 | 76 | | DGCNN | 1.22 | 77 | | PointConv | | 78 | | KPConv | | 79 | 80 | ### 分割训练效果测试 81 | 82 | | Model | Input | pIoU | 83 | | ----------- | ----------------------------- | ---- | 84 | | PointNet | 2048 xyz + cls label | 83.5 | 85 | | PointNet ++ | 2048 xyz + cls label + normal | 85.0 | 86 | | PointCNN | 2048 xyz + normal | 86.0 | 87 | | DGCNN | 2048 xyz + cls label | 85.1 | 88 | | PointConv | 2048 xyz | 85.4 | 89 | 90 | ### 分割训练时间测试 91 | 92 | | Model | Speed up ratio (Compare with Pytorch) | 93 | | ----------- | ------------------------------------- | 94 | | PointNet | 1.06 | 95 | | PointNet ++ | 1.85 | 96 | | PointCNN | None (No pytorch implementation) | 97 | | DGCNN | 1.05 | 98 | | PointConv | None (No pytorch implementation) | 99 | 100 | ## 目录结构 101 | 102 | ``` 103 | . 104 | ├── data_utils # 数据相关工具 105 | │ ├── data # 数据存放路径 106 | │ ├── modelnet40_loader.py 107 | │   └── shapenet_loader.py 108 | ├── misc 109 | │   ├── layers.py 110 | │   ├── ops.py 111 | │   ├── pointconv_utils.py 112 | │   └── utils.py 113 | ├── networks 114 | │   ├── cls 115 | │   │   ├── dgcnn.py 116 | │   │   ├── pointcnn.py 117 | │   │   ├── pointconv.py 118 | │   │   ├── pointnet2.py 119 | │   │   └── pointnet.py 120 | │   └── seg 121 | │   ├── dgcnn_partseg.py 122 | │   ├── pointcnn_partseg.py 123 | │   ├── pointconv_partseg.py 124 | │   ├── pointnet2_partseg.py 125 | │   └── pointnet_partseg.py 126 | 127 | ├── README.md 128 | ├── run_cls.sh 129 | ├── run_partseg.sh 130 | ├── train_cls.py 131 | └── train_partseg.py 132 | ``` 133 | 其他点云开源工作： 134 | PCT: https://github.com/MenghaoGuo/PCT 135 | 136 | 非常欢迎您使用计图的点云库进行相关的研究，如在使用中有问题，欢迎提交 issues。 137 | ## Reference code : 138 | https://github.com/AnTao97/dgcnn.pytorch 139 | -------------------------------------------------------------------------------- /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 cpp neighbors 9 | cd cpp_neighbors 10 | python3 setup.py build_ext --inplace 11 | cd .. -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | py setup.py build_ext --inplace 3 | 4 | 5 | pause -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/lib.linux-x86_64-3.8/radius_neighbors.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/lib.linux-x86_64-3.8/radius_neighbors.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/cpp_wrappers/cpp_utils/cloud/cloud.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/cpp_wrappers/cpp_utils/cloud/cloud.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/neighbors/neighbors.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/neighbors/neighbors.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/wrapper.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.8/wrapper.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/cpp_wrappers/cpp_utils/cloud/cloud.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/cpp_wrappers/cpp_utils/cloud/cloud.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/neighbors/neighbors.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/neighbors/neighbors.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/wrapper.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/build/temp.linux-x86_64-3.9/wrapper.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "neighbors.h" 3 | 4 | 5 | void brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose) 6 | { 7 | 8 | // Initialize variables 9 | // ****************** 10 | 11 | // square radius 12 | float r2 = radius * radius; 13 | 14 | // indices 15 | int i0 = 0; 16 | 17 | // Counting vector 18 | int max_count = 0; 19 | vector> tmp(queries.size()); 20 | 21 | // Search neigbors indices 22 | // *********************** 23 | 24 | for (auto& p0 : queries) 25 | { 26 | int i = 0; 27 | for (auto& p : supports) 28 | { 29 | if ((p0 - p).sq_norm() < r2) 30 | { 31 | tmp[i0].push_back(i); 32 | if (tmp[i0].size() > max_count) 33 | max_count = tmp[i0].size(); 34 | } 35 | i++; 36 | } 37 | i0++; 38 | } 39 | 40 | // Reserve the memory 41 | neighbors_indices.resize(queries.size() * max_count); 42 | i0 = 0; 43 | for (auto& inds : tmp) 44 | { 45 | for (int j = 0; j < max_count; j++) 46 | { 47 | if (j < inds.size()) 48 | neighbors_indices[i0 * max_count + j] = inds[j]; 49 | else 50 | neighbors_indices[i0 * max_count + j] = -1; 51 | } 52 | i0++; 53 | } 54 | 55 | return; 56 | } 57 | 58 | void ordered_neighbors(vector& queries, 59 | vector& supports, 60 | vector& neighbors_indices, 61 | float radius) 62 | { 63 | 64 | // Initialize variables 65 | // ****************** 66 | 67 | // square radius 68 | float r2 = radius * radius; 69 | 70 | // indices 71 | int i0 = 0; 72 | 73 | // Counting vector 74 | int max_count = 0; 75 | float d2; 76 | vector> tmp(queries.size()); 77 | vector> dists(queries.size()); 78 | 79 | // Search neigbors indices 80 | // *********************** 81 | 82 | for (auto& p0 : queries) 83 | { 84 | int i = 0; 85 | for (auto& p : supports) 86 | { 87 | d2 = (p0 - p).sq_norm(); 88 | if (d2 < r2) 89 | { 90 | // Find order of the new point 91 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 92 | int index = std::distance(dists[i0].begin(), it); 93 | 94 | // Insert element 95 | dists[i0].insert(it, d2); 96 | tmp[i0].insert(tmp[i0].begin() + index, i); 97 | 98 | // Update max count 99 | if (tmp[i0].size() > max_count) 100 | max_count = tmp[i0].size(); 101 | } 102 | i++; 103 | } 104 | i0++; 105 | } 106 | 107 | // Reserve the memory 108 | neighbors_indices.resize(queries.size() * max_count); 109 | i0 = 0; 110 | for (auto& inds : tmp) 111 | { 112 | for (int j = 0; j < max_count; j++) 113 | { 114 | if (j < inds.size()) 115 | neighbors_indices[i0 * max_count + j] = inds[j]; 116 | else 117 | neighbors_indices[i0 * max_count + j] = -1; 118 | } 119 | i0++; 120 | } 121 | 122 | return; 123 | } 124 | 125 | void batch_ordered_neighbors(vector& queries, 126 | vector& supports, 127 | vector& q_batches, 128 | vector& s_batches, 129 | vector& neighbors_indices, 130 | float radius) 131 | { 132 | 133 | // Initialize variables 134 | // ****************** 135 | 136 | // square radius 137 | float r2 = radius * radius; 138 | 139 | // indices 140 | int i0 = 0; 141 | 142 | // Counting vector 143 | int max_count = 0; 144 | float d2; 145 | vector> tmp(queries.size()); 146 | vector> dists(queries.size()); 147 | 148 | // batch index 149 | int b = 0; 150 | int sum_qb = 0; 151 | int sum_sb = 0; 152 | 153 | 154 | // Search neigbors indices 155 | // *********************** 156 | 157 | for (auto& p0 : queries) 158 | { 159 | // Check if we changed batch 160 | if (i0 == sum_qb + q_batches[b]) 161 | { 162 | sum_qb += q_batches[b]; 163 | sum_sb += s_batches[b]; 164 | b++; 165 | } 166 | 167 | // Loop only over the supports of current batch 168 | vector::iterator p_it; 169 | int i = 0; 170 | for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ ) 171 | { 172 | d2 = (p0 - *p_it).sq_norm(); 173 | if (d2 < r2) 174 | { 175 | // Find order of the new point 176 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 177 | int index = std::distance(dists[i0].begin(), it); 178 | 179 | // Insert element 180 | dists[i0].insert(it, d2); 181 | tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i); 182 | 183 | // Update max count 184 | if (tmp[i0].size() > max_count) 185 | max_count = tmp[i0].size(); 186 | } 187 | i++; 188 | } 189 | i0++; 190 | } 191 | 192 | // Reserve the memory 193 | neighbors_indices.resize(queries.size() * max_count); 194 | i0 = 0; 195 | for (auto& inds : tmp) 196 | { 197 | for (int j = 0; j < max_count; j++) 198 | { 199 | if (j < inds.size()) 200 | neighbors_indices[i0 * max_count + j] = inds[j]; 201 | else 202 | neighbors_indices[i0 * max_count + j] = supports.size(); 203 | } 204 | i0++; 205 | } 206 | 207 | return; 208 | } 209 | 210 | 211 | void batch_nanoflann_neighbors(vector& queries, 212 | vector& supports, 213 | vector& q_batches, 214 | vector& s_batches, 215 | vector& neighbors_indices, 216 | float radius) 217 | { 218 | 219 | // Initialize variables 220 | // ****************** 221 | 222 | // indices 223 | int i0 = 0; 224 | 225 | // Square radius 226 | float r2 = radius * radius; 227 | 228 | // Counting vector 229 | int max_count = 0; 230 | float d2; 231 | vector>> all_inds_dists(queries.size()); 232 | 233 | // batch index 234 | int b = 0; 235 | int sum_qb = 0; 236 | int sum_sb = 0; 237 | 238 | // Nanoflann related variables 239 | // *************************** 240 | 241 | // CLoud variable 242 | PointCloud current_cloud; 243 | 244 | // Tree parameters 245 | nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); 246 | 247 | // KDTree type definition 248 | typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , 249 | PointCloud, 250 | 3 > my_kd_tree_t; 251 | 252 | // Pointer to trees 253 | my_kd_tree_t* index; 254 | 255 | // Build KDTree for the first batch element 256 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 257 | index = new my_kd_tree_t(3, current_cloud, tree_params); 258 | index->buildIndex(); 259 | 260 | 261 | // Search neigbors indices 262 | // *********************** 263 | 264 | // Search params 265 | nanoflann::SearchParams search_params; 266 | search_params.sorted = true; 267 | 268 | for (auto& p0 : queries) 269 | { 270 | 271 | // Check if we changed batch 272 | if (i0 == sum_qb + q_batches[b]) 273 | { 274 | sum_qb += q_batches[b]; 275 | sum_sb += s_batches[b]; 276 | b++; 277 | 278 | // Change the points 279 | current_cloud.pts.clear(); 280 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 281 | 282 | // Build KDTree of the current element of the batch 283 | delete index; 284 | index = new my_kd_tree_t(3, current_cloud, tree_params); 285 | index->buildIndex(); 286 | } 287 | 288 | // Initial guess of neighbors size 289 | all_inds_dists[i0].reserve(max_count); 290 | 291 | // Find neighbors 292 | float query_pt[3] = { p0.x, p0.y, p0.z}; 293 | size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params); 294 | 295 | // Update max count 296 | if (nMatches > max_count) 297 | max_count = nMatches; 298 | 299 | // Increment query idx 300 | i0++; 301 | } 302 | 303 | // Reserve the memory 304 | neighbors_indices.resize(queries.size() * max_count); 305 | i0 = 0; 306 | sum_sb = 0; 307 | sum_qb = 0; 308 | b = 0; 309 | for (auto& inds_dists : all_inds_dists) 310 | { 311 | // Check if we changed batch 312 | if (i0 == sum_qb + q_batches[b]) 313 | { 314 | sum_qb += q_batches[b]; 315 | sum_sb += s_batches[b]; 316 | b++; 317 | } 318 | 319 | for (int j = 0; j < max_count; j++) 320 | { 321 | if (j < inds_dists.size()) 322 | neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb; 323 | else 324 | neighbors_indices[i0 * max_count + j] = supports.size(); 325 | } 326 | i0++; 327 | } 328 | 329 | delete index; 330 | 331 | return; 332 | } 333 | 334 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/neighbors/neighbors.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | #include "../../cpp_utils/nanoflann/nanoflann.hpp" 5 | 6 | #include 7 | #include 8 | 9 | using namespace std; 10 | 11 | 12 | void ordered_neighbors(vector& queries, 13 | vector& supports, 14 | vector& neighbors_indices, 15 | float radius); 16 | 17 | void batch_ordered_neighbors(vector& queries, 18 | vector& supports, 19 | vector& q_batches, 20 | vector& s_batches, 21 | vector& neighbors_indices, 22 | float radius); 23 | 24 | void batch_nanoflann_neighbors(vector& queries, 25 | vector& supports, 26 | vector& q_batches, 27 | vector& s_batches, 28 | vector& neighbors_indices, 29 | float radius); 30 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/radius_neighbors.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/radius_neighbors.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/radius_neighbors.cpython-39-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_neighbors/radius_neighbors.cpython-39-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/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 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 11 | "neighbors/neighbors.cpp", 12 | "wrapper.cpp"] 13 | 14 | module = Extension(name="radius_neighbors", 15 | sources=SOURCES, 16 | extra_compile_args=['-std=c++11', 17 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 18 | 19 | 20 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "neighbors/neighbors.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides two methods to compute radius neighbors from pointclouds or batch of pointclouds"; 12 | 13 | static char batch_query_docstring[] = "Method to get radius neighbors in a batch of stacked pointclouds"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *batch_neighbors(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "batch_query", (PyCFunction)batch_neighbors, METH_VARARGS | METH_KEYWORDS, batch_query_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 | "radius_neighbors", // 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_radius_neighbors(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Definition of the batch_subsample method 56 | // ********************************** 57 | 58 | static PyObject* batch_neighbors(PyObject* self, PyObject* args, PyObject* keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject* queries_obj = NULL; 66 | PyObject* supports_obj = NULL; 67 | PyObject* q_batches_obj = NULL; 68 | PyObject* s_batches_obj = NULL; 69 | 70 | // Keywords containers 71 | static char* kwlist[] = { "queries", "supports", "q_batches", "s_batches", "radius", NULL }; 72 | float radius = 0.1; 73 | 74 | // Parse the input 75 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "OOOO|$f", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &radius)) 76 | { 77 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 78 | return NULL; 79 | } 80 | 81 | 82 | // Interpret the input objects as numpy arrays. 83 | PyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY); 84 | PyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY); 85 | PyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY); 86 | PyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY); 87 | 88 | // Verify data was load correctly. 89 | if (queries_array == NULL) 90 | { 91 | Py_XDECREF(queries_array); 92 | Py_XDECREF(supports_array); 93 | Py_XDECREF(q_batches_array); 94 | Py_XDECREF(s_batches_array); 95 | PyErr_SetString(PyExc_RuntimeError, "Error converting query points to numpy arrays of type float32"); 96 | return NULL; 97 | } 98 | if (supports_array == NULL) 99 | { 100 | Py_XDECREF(queries_array); 101 | Py_XDECREF(supports_array); 102 | Py_XDECREF(q_batches_array); 103 | Py_XDECREF(s_batches_array); 104 | PyErr_SetString(PyExc_RuntimeError, "Error converting support points to numpy arrays of type float32"); 105 | return NULL; 106 | } 107 | if (q_batches_array == NULL) 108 | { 109 | Py_XDECREF(queries_array); 110 | Py_XDECREF(supports_array); 111 | Py_XDECREF(q_batches_array); 112 | Py_XDECREF(s_batches_array); 113 | PyErr_SetString(PyExc_RuntimeError, "Error converting query batches to numpy arrays of type int32"); 114 | return NULL; 115 | } 116 | if (s_batches_array == NULL) 117 | { 118 | Py_XDECREF(queries_array); 119 | Py_XDECREF(supports_array); 120 | Py_XDECREF(q_batches_array); 121 | Py_XDECREF(s_batches_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting support batches to numpy arrays of type int32"); 123 | return NULL; 124 | } 125 | 126 | // Check that the input array respect the dims 127 | if ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3) 128 | { 129 | Py_XDECREF(queries_array); 130 | Py_XDECREF(supports_array); 131 | Py_XDECREF(q_batches_array); 132 | Py_XDECREF(s_batches_array); 133 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : query.shape is not (N, 3)"); 134 | return NULL; 135 | } 136 | if ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3) 137 | { 138 | Py_XDECREF(queries_array); 139 | Py_XDECREF(supports_array); 140 | Py_XDECREF(q_batches_array); 141 | Py_XDECREF(s_batches_array); 142 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : support.shape is not (N, 3)"); 143 | return NULL; 144 | } 145 | if ((int)PyArray_NDIM(q_batches_array) > 1) 146 | { 147 | Py_XDECREF(queries_array); 148 | Py_XDECREF(supports_array); 149 | Py_XDECREF(q_batches_array); 150 | Py_XDECREF(s_batches_array); 151 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : queries_batches.shape is not (B,) "); 152 | return NULL; 153 | } 154 | if ((int)PyArray_NDIM(s_batches_array) > 1) 155 | { 156 | Py_XDECREF(queries_array); 157 | Py_XDECREF(supports_array); 158 | Py_XDECREF(q_batches_array); 159 | Py_XDECREF(s_batches_array); 160 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : supports_batches.shape is not (B,) "); 161 | return NULL; 162 | } 163 | if ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0)) 164 | { 165 | Py_XDECREF(queries_array); 166 | Py_XDECREF(supports_array); 167 | Py_XDECREF(q_batches_array); 168 | Py_XDECREF(s_batches_array); 169 | PyErr_SetString(PyExc_RuntimeError, "Wrong number of batch elements: different for queries and supports "); 170 | return NULL; 171 | } 172 | 173 | // Number of points 174 | int Nq = (int)PyArray_DIM(queries_array, 0); 175 | int Ns= (int)PyArray_DIM(supports_array, 0); 176 | 177 | // Number of batches 178 | int Nb = (int)PyArray_DIM(q_batches_array, 0); 179 | 180 | // Call the C++ function 181 | // ********************* 182 | 183 | // Convert PyArray to Cloud C++ class 184 | vector queries; 185 | vector supports; 186 | vector q_batches; 187 | vector s_batches; 188 | queries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq); 189 | supports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns); 190 | q_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb); 191 | s_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb); 192 | 193 | // Create result containers 194 | vector neighbors_indices; 195 | 196 | // Compute results 197 | //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 198 | batch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 199 | 200 | // Check result 201 | if (neighbors_indices.size() < 1) 202 | { 203 | PyErr_SetString(PyExc_RuntimeError, "Error"); 204 | return NULL; 205 | } 206 | 207 | // Manage outputs 208 | // ************** 209 | 210 | // Maximal number of neighbors 211 | int max_neighbors = neighbors_indices.size() / Nq; 212 | 213 | // Dimension of output containers 214 | npy_intp* neighbors_dims = new npy_intp[2]; 215 | neighbors_dims[0] = Nq; 216 | neighbors_dims[1] = max_neighbors; 217 | 218 | // Create output array 219 | PyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_INT); 220 | PyObject* ret = NULL; 221 | 222 | // Fill output array with values 223 | size_t size_in_bytes = Nq * max_neighbors * sizeof(int); 224 | memcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes); 225 | 226 | // Merge results 227 | ret = Py_BuildValue("N", res_obj); 228 | 229 | // Clean up 230 | // ******** 231 | 232 | Py_XDECREF(queries_array); 233 | Py_XDECREF(supports_array); 234 | Py_XDECREF(q_batches_array); 235 | Py_XDECREF(s_batches_array); 236 | 237 | return ret; 238 | } 239 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | py setup.py build_ext --inplace 3 | 4 | 5 | pause -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/lib.linux-x86_64-3.8/grid_subsampling.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/lib.linux-x86_64-3.8/grid_subsampling.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/cpp_wrappers/cpp_utils/cloud/cloud.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/cpp_wrappers/cpp_utils/cloud/cloud.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/grid_subsampling/grid_subsampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/grid_subsampling/grid_subsampling.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/wrapper.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.8/wrapper.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/cpp_wrappers/cpp_utils/cloud/cloud.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/cpp_wrappers/cpp_utils/cloud/cloud.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/grid_subsampling/grid_subsampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/grid_subsampling/grid_subsampling.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/wrapper.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/build/temp.linux-x86_64-3.9/wrapper.o -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/grid_subsampling.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling.cpython-39-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/cpp_wrappers/cpp_subsampling/grid_subsampling.cpython-39-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /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 | // Initialize 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 | // Initialize 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 | 108 | 109 | void batch_grid_subsampling(vector& original_points, 110 | vector& subsampled_points, 111 | vector& original_features, 112 | vector& subsampled_features, 113 | vector& original_classes, 114 | vector& subsampled_classes, 115 | vector& original_batches, 116 | vector& subsampled_batches, 117 | float sampleDl, 118 | int max_p) 119 | { 120 | // Initialize variables 121 | // ****************** 122 | 123 | int b = 0; 124 | int sum_b = 0; 125 | 126 | // Number of points in the cloud 127 | size_t N = original_points.size(); 128 | 129 | // Dimension of the features 130 | size_t fdim = original_features.size() / N; 131 | size_t ldim = original_classes.size() / N; 132 | 133 | // Handle max_p = 0 134 | if (max_p < 1) 135 | max_p = N; 136 | 137 | // Loop over batches 138 | // ***************** 139 | 140 | for (b = 0; b < original_batches.size(); b++) 141 | { 142 | 143 | // Extract batch points features and labels 144 | vector b_o_points = vector(original_points.begin () + sum_b, 145 | original_points.begin () + sum_b + original_batches[b]); 146 | 147 | vector b_o_features; 148 | if (original_features.size() > 0) 149 | { 150 | b_o_features = vector(original_features.begin () + sum_b * fdim, 151 | original_features.begin () + (sum_b + original_batches[b]) * fdim); 152 | } 153 | 154 | vector b_o_classes; 155 | if (original_classes.size() > 0) 156 | { 157 | b_o_classes = vector(original_classes.begin () + sum_b * ldim, 158 | original_classes.begin () + sum_b + original_batches[b] * ldim); 159 | } 160 | 161 | 162 | // Create result containers 163 | vector b_s_points; 164 | vector b_s_features; 165 | vector b_s_classes; 166 | 167 | // Compute subsampling on current batch 168 | grid_subsampling(b_o_points, 169 | b_s_points, 170 | b_o_features, 171 | b_s_features, 172 | b_o_classes, 173 | b_s_classes, 174 | sampleDl, 175 | 0); 176 | 177 | // Stack batches points features and labels 178 | // **************************************** 179 | 180 | // If too many points remove some 181 | if (b_s_points.size() <= max_p) 182 | { 183 | subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.end()); 184 | 185 | if (original_features.size() > 0) 186 | subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.end()); 187 | 188 | if (original_classes.size() > 0) 189 | subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.end()); 190 | 191 | subsampled_batches.push_back(b_s_points.size()); 192 | } 193 | else 194 | { 195 | subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.begin() + max_p); 196 | 197 | if (original_features.size() > 0) 198 | subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.begin() + max_p * fdim); 199 | 200 | if (original_classes.size() > 0) 201 | subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.begin() + max_p * ldim); 202 | 203 | subsampled_batches.push_back(max_p); 204 | } 205 | 206 | // Stack new batch lengths 207 | sum_b += original_batches[b]; 208 | } 209 | 210 | return; 211 | } 212 | -------------------------------------------------------------------------------- /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 | void grid_subsampling(vector& original_points, 83 | vector& subsampled_points, 84 | vector& original_features, 85 | vector& subsampled_features, 86 | vector& original_classes, 87 | vector& subsampled_classes, 88 | float sampleDl, 89 | int verbose); 90 | 91 | void batch_grid_subsampling(vector& original_points, 92 | vector& subsampled_points, 93 | vector& original_features, 94 | vector& subsampled_features, 95 | vector& original_classes, 96 | vector& subsampled_classes, 97 | vector& original_batches, 98 | vector& subsampled_batches, 99 | float sampleDl, 100 | int max_p); 101 | 102 | -------------------------------------------------------------------------------- /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 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 11 | "grid_subsampling/grid_subsampling.cpp", 12 | "wrapper.cpp"] 13 | 14 | module = Extension(name="grid_subsampling", 15 | sources=SOURCES, 16 | extra_compile_args=['-std=c++11', 17 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 18 | 19 | 20 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /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 batch of stacked pointclouds"; 12 | 13 | static char subsample_docstring[] = "function subsampling a pointcloud"; 14 | 15 | static char subsample_batch_docstring[] = "function subsampling a batch of stacked pointclouds"; 16 | 17 | 18 | // Declare the functions 19 | // ********************* 20 | 21 | static PyObject *cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds); 22 | static PyObject *batch_subsampling(PyObject *self, PyObject *args, PyObject *keywds); 23 | 24 | 25 | // Specify the members of the module 26 | // ********************************* 27 | 28 | static PyMethodDef module_methods[] = 29 | { 30 | { "subsample", (PyCFunction)cloud_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_docstring }, 31 | { "subsample_batch", (PyCFunction)batch_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_batch_docstring }, 32 | {NULL, NULL, 0, NULL} 33 | }; 34 | 35 | 36 | // Initialize the module 37 | // ********************* 38 | 39 | static struct PyModuleDef moduledef = 40 | { 41 | PyModuleDef_HEAD_INIT, 42 | "grid_subsampling", // m_name 43 | module_docstring, // m_doc 44 | -1, // m_size 45 | module_methods, // m_methods 46 | NULL, // m_reload 47 | NULL, // m_traverse 48 | NULL, // m_clear 49 | NULL, // m_free 50 | }; 51 | 52 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 53 | { 54 | import_array(); 55 | return PyModule_Create(&moduledef); 56 | } 57 | 58 | 59 | // Definition of the batch_subsample method 60 | // ********************************** 61 | 62 | static PyObject* batch_subsampling(PyObject* self, PyObject* args, PyObject* keywds) 63 | { 64 | 65 | // Manage inputs 66 | // ************* 67 | 68 | // Args containers 69 | PyObject* points_obj = NULL; 70 | PyObject* features_obj = NULL; 71 | PyObject* classes_obj = NULL; 72 | PyObject* batches_obj = NULL; 73 | 74 | // Keywords containers 75 | static char* kwlist[] = { "points", "batches", "features", "classes", "sampleDl", "method", "max_p", "verbose", NULL }; 76 | float sampleDl = 0.1; 77 | const char* method_buffer = "barycenters"; 78 | int verbose = 0; 79 | int max_p = 0; 80 | 81 | // Parse the input 82 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "OO|$OOfsii", kwlist, &points_obj, &batches_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &max_p, &verbose)) 83 | { 84 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 85 | return NULL; 86 | } 87 | 88 | // Get the method argument 89 | string method(method_buffer); 90 | 91 | // Interpret method 92 | if (method.compare("barycenters") && method.compare("voxelcenters")) 93 | { 94 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 95 | return NULL; 96 | } 97 | 98 | // Check if using features or classes 99 | bool use_feature = true, use_classes = true; 100 | if (features_obj == NULL) 101 | use_feature = false; 102 | if (classes_obj == NULL) 103 | use_classes = false; 104 | 105 | // Interpret the input objects as numpy arrays. 106 | PyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 107 | PyObject* batches_array = PyArray_FROM_OTF(batches_obj, NPY_INT, NPY_IN_ARRAY); 108 | PyObject* features_array = NULL; 109 | PyObject* classes_array = NULL; 110 | if (use_feature) 111 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 112 | if (use_classes) 113 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 114 | 115 | // Verify data was load correctly. 116 | if (points_array == NULL) 117 | { 118 | Py_XDECREF(points_array); 119 | Py_XDECREF(batches_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (batches_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(batches_array); 129 | Py_XDECREF(classes_array); 130 | Py_XDECREF(features_array); 131 | PyErr_SetString(PyExc_RuntimeError, "Error converting input batches to numpy arrays of type int32"); 132 | return NULL; 133 | } 134 | if (use_feature && features_array == NULL) 135 | { 136 | Py_XDECREF(points_array); 137 | Py_XDECREF(batches_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 141 | return NULL; 142 | } 143 | if (use_classes && classes_array == NULL) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(batches_array); 147 | Py_XDECREF(classes_array); 148 | Py_XDECREF(features_array); 149 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 150 | return NULL; 151 | } 152 | 153 | // Check that the input array respect the dims 154 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 155 | { 156 | Py_XDECREF(points_array); 157 | Py_XDECREF(batches_array); 158 | Py_XDECREF(classes_array); 159 | Py_XDECREF(features_array); 160 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 161 | return NULL; 162 | } 163 | if ((int)PyArray_NDIM(batches_array) > 1) 164 | { 165 | Py_XDECREF(points_array); 166 | Py_XDECREF(batches_array); 167 | Py_XDECREF(classes_array); 168 | Py_XDECREF(features_array); 169 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : batches.shape is not (B,) "); 170 | return NULL; 171 | } 172 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 173 | { 174 | Py_XDECREF(points_array); 175 | Py_XDECREF(batches_array); 176 | Py_XDECREF(classes_array); 177 | Py_XDECREF(features_array); 178 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 179 | return NULL; 180 | } 181 | 182 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 183 | { 184 | Py_XDECREF(points_array); 185 | Py_XDECREF(batches_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 | // Number of points 193 | int N = (int)PyArray_DIM(points_array, 0); 194 | 195 | // Number of batches 196 | int Nb = (int)PyArray_DIM(batches_array, 0); 197 | 198 | // Dimension of the features 199 | int fdim = 0; 200 | if (use_feature) 201 | fdim = (int)PyArray_DIM(features_array, 1); 202 | 203 | //Dimension of labels 204 | int ldim = 1; 205 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 206 | ldim = (int)PyArray_DIM(classes_array, 1); 207 | 208 | // Check that the input array respect the number of points 209 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 210 | { 211 | Py_XDECREF(points_array); 212 | Py_XDECREF(batches_array); 213 | Py_XDECREF(classes_array); 214 | Py_XDECREF(features_array); 215 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 216 | return NULL; 217 | } 218 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 219 | { 220 | Py_XDECREF(points_array); 221 | Py_XDECREF(batches_array); 222 | Py_XDECREF(classes_array); 223 | Py_XDECREF(features_array); 224 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 225 | return NULL; 226 | } 227 | 228 | 229 | // Call the C++ function 230 | // ********************* 231 | 232 | // Create pyramid 233 | if (verbose > 0) 234 | cout << "Computing cloud pyramid with support points: " << endl; 235 | 236 | 237 | // Convert PyArray to Cloud C++ class 238 | vector original_points; 239 | vector original_batches; 240 | vector original_features; 241 | vector original_classes; 242 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 243 | original_batches = vector((int*)PyArray_DATA(batches_array), (int*)PyArray_DATA(batches_array) + Nb); 244 | if (use_feature) 245 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim); 246 | if (use_classes) 247 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim); 248 | 249 | // Subsample 250 | vector subsampled_points; 251 | vector subsampled_features; 252 | vector subsampled_classes; 253 | vector subsampled_batches; 254 | batch_grid_subsampling(original_points, 255 | subsampled_points, 256 | original_features, 257 | subsampled_features, 258 | original_classes, 259 | subsampled_classes, 260 | original_batches, 261 | subsampled_batches, 262 | sampleDl, 263 | max_p); 264 | 265 | // Check result 266 | if (subsampled_points.size() < 1) 267 | { 268 | PyErr_SetString(PyExc_RuntimeError, "Error"); 269 | return NULL; 270 | } 271 | 272 | // Manage outputs 273 | // ************** 274 | 275 | // Dimension of input containers 276 | npy_intp* point_dims = new npy_intp[2]; 277 | point_dims[0] = subsampled_points.size(); 278 | point_dims[1] = 3; 279 | npy_intp* feature_dims = new npy_intp[2]; 280 | feature_dims[0] = subsampled_points.size(); 281 | feature_dims[1] = fdim; 282 | npy_intp* classes_dims = new npy_intp[2]; 283 | classes_dims[0] = subsampled_points.size(); 284 | classes_dims[1] = ldim; 285 | npy_intp* batches_dims = new npy_intp[1]; 286 | batches_dims[0] = Nb; 287 | 288 | // Create output array 289 | PyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 290 | PyObject* res_batches_obj = PyArray_SimpleNew(1, batches_dims, NPY_INT); 291 | PyObject* res_features_obj = NULL; 292 | PyObject* res_classes_obj = NULL; 293 | PyObject* ret = NULL; 294 | 295 | // Fill output array with values 296 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 297 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 298 | size_in_bytes = Nb * sizeof(int); 299 | memcpy(PyArray_DATA(res_batches_obj), subsampled_batches.data(), size_in_bytes); 300 | if (use_feature) 301 | { 302 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 303 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 304 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 305 | } 306 | if (use_classes) 307 | { 308 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 309 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 310 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 311 | } 312 | 313 | 314 | // Merge results 315 | if (use_feature && use_classes) 316 | ret = Py_BuildValue("NNNN", res_points_obj, res_batches_obj, res_features_obj, res_classes_obj); 317 | else if (use_feature) 318 | ret = Py_BuildValue("NNN", res_points_obj, res_batches_obj, res_features_obj); 319 | else if (use_classes) 320 | ret = Py_BuildValue("NNN", res_points_obj, res_batches_obj, res_classes_obj); 321 | else 322 | ret = Py_BuildValue("NN", res_points_obj, res_batches_obj); 323 | 324 | // Clean up 325 | // ******** 326 | 327 | Py_DECREF(points_array); 328 | Py_DECREF(batches_array); 329 | Py_XDECREF(features_array); 330 | Py_XDECREF(classes_array); 331 | 332 | return ret; 333 | } 334 | 335 | // Definition of the subsample method 336 | // **************************************** 337 | 338 | static PyObject* cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds) 339 | { 340 | 341 | // Manage inputs 342 | // ************* 343 | 344 | // Args containers 345 | PyObject* points_obj = NULL; 346 | PyObject* features_obj = NULL; 347 | PyObject* classes_obj = NULL; 348 | 349 | // Keywords containers 350 | static char* kwlist[] = { "points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 351 | float sampleDl = 0.1; 352 | const char* method_buffer = "barycenters"; 353 | int verbose = 0; 354 | 355 | // Parse the input 356 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 357 | { 358 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 359 | return NULL; 360 | } 361 | 362 | // Get the method argument 363 | string method(method_buffer); 364 | 365 | // Interpret method 366 | if (method.compare("barycenters") && method.compare("voxelcenters")) 367 | { 368 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 369 | return NULL; 370 | } 371 | 372 | // Check if using features or classes 373 | bool use_feature = true, use_classes = true; 374 | if (features_obj == NULL) 375 | use_feature = false; 376 | if (classes_obj == NULL) 377 | use_classes = false; 378 | 379 | // Interpret the input objects as numpy arrays. 380 | PyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 381 | PyObject* features_array = NULL; 382 | PyObject* classes_array = NULL; 383 | if (use_feature) 384 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 385 | if (use_classes) 386 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 387 | 388 | // Verify data was load correctly. 389 | if (points_array == NULL) 390 | { 391 | Py_XDECREF(points_array); 392 | Py_XDECREF(classes_array); 393 | Py_XDECREF(features_array); 394 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 395 | return NULL; 396 | } 397 | if (use_feature && features_array == NULL) 398 | { 399 | Py_XDECREF(points_array); 400 | Py_XDECREF(classes_array); 401 | Py_XDECREF(features_array); 402 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 403 | return NULL; 404 | } 405 | if (use_classes && classes_array == NULL) 406 | { 407 | Py_XDECREF(points_array); 408 | Py_XDECREF(classes_array); 409 | Py_XDECREF(features_array); 410 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 411 | return NULL; 412 | } 413 | 414 | // Check that the input array respect the dims 415 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 416 | { 417 | Py_XDECREF(points_array); 418 | Py_XDECREF(classes_array); 419 | Py_XDECREF(features_array); 420 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 421 | return NULL; 422 | } 423 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 424 | { 425 | Py_XDECREF(points_array); 426 | Py_XDECREF(classes_array); 427 | Py_XDECREF(features_array); 428 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 429 | return NULL; 430 | } 431 | 432 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 433 | { 434 | Py_XDECREF(points_array); 435 | Py_XDECREF(classes_array); 436 | Py_XDECREF(features_array); 437 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 438 | return NULL; 439 | } 440 | 441 | // Number of points 442 | int N = (int)PyArray_DIM(points_array, 0); 443 | 444 | // Dimension of the features 445 | int fdim = 0; 446 | if (use_feature) 447 | fdim = (int)PyArray_DIM(features_array, 1); 448 | 449 | //Dimension of labels 450 | int ldim = 1; 451 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 452 | ldim = (int)PyArray_DIM(classes_array, 1); 453 | 454 | // Check that the input array respect the number of points 455 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 456 | { 457 | Py_XDECREF(points_array); 458 | Py_XDECREF(classes_array); 459 | Py_XDECREF(features_array); 460 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 461 | return NULL; 462 | } 463 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 464 | { 465 | Py_XDECREF(points_array); 466 | Py_XDECREF(classes_array); 467 | Py_XDECREF(features_array); 468 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 469 | return NULL; 470 | } 471 | 472 | 473 | // Call the C++ function 474 | // ********************* 475 | 476 | // Create pyramid 477 | if (verbose > 0) 478 | cout << "Computing cloud pyramid with support points: " << endl; 479 | 480 | 481 | // Convert PyArray to Cloud C++ class 482 | vector original_points; 483 | vector original_features; 484 | vector original_classes; 485 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 486 | if (use_feature) 487 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim); 488 | if (use_classes) 489 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim); 490 | 491 | // Subsample 492 | vector subsampled_points; 493 | vector subsampled_features; 494 | vector subsampled_classes; 495 | grid_subsampling(original_points, 496 | subsampled_points, 497 | original_features, 498 | subsampled_features, 499 | original_classes, 500 | subsampled_classes, 501 | sampleDl, 502 | verbose); 503 | 504 | // Check result 505 | if (subsampled_points.size() < 1) 506 | { 507 | PyErr_SetString(PyExc_RuntimeError, "Error"); 508 | return NULL; 509 | } 510 | 511 | // Manage outputs 512 | // ************** 513 | 514 | // Dimension of input containers 515 | npy_intp* point_dims = new npy_intp[2]; 516 | point_dims[0] = subsampled_points.size(); 517 | point_dims[1] = 3; 518 | npy_intp* feature_dims = new npy_intp[2]; 519 | feature_dims[0] = subsampled_points.size(); 520 | feature_dims[1] = fdim; 521 | npy_intp* classes_dims = new npy_intp[2]; 522 | classes_dims[0] = subsampled_points.size(); 523 | classes_dims[1] = ldim; 524 | 525 | // Create output array 526 | PyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 527 | PyObject* res_features_obj = NULL; 528 | PyObject* res_classes_obj = NULL; 529 | PyObject* ret = NULL; 530 | 531 | // Fill output array with values 532 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 533 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 534 | if (use_feature) 535 | { 536 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 537 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 538 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 539 | } 540 | if (use_classes) 541 | { 542 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 543 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 544 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 545 | } 546 | 547 | 548 | // Merge results 549 | if (use_feature && use_classes) 550 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 551 | else if (use_feature) 552 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 553 | else if (use_classes) 554 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 555 | else 556 | ret = Py_BuildValue("N", res_points_obj); 557 | 558 | // Clean up 559 | // ******** 560 | 561 | Py_DECREF(points_array); 562 | Py_XDECREF(features_array); 563 | Py_XDECREF(classes_array); 564 | 565 | return ret; 566 | } -------------------------------------------------------------------------------- /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 | // Initialize 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 | // Initialize limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /cpp_wrappers/cpp_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 | struct PointCloud 151 | { 152 | 153 | std::vector pts; 154 | 155 | // Must return the number of data points 156 | inline size_t kdtree_get_point_count() const { return pts.size(); } 157 | 158 | // Returns the dim'th component of the idx'th point in the class: 159 | // Since this is inlined and the "dim" argument is typically an immediate value, the 160 | // "if/else's" are actually solved at compile time. 161 | inline float kdtree_get_pt(const size_t idx, const size_t dim) const 162 | { 163 | if (dim == 0) return pts[idx].x; 164 | else if (dim == 1) return pts[idx].y; 165 | else return pts[idx].z; 166 | } 167 | 168 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 169 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 170 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 171 | template 172 | bool kdtree_get_bbox(BBOX& /* bb */) const { return false; } 173 | 174 | }; 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | -------------------------------------------------------------------------------- /data_utils/modelnet40_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | from pathlib import Path 4 | 5 | import lmdb 6 | import msgpack_numpy 7 | import numpy as np 8 | from tqdm import tqdm 9 | 10 | from jittor.dataset.dataset import Dataset 11 | 12 | BASE_DIR = Path(__file__).parent 13 | 14 | class ModelNet40(Dataset): 15 | def __init__(self, n_points: int, train: bool, batch_size=1, shuffle=False): 16 | super().__init__() 17 | self.n_points = n_points 18 | self.batch_size = batch_size 19 | self.train = train 20 | self.shuffle = shuffle 21 | 22 | self.path = BASE_DIR / 'data' / 'modelnet40_normal_resampled' 23 | self.cache = BASE_DIR / 'data' / 'modelnet40_normal_resampled_cache' 24 | self.cache.mkdir(exist_ok=True) 25 | 26 | if not self.path.exists(): 27 | self.path.mkdir(parents=True) 28 | self.url = ( 29 | "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip" 30 | ) 31 | zipfile = self.path / '..' / 'modelnet40_normal_resampled.zip' 32 | 33 | if not zipfile.exists(): 34 | subprocess.check_call([ 35 | 'curl', self.url, '-o', str(zipfile), '-k' 36 | ]) 37 | 38 | subprocess.check_call([ 39 | 'unzip', str(zipfile), '-d', str(self.path / '..') 40 | ]) 41 | 42 | cats_file = self.path / 'modelnet40_shape_names.txt' 43 | with cats_file.open() as f: 44 | cats = [line.rstrip() for line in f.readlines()] 45 | self.classes = dict(zip(cats, range(len(cats)))) 46 | 47 | train = 'train' if self.train else 'test' 48 | shapes_file = self.path / f'modelnet40_{train}.txt' 49 | with shapes_file.open() as f: 50 | self.shapes = [] 51 | for line in f.readlines(): 52 | shape_id = line.rstrip() 53 | shape_name = '_'.join(shape_id.split('_')[0:-1]) 54 | self.shapes.append(( 55 | shape_name, 56 | shape_id + '.txt', 57 | )) 58 | 59 | self.lmdb_file = self.cache / train 60 | self.lmdb_env = None 61 | if not self.lmdb_file.exists(): 62 | # create lmdb file 63 | with lmdb.open(str(self.lmdb_file), map_size=1 << 36) as lmdb_env: 64 | with lmdb_env.begin(write=True) as txn: 65 | for i, (shape_name, shape_file) in enumerate(tqdm(self.shapes)): 66 | shape_path = self.path / shape_name / shape_file 67 | pts = np.loadtxt(shape_path, delimiter=',', dtype=np.float32) 68 | cls = self.classes[shape_name] 69 | 70 | txn.put( 71 | str(i).encode(), 72 | msgpack_numpy.packb(pts, use_bin_type=True), 73 | ) 74 | 75 | self.set_attrs( 76 | batch_size=self.batch_size, 77 | total_len=len(self.shapes), 78 | shuffle=self.shuffle 79 | ) 80 | 81 | def __getitem__(self, idx): 82 | if self.lmdb_env is None: 83 | self.lmdb_env = lmdb.open(str(self.lmdb_file), map_size=1<<36, 84 | readonly=True, lock=False) 85 | 86 | shape_name, _ = self.shapes[idx] 87 | with self.lmdb_env.begin(buffers=True) as txn: 88 | pts = msgpack_numpy.unpackb(txn.get(str(idx).encode()), raw=False) 89 | 90 | pt_idxs = np.arange(0, self.n_points) 91 | np.random.shuffle(pt_idxs) 92 | 93 | pts = pts[pt_idxs, :] 94 | pts, normals = pts[:, :3], pts[:, 3:] 95 | 96 | if self.train: 97 | pts = self.translate_pointcloud(self.normalize_pointclouds(pts)) 98 | # pts, normals = self.random_point_dropout(pts, normals) 99 | else : 100 | pts = self.normalize_pointclouds(pts) 101 | 102 | return pts, normals, self.classes[shape_name] 103 | 104 | def random_point_dropout(self, pc, normal, max_dropout_ratio=0.875): 105 | ''' pc: Nx3 ''' 106 | # for b in range(batch_pc.shape[0]): 107 | dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875 108 | drop_idx = np.where(np.random.random((pc.shape[0]))<=dropout_ratio)[0] 109 | if len(drop_idx)>0: 110 | pc[drop_idx,:] = pc[0,:] # set to the first point 111 | normal[drop_idx,:] = normal[0,:] # normal set to the first point 112 | 113 | return pc, normal 114 | 115 | def collect_batch(self, batch): 116 | pts = np.stack([b[0] for b in batch], axis=0) 117 | normals = np.stack([b[1] for b in batch], axis=0) 118 | cls = np.stack([b[2] for b in batch]) 119 | return pts, cls 120 | 121 | def normalize_pointclouds(self, pts): 122 | pts = pts - pts.mean(axis=0) 123 | scale = np.sqrt((pts ** 2).sum(axis=1).max()) 124 | pts = pts / scale 125 | return pts 126 | 127 | 128 | def translate_pointcloud(self, pointcloud): 129 | xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) 130 | xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) 131 | translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') 132 | return translated_pointcloud 133 | 134 | 135 | if __name__ == '__main__': 136 | modelnet40 = ModelNet40(n_points=2048, train=True, batch_size=32, shuffle=True) 137 | for pts, normals, cls in modelnet40: 138 | print (pts.size(), normals.size()) 139 | break 140 | -------------------------------------------------------------------------------- /data_utils/shapenet_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import glob 4 | import h5py 5 | import numpy as np 6 | 7 | # import jittor related 8 | import jittor as jt 9 | from jittor.dataset.dataset import Dataset 10 | 11 | 12 | def download_shapenetpart(): 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | DATA_DIR = os.path.join(BASE_DIR, 'data') 15 | if not os.path.exists(DATA_DIR): 16 | os.mkdir(DATA_DIR) 17 | if not os.path.exists(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data')): 18 | www = 'https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip' 19 | zipfile = os.path.basename(www) 20 | os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile)) 21 | os.system('mv %s %s' % (zipfile[:-4], os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data'))) 22 | os.system('rm %s' % (zipfile)) 23 | 24 | 25 | def load_data_partseg(partition): 26 | download_shapenetpart() 27 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 28 | DATA_DIR = os.path.join(BASE_DIR, 'data') 29 | all_data = [] 30 | all_label = [] 31 | all_seg = [] 32 | if partition == 'trainval': 33 | file = glob.glob(os.path.join(DATA_DIR, 'shapenet*hdf5*', '*train*.h5')) \ 34 | + glob.glob(os.path.join(DATA_DIR, 'shapenet*hdf5*', '*val*.h5')) 35 | else: 36 | file = glob.glob(os.path.join(DATA_DIR, 'shapenet*hdf5*', '*%s*.h5'%partition)) 37 | for h5_name in file: 38 | f = h5py.File(h5_name, 'r+') 39 | data = f['data'][:].astype('float32') 40 | label = f['label'][:].astype('int64') 41 | seg = f['pid'][:].astype('int64') 42 | f.close() 43 | all_data.append(data) 44 | all_label.append(label) 45 | all_seg.append(seg) 46 | all_data = np.concatenate(all_data, axis=0) 47 | all_label = np.concatenate(all_label, axis=0) 48 | all_seg = np.concatenate(all_seg, axis=0) 49 | return all_data, all_label, all_seg 50 | 51 | 52 | class ShapeNetPart(Dataset): 53 | def __init__(self, num_points, partition='train', class_choice=None, batch_size=16, shuffle=False): 54 | super().__init__() 55 | self.batch_size = batch_size 56 | self.shuffle = shuffle 57 | 58 | self.data, self.label, self.seg = load_data_partseg(partition) 59 | self.cat2id = {'airplane': 0, 'bag': 1, 'cap': 2, 'car': 3, 'chair': 4, 60 | 'earphone': 5, 'guitar': 6, 'knife': 7, 'lamp': 8, 'laptop': 9, 61 | 'motor': 10, 'mug': 11, 'pistol': 12, 'rocket': 13, 'skateboard': 14, 'table': 15} 62 | self.seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 63 | self.index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 64 | self.num_points = num_points 65 | self.partition = partition 66 | self.class_choice = class_choice 67 | 68 | if self.class_choice != None: 69 | id_choice = self.cat2id[self.class_choice] 70 | indices = (self.label == id_choice).squeeze() 71 | self.data = self.data[indices] 72 | self.label = self.label[indices] 73 | self.seg = self.seg[indices] 74 | self.seg_num_all = self.seg_num[id_choice] 75 | self.seg_start_index = self.index_start[id_choice] 76 | else: 77 | self.seg_num_all = 50 78 | self.seg_start_index = 0 79 | 80 | total_lens = self.data.shape[0] 81 | # print (total_lens) 82 | self.set_attrs( 83 | batch_size=self.batch_size, 84 | total_len=total_lens, 85 | shuffle=self.shuffle, 86 | drop_last=True, 87 | ) 88 | 89 | 90 | def __getitem__(self, item): 91 | pointcloud = self.data[item][:self.num_points] 92 | label = self.label[item] 93 | seg = self.seg[item][:self.num_points] 94 | if self.partition == 'trainval': 95 | # pointcloud = translate_pointcloud(pointcloud) 96 | indices = list(range(pointcloud.shape[0])) 97 | np.random.shuffle(indices) 98 | pointcloud = pointcloud[indices] 99 | seg = seg[indices] 100 | return pointcloud, label, seg 101 | 102 | # def collect_batch(self, batch): 103 | # pts = np.stack([b[0] for b in batch], axis=0) 104 | # label = np.stack([b[1] for b in batch], axis=0) 105 | # seg = np.stack([b[2] for b in batch], axis=0) 106 | # return pts, label, seg 107 | 108 | 109 | 110 | if __name__ == '__main__': 111 | 112 | trainval = ShapeNetPart(2048, 'trainval', batch_size=16, shuffle=True) 113 | # test = ShapeNetPart(2048, 'test', batch_size=16) 114 | data, label, seg = trainval[0] 115 | print(data.shape) 116 | print(label.shape) 117 | print(seg.shape) 118 | for data, label, seg in trainval: 119 | print(data.shape) 120 | print(label.shape) 121 | print(seg.shape) 122 | break 123 | 124 | 125 | -------------------------------------------------------------------------------- /misc/layers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jittor as jt 3 | from jittor import nn 4 | from jittor.contrib import concat 5 | import math 6 | from typing import Tuple, Callable, Optional, Union 7 | from misc.ops import FurthestPointSampler, KNN 8 | import time 9 | 10 | 11 | class STN3d(nn.Module): 12 | def __init__(self): 13 | super(STN3d, self).__init__() 14 | self.conv1 = nn.Conv1d(3, 64, 1) 15 | self.conv2 = nn.Conv1d(64, 128, 1) 16 | self.conv3 = nn.Conv1d(128, 1024, 1) 17 | self.fc1 = nn.Linear(1024, 512) 18 | self.fc2 = nn.Linear(512, 256) 19 | self.fc3 = nn.Linear(256, 9) 20 | self.relu = nn.ReLU() 21 | 22 | self.bn1 = nn.BatchNorm1d(64) 23 | self.bn2 = nn.BatchNorm1d(128) 24 | self.bn3 = nn.BatchNorm1d(1024) 25 | self.bn4 = nn.BatchNorm1d(512) 26 | self.bn5 = nn.BatchNorm1d(256) 27 | 28 | 29 | def execute(self, x): 30 | batchsize = x.shape[0] 31 | # print ('x shape =', x.shape) 32 | x = self.conv1(x) 33 | # print ('x shape =', x.shape) 34 | x = self.bn1(x) 35 | x = self.relu(x) 36 | x = self.relu(self.bn2(self.conv2(x))) 37 | x = self.relu(self.bn3(self.conv3(x))) 38 | # print ('before max shape is', x.shape) 39 | x = jt.max(x, 2) 40 | 41 | # print ('after max shape is', x.shape) 42 | x = x.reshape(-1, 1024) 43 | 44 | #x = self.relu(self.bn4(self.fc1(x))) 45 | x = self.fc1(x) 46 | x = self.bn4(x) 47 | x = self.relu(x) 48 | x = self.relu(self.bn5(self.fc2(x))) 49 | x = self.fc3(x) 50 | 51 | iden = ((jt.array(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).reshape(1,9)).repeat(batchsize, 1) 52 | # print (iden.shape) 53 | x = x + iden 54 | # print (x.shape) 55 | x = x.reshape(-1, 3, 3) 56 | return x 57 | 58 | 59 | class STNkd(nn.Module): 60 | def __init__(self, k=64): 61 | super(STNkd, self).__init__() 62 | self.conv1 = nn.Conv1d(k, 64, 1) 63 | self.conv2 = nn.Conv1d(64, 128, 1) 64 | self.conv3 = nn.Conv1d(128, 1024, 1) 65 | self.fc1 = nn.Linear(1024, 512) 66 | self.fc2 = nn.Linear(512, 256) 67 | self.fc3 = nn.Linear(256, k*k) 68 | self.relu = nn.ReLU() 69 | 70 | self.bn1 = nn.BatchNorm1d(64) 71 | self.bn2 = nn.BatchNorm1d(128) 72 | self.bn3 = nn.BatchNorm1d(1024) 73 | self.bn4 = nn.BatchNorm1d(512) 74 | self.bn5 = nn.BatchNorm1d(256) 75 | self.k = k 76 | 77 | def execute(self, x): 78 | batchsize = x.size()[0] 79 | x = self.relu(self.bn1(self.conv1(x))) 80 | x = self.relu(self.bn2(self.conv2(x))) 81 | x = self.relu(self.bn3(self.conv3(x))) 82 | x = jt.max(x, 2) 83 | x = x.reshape(-1, 1024) 84 | 85 | x = self.relu(self.bn4(self.fc1(x))) 86 | x = self.relu(self.bn5(self.fc2(x))) 87 | x = self.fc3(x) 88 | 89 | iden = jt.array(np.eye(self.k).flatten().astype(np.float32)).reshape(1,self.k*self.k) 90 | x = x + iden 91 | x = x.reshape(-1, self.k, self.k) 92 | return x 93 | 94 | 95 | 96 | 97 | def EndChannels(f, make_contiguous=False): 98 | """ Class decorator to apply 2D convolution along end channels. """ 99 | 100 | class WrappedLayer(nn.Module): 101 | 102 | def __init__(self): 103 | super(WrappedLayer, self).__init__() 104 | self.f = f 105 | 106 | def execute(self, x): 107 | x = x.permute(0,3,1,2) 108 | x = self.f(x) 109 | x = x.permute(0,2,3,1) 110 | return x 111 | 112 | return WrappedLayer() 113 | 114 | 115 | def EndChannels1d(f, make_contiguous=False): 116 | """ Class decorator to apply 2D convolution along end channels. """ 117 | 118 | class WrappedLayer(nn.Module): 119 | 120 | def __init__(self): 121 | super(WrappedLayer, self).__init__() 122 | self.f = f 123 | 124 | def execute(self, x): 125 | x = x.permute(0,2,1) 126 | x = self.f(x) 127 | x = x.permute(0,2,1) 128 | return x 129 | 130 | return WrappedLayer() 131 | 132 | 133 | class SepConv(nn.Module): 134 | """ Depthwise separable convolution with optional activation and batch normalization""" 135 | 136 | def __init__(self, in_channels : int, out_channels : int, 137 | kernel_size : Union[int, Tuple[int, int]], 138 | depth_multiplier = 1, with_bn = True, 139 | activation = nn.ReLU()) -> None: 140 | """ 141 | :param in_channels: Length of input featuers (first dimension). 142 | :param out_channels: Length of output features (first dimension). 143 | :param kernel_size: Size of convolutional kernel. 144 | :depth_multiplier: Depth multiplier for middle part of separable convolution. 145 | :param with_bn: Whether or not to apply batch normalization. 146 | :param activation: Activation function. 147 | """ 148 | super(SepConv, self).__init__() 149 | 150 | self.conv = nn.Sequential( 151 | nn.Conv(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels), 152 | nn.Conv(in_channels * depth_multiplier, out_channels, 1, bias = not with_bn) 153 | ) 154 | 155 | self.activation = activation 156 | self.bn = nn.BatchNorm(out_channels, momentum = 0.9) if with_bn else None 157 | 158 | def execute(self, x): 159 | """ 160 | :param x: Any input tensor that can be input into nn.Conv2d. 161 | :return: Tensor with depthwise separable convolutional layer and 162 | optional activation and batchnorm applied. 163 | """ 164 | x = self.conv(x) 165 | if self.activation: 166 | x = self.activation(x) 167 | if self.bn: 168 | x = self.bn(x) 169 | return x 170 | 171 | 172 | 173 | class Conv(nn.Module): 174 | """ 175 | 2D convolutional layer with optional activation and batch normalization. 176 | """ 177 | 178 | def __init__(self, in_channels : int, out_channels : int, 179 | kernel_size, with_bn = True, 180 | activation = nn.ReLU()) -> None: 181 | """ 182 | :param in_channels: Length of input featuers (first dimension). 183 | :param out_channels: Length of output features (first dimension). 184 | :param kernel_size: Size of convolutional kernel. 185 | :param with_bn: Whether or not to apply batch normalization. 186 | :param activation: Activation function. 187 | """ 188 | super(Conv, self).__init__() 189 | 190 | self.conv = nn.Conv(in_channels, out_channels, kernel_size, bias = not with_bn) 191 | self.activation = activation 192 | self.bn = nn.BatchNorm(out_channels, momentum = 0.9) if with_bn else None 193 | 194 | def execute(self, x): 195 | """ 196 | :param x: Any input tensor that can be input into nn.Conv2d. 197 | :return: Tensor with convolutional layer and optional activation and batchnorm applied. 198 | """ 199 | x = self.conv(x) 200 | if self.activation: 201 | x = self.activation(x) 202 | if self.bn: 203 | x = self.bn(x) 204 | return x 205 | 206 | 207 | 208 | class Dense_Conv1d(nn.Module): 209 | def __init__(self, in_features : int, out_features : int, 210 | drop_rate : int = 0, with_bn : bool = True, 211 | activation = nn.ReLU() 212 | ) -> None: 213 | """ 214 | :param in_features: Length of input featuers (last dimension). 215 | :param out_features: Length of output features (last dimension). 216 | :param drop_rate: Drop rate to be applied after activation. 217 | :param with_bn: Whether or not to apply batch normalization. 218 | :param activation: Activation function. 219 | """ 220 | super(Dense_Conv1d, self).__init__() 221 | 222 | self.linear = nn.Conv1d(in_features, out_features, 1) 223 | self.activation = activation 224 | self.with_bn = with_bn 225 | self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None 226 | self.bn = nn.BatchNorm1d(out_features) if with_bn else None 227 | 228 | def execute(self, x): 229 | x = self.linear(x) 230 | 231 | if self.with_bn: 232 | x = self.bn(x) 233 | if self.activation: 234 | x = self.activation(x) 235 | if self.drop: 236 | x = self.drop(x) 237 | return x 238 | 239 | 240 | class Dense_Conv2d(nn.Module): 241 | def __init__(self, in_features : int, out_features : int, 242 | drop_rate : int = 0, with_bn : bool = True, 243 | activation = nn.ReLU(), groups=1) -> None: 244 | """ 245 | :param in_features: Length of input featuers (last dimension). 246 | :param out_features: Length of output features (last dimension). 247 | :param drop_rate: Drop rate to be applied after activation. 248 | :param with_bn: Whether or not to apply batch normalization. 249 | :param activation: Activation function. 250 | """ 251 | super(Dense_Conv2d, self).__init__() 252 | 253 | self.linear = nn.Conv(in_features, out_features, 1, groups=groups) 254 | self.activation = activation 255 | self.with_bn = with_bn 256 | self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None 257 | self.bn = nn.BatchNorm(out_features) if with_bn else None 258 | 259 | def execute(self, x): 260 | x = self.linear(x) 261 | 262 | # print ('before bn shape', x.shape) 263 | if self.with_bn: 264 | x = self.bn(x) 265 | if self.activation: 266 | x = self.activation(x) 267 | if self.drop: 268 | x = self.drop(x) 269 | 270 | return x 271 | 272 | 273 | class RandPointCNN_Decoder(nn.Module): 274 | """ PointCNN with randomly subsampled representative points. """ 275 | 276 | def __init__(self, C_in : int, C_out : int, C_last : int, dims : int, K : int, D : int, P : int) -> None: 277 | """ See documentation for PointCNN. """ 278 | super(RandPointCNN_Decoder, self).__init__() 279 | self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P) 280 | self.P = P 281 | self.conv_fuse = EndChannels1d(Dense_Conv1d(C_out + C_last, C_out)) 282 | # if self.P > 0: 283 | # self.sampler = FurthestPointSampler(self.P) 284 | 285 | def execute(self, x_l, x_h): # (N, P, C_out) 286 | """ 287 | Given a point cloud, and its corresponding features, return a new set 288 | of randomly-sampled representative points with features projected from 289 | the point cloud. 290 | :param x: (pts, fts) where 291 | - pts: Regional point cloud such that fts[:,p_idx,:] is the 292 | feature associated with pts[:,p_idx,:]. 293 | - fts: Regional features such that pts[:,p_idx,:] is the feature 294 | associated with fts[:,p_idx,:]. 295 | :return: Randomly subsampled points and their features. 296 | """ 297 | pts_l, fts_l = x_l 298 | pts_h, fts_h = x_h 299 | rep_pts_fts = self.pointcnn((pts_h, pts_l, fts_l)) 300 | concat_feature = concat ((rep_pts_fts, fts_h), dim=2) 301 | rep_pts_fts = self.conv_fuse(concat_feature) 302 | return pts_h, rep_pts_fts 303 | 304 | 305 | class RandPointCNN(nn.Module): 306 | """ PointCNN with randomly subsampled representative points. """ 307 | 308 | def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int) -> None: 309 | """ See documentation for PointCNN. """ 310 | super(RandPointCNN, self).__init__() 311 | self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P) 312 | self.P = P 313 | if self.P > 0: 314 | self.sampler = FurthestPointSampler(self.P) 315 | 316 | def execute(self, x): # (N, P, C_out) 317 | """ 318 | Given a point cloud, and its corresponding features, return a new set 319 | of randomly-sampled representative points with features projected from 320 | the point cloud. 321 | :param x: (pts, fts) where 322 | - pts: Regional point cloud such that fts[:,p_idx,:] is the 323 | feature associated with pts[:,p_idx,:]. 324 | - fts: Regional features such that pts[:,p_idx,:] is the feature 325 | associated with fts[:,p_idx,:]. 326 | :return: Randomly subsampled points and their features. 327 | """ 328 | pts, fts = x 329 | if 0 < self.P < pts.size()[1]: 330 | rep_pts = self.sampler(pts) 331 | # idx = np.random.choice(pts.size()[1], self.P, replace = False).tolist() 332 | # rep_pts = pts[:,idx,:] 333 | 334 | else: 335 | rep_pts = pts 336 | rep_pts_fts = self.pointcnn((rep_pts, pts, fts)) 337 | return rep_pts, rep_pts_fts 338 | 339 | 340 | 341 | class PointCNN(nn.Module): 342 | """ Pointwise convolutional model. """ 343 | 344 | def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int) -> None: 345 | """ 346 | :param C_in: Input dimension of the points' features. 347 | :param C_out: Output dimension of the representative point features. 348 | :param dims: Spatial dimensionality of points. 349 | :param K: Number of neighbors to convolve over. 350 | :param D: "Spread" of neighboring points. 351 | :param P: Number of representative points. 352 | :param r_indices_func: Selector function of the type, 353 | INPUTS 354 | rep_pts : Representative points. 355 | pts : Point cloud. 356 | K : Number of points for each region. 357 | D : "Spread" of neighboring points. 358 | 359 | OUTPUT 360 | pts_idx : Array of indices into pts such that pts[pts_idx] is the set 361 | of points in the "region" around rep_pt. 362 | """ 363 | super(PointCNN, self).__init__() 364 | 365 | C_mid = C_out // 2 if C_in == 0 else C_out // 4 366 | 367 | if C_in == 0: 368 | depth_multiplier = 4 369 | else: 370 | # depth_multiplier = min(int(np.ceil(C_out / C_in)), 4) 371 | depth_multiplier = int(np.ceil(C_out / C_in)) 372 | 373 | self.knn = KNN(K * D) 374 | self.dense = EndChannels1d(Dense_Conv1d(C_in, C_out // 2)) if C_in != 0 else None 375 | 376 | self.x_conv = XConv(C_out // 2 if C_in != 0 else C_in, C_out, dims, K, P, C_mid, depth_multiplier) 377 | 378 | self.D = D 379 | self.K = K 380 | 381 | def select_region(self, pts, # (N, x, dims) 382 | pts_idx): # (P, K, dims) 383 | 384 | regions = jt.stack([ 385 | pts[n][idx,:] for n, idx in enumerate(jt.misc.unbind(pts_idx, dim = 0)) 386 | ], dim = 0) 387 | 388 | return regions 389 | 390 | def execute(self, x): # (N, P, C_out) 391 | 392 | rep_pts, pts, fts = x 393 | # print ('input size =', rep_pts.shape, fts.shape, fts.shape) 394 | fts = self.dense(fts) if fts is not None else fts # B, N, D 395 | tmp_rep_pts = rep_pts.permute(0, 2, 1) 396 | tmp_pts = pts.permute(0, 2, 1) 397 | # pts_idx = knn_indices_func_gpu(rep_pts, pts, self.K, self.D) 398 | pts_idx = self.knn(tmp_rep_pts, tmp_pts) # b, d, n 399 | pts_idx = pts_idx[:,0::self.D, :] 400 | pts_idx = pts_idx.permute(0, 2, 1) 401 | # jt.sync_all(True) 402 | # start_time = time.time() 403 | pts_regional = self.select_region(pts, pts_idx) 404 | fts_regional = self.select_region(fts, pts_idx) if fts is not None else fts 405 | # jt.sync_all(True) 406 | # end_time = time.time() 407 | # print ('select region run time =', end_time - start_time) 408 | fts_p = self.x_conv((rep_pts, pts_regional, fts_regional)) 409 | return fts_p 410 | 411 | 412 | 413 | class XConv(nn.Module): 414 | """ Convolution over a single point and its neighbors. """ 415 | 416 | def __init__(self, C_in : int, C_out : int, dims : int, K : int, 417 | P : int, C_mid : int, depth_multiplier : int) -> None: 418 | super(XConv, self).__init__() 419 | 420 | self.C_in = C_in 421 | self.C_mid = C_mid 422 | self.dims = dims 423 | self.K = K 424 | 425 | self.P = P 426 | 427 | # Additional processing layers 428 | # self.pts_layernorm = LayerNorm(2, momentum = 0.9) 429 | 430 | # Main dense linear layers 431 | self.dense1 = Dense_Conv2d(dims, C_mid) 432 | self.dense2 = Dense_Conv2d(C_mid, C_mid) 433 | 434 | # Layers to generate X 435 | self.x_trans_0 = Conv( 436 | in_channels = dims, 437 | out_channels = K*K, 438 | kernel_size = (1, K), 439 | with_bn = True) 440 | self.x_trans_1 = Dense_Conv2d(K*K, K*K, with_bn = True, groups=1) 441 | self.x_trans_2 = Dense_Conv2d(K*K, K*K, with_bn = False, activation = None, groups=1) 442 | 443 | self.end_conv = EndChannels(SepConv( 444 | in_channels = C_mid + C_in, 445 | out_channels = C_out, 446 | kernel_size = (1, K), 447 | depth_multiplier = depth_multiplier 448 | )) 449 | 450 | def execute(self, x): # (N, K, C_out) 451 | """ 452 | Applies XConv to the input data. 453 | :param x: (rep_pt, pts, fts) where 454 | - rep_pt: Representative point. 455 | - pts: Regional point cloud such that fts[:,p_idx,:] is the feature 456 | associated with pts[:,p_idx,:]. 457 | - fts: Regional features such that pts[:,p_idx,:] is the feature 458 | associated with fts[:,p_idx,:]. 459 | :return: Features aggregated into point rep_pt. 460 | """ 461 | rep_pt, pts, fts = x # b, n, c // b ,n k, c // b, n, k, d 462 | if fts is not None: 463 | assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0]) # Check N is equal. 464 | assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1]) # Check P is equal. 465 | assert(pts.size()[2] == fts.size()[2] == self.K) # Check K is equal. 466 | assert(fts.size()[3] == self.C_in) # Check C_in is equal. 467 | else: 468 | assert(rep_pt.size()[0] == pts.size()[0]) # Check N is equal. 469 | assert(rep_pt.size()[1] == pts.size()[1]) # Check P is equal. 470 | assert(pts.size()[2] == self.K) # Check K is equal. 471 | assert(rep_pt.size()[2] == pts.size()[3] == self.dims) # Check dims is equal. 472 | 473 | N = pts.size()[0] 474 | P = rep_pt.size()[1] # (N, P, K, dims) 475 | p_center = jt.unsqueeze(rep_pt, dim = 2) # (N, P, 1, dims) 476 | # print (p_center.size()) # 477 | # Move pts to local coordinate system of rep_pt. 478 | pts_local = pts - p_center.repeat(1, 1, self.K, 1) # (N, P, K, dims) 479 | # pts_local = self.pts_layernorm(pts - p_center) 480 | 481 | # Individually lift each point into C_mid space. 482 | # print (pts_local.size(), 'before size') 483 | pts_local = pts_local.permute(0, 3, 1, 2) # N, dim, P, K 484 | fts_lifted0 = self.dense1(pts_local) # ? 485 | # print (.size(), 'after size') 486 | fts_lifted = self.dense2(fts_lifted0) # N, C_mid, P, K 487 | 488 | fts = fts.permute(0, 3, 1, 2) 489 | if fts is None: 490 | fts_cat = fts_lifted 491 | else: 492 | fts_cat = concat((fts_lifted, fts), 1) # (N, C_mid + C_in, P, K) 493 | 494 | # Learn the (N, K, K) X-transformation matrix. 495 | X_shape = (N, P, self.K, self.K) 496 | # X = self.x_trans(pts_local) # N, K*K, 1, P 497 | x = self.x_trans_0(pts_local) 498 | x = self.x_trans_1(x) 499 | X = self.x_trans_2(x) 500 | 501 | # print ('X size ', X.size()) 502 | X = X.permute(0, 2, 3, 1) # n p 1 k 503 | X = X.view(X_shape) # N, P, K, K 504 | 505 | 506 | # print (fts_cat.shape) 507 | fts_cat = fts_cat.permute(0, 2, 3, 1) 508 | fts_X = jt.matmul(X, fts_cat) # 509 | 510 | # print ('fts X size =', fts_X.shape) 511 | 512 | fts_p = self.end_conv(fts_X).squeeze(dim = 2) 513 | # print ('xxxxxxxxxxx') 514 | # print ('result size') 515 | # print (fts_X.size(), fts_p.size()) 516 | 517 | return fts_p 518 | -------------------------------------------------------------------------------- /misc/pointconv_utils.py: -------------------------------------------------------------------------------- 1 | 2 | from time import time 3 | import numpy as np 4 | import random 5 | #from sklearn.neighbors.kde import KernelDensity 6 | 7 | import jittor as jt 8 | from jittor import nn 9 | from jittor import Module 10 | from jittor import init 11 | from jittor.contrib import concat 12 | 13 | def timeit(tag, t): 14 | print("{}: {}s".format(tag, time() - t)) 15 | return time() 16 | 17 | def topk(input, k, dim=None, largest=True, sorted=True): 18 | if dim is None: 19 | dim = -1 20 | if dim<0: 21 | dim+=input.ndim 22 | 23 | transpose_dims = [i for i in range(input.ndim)] 24 | transpose_dims[0] = dim 25 | transpose_dims[dim] = 0 26 | input = input.transpose(transpose_dims) 27 | index,values = jt.argsort(input,dim=0,descending=largest) 28 | indices = index[:k] 29 | values = values[:k] 30 | indices = indices.transpose(transpose_dims) 31 | values = values.transpose(transpose_dims) 32 | return values,indices 33 | 34 | def square_distance(src, dst): 35 | """ 36 | Calculate Euclid distance between each two points. 37 | src^T * dst = xn * xm + yn * ym + zn * zm； 38 | sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; 39 | sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; 40 | dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 41 | = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst 42 | Input: 43 | src: source points, [B, N, C] 44 | dst: target points, [B, M, C] 45 | Output: 46 | dist: per-point square distance, [B, N, M] 47 | """ 48 | B, N, _ = src.shape 49 | _, M, _ = dst.shape 50 | dist = -2 * jt.matmul(src, dst.permute(0, 2, 1)) 51 | dist += jt.sum(src ** 2, -1).view(B, N, 1) 52 | dist += jt.sum(dst ** 2, -1).view(B, 1, M) 53 | return dist 54 | 55 | def index_points(points, idx): 56 | """ 57 | Input: 58 | points: input points data, [B, N, C] 59 | idx: sample index data, [B, S] 60 | Return: 61 | new_points:, indexed points data, [B, S, C] 62 | """ 63 | #device = points.device 64 | B = points.shape[0] 65 | view_shape = list(idx.shape) 66 | view_shape[1:] = [1] * (len(view_shape) - 1) 67 | repeat_shape = list(idx.shape) 68 | repeat_shape[0] = 1 69 | batch_indices = np.arange(B, dtype='l') 70 | batch_indices = jt.array(batch_indices).view(view_shape).repeat(repeat_shape) 71 | new_points = points[batch_indices, idx, :] 72 | return new_points 73 | 74 | def farthest_point_sample(xyz, npoint): 75 | """ 76 | Input: 77 | xyz: pointcloud data, [B, N, C] 78 | npoint: number of samples 79 | Return: 80 | centroids: sampled pointcloud index, [B, npoint] 81 | """ 82 | #import ipdb; ipdb.set_trace() 83 | #device = xyz.device 84 | B, N, C = xyz.shape 85 | centroids = jt.zeros((B, npoint)) 86 | distance = jt.ones((B, N)) * 1e10 87 | 88 | farthest = np.random.randint(0, N, B, dtype='l') 89 | batch_indices = np.arange(B, dtype='l') 90 | farthest = jt.array(farthest) 91 | batch_indices = jt.array(batch_indices) 92 | # jt.sync_all(True) 93 | # print (xyz.shape, farthest.shape, batch_indices.shape, centroids.shape, distance.shape) 94 | for i in range(npoint): 95 | centroids[:, i] = farthest 96 | centroid = xyz[batch_indices, farthest, :] 97 | centroid = centroid.view(B, 1, 3) 98 | 99 | dist = jt.sum((xyz - centroid.repeat(1, N, 1)) ** 2, 2) 100 | mask = dist < distance 101 | # distance = mask.ternary(distance, dist) 102 | # print (mask.size()) 103 | 104 | if mask.sum().data[0] > 0: 105 | distance[mask] = dist[mask] # bug if mask.sum() == 0 106 | 107 | farthest = jt.argmax(distance, 1)[0] 108 | # print (farthest) 109 | # print (farthest.shape) 110 | # B, N, C = xyz.size() 111 | # sample_list = random.sample(range(0, N), npoint) 112 | # centroids = jt.zeros((1, npoint)) 113 | # centroids[0,:] = jt.array(sample_list) 114 | # centroids = centroids.view(1, -1).repeat(B, 1) 115 | # x_center = x[:,sample_list, :] 116 | return centroids 117 | 118 | 119 | 120 | def knn_point(nsample, xyz, new_xyz): 121 | """ 122 | Input: 123 | nsample: max sample number in local region 124 | xyz: all points, [B, N, C] 125 | new_xyz: query points, [B, S, C] 126 | Return: 127 | group_idx: grouped points index, [B, S, nsample] 128 | """ 129 | sqrdists = square_distance(new_xyz, xyz) 130 | _, group_idx = topk(sqrdists, nsample, dim = -1, largest=False, sorted=False) 131 | return group_idx 132 | 133 | def sample_and_group(npoint, nsample, xyz, points, density_scale = None): 134 | """ 135 | Input: 136 | npoint: 137 | nsample: 138 | xyz: input points position data, [B, N, C] 139 | points: input points data, [B, N, D] 140 | Return: 141 | new_xyz: sampled points position data, [B, 1, C] 142 | new_points: sampled points data, [B, 1, N, C+D] 143 | """ 144 | B, N, C = xyz.shape 145 | S = npoint 146 | 147 | fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C] 148 | # jt.sync_all(True) 149 | # print ('11111111111111111') 150 | new_xyz = index_points(xyz, fps_idx) 151 | # jt.sync_all(True) 152 | # print ('2222222222222222222') 153 | idx = knn_point(nsample, xyz, new_xyz) 154 | # jt.sync_all(True) 155 | # print ('333333333333333333') 156 | grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C] 157 | grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C) 158 | if points is not None: 159 | grouped_points = index_points(points, idx) 160 | new_points = concat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D] 161 | else: 162 | new_points = grouped_xyz_norm 163 | # jt.sync_all(True) 164 | # print ('44444444444444444444444') 165 | 166 | if density_scale is None: 167 | return new_xyz, new_points, grouped_xyz_norm, idx 168 | else: 169 | grouped_density = index_points(density_scale, idx) 170 | return new_xyz, new_points, grouped_xyz_norm, idx, grouped_density 171 | 172 | 173 | 174 | def compute_density(xyz, bandwidth): 175 | ''' 176 | xyz: input points position data, [B, N, C] 177 | ''' 178 | #import ipdb; ipdb.set_trace() 179 | B, N, C = xyz.shape 180 | sqrdists = square_distance(xyz, xyz) 181 | gaussion_density = jt.exp(- sqrdists / (2.0 * bandwidth * bandwidth)) / (2.5 * bandwidth) 182 | xyz_density = gaussion_density.mean(dim = -1) 183 | 184 | return xyz_density 185 | 186 | class DensityNet(nn.Module): 187 | def __init__(self, hidden_unit = [8, 8]): 188 | super(DensityNet, self).__init__() 189 | self.mlp_convs = nn.ModuleList() 190 | self.mlp_bns = nn.ModuleList() 191 | 192 | self.mlp_convs.append(nn.Conv1d(1, hidden_unit[0], 1)) 193 | self.mlp_bns.append(nn.BatchNorm1d(hidden_unit[0])) 194 | for i in range(1, len(hidden_unit)): 195 | self.mlp_convs.append(nn.Conv1d(hidden_unit[i - 1], hidden_unit[i], 1)) 196 | self.mlp_bns.append(nn.BatchNorm1d(hidden_unit[i])) 197 | self.mlp_convs.append(nn.Conv1d(hidden_unit[-1], 1, 1)) 198 | self.mlp_bns.append(nn.BatchNorm1d(1)) 199 | self.sigmoid = nn.Sigmoid() 200 | self.relu = nn.ReLU() 201 | def execute(self, xyz_density): 202 | B, N = xyz_density.shape 203 | density_scale = xyz_density.unsqueeze(1) 204 | 205 | # for i, conv in enumerate(self.mlp_convs): 206 | for i in range(len(self.mlp_convs)): 207 | # print ('xxxxxxxxxxx', i, len(self.mlp_bns), len(self.mlp_convs)) 208 | bn = self.mlp_bns[i] 209 | conv = self.mlp_convs[i] 210 | # print ('after get bn') 211 | density_scale = bn(conv(density_scale)) 212 | # print ('after get desity scale') 213 | if i == len(self.mlp_convs): 214 | density_scale = self.sigmoid(density_scale) + 0.5 215 | else: 216 | density_scale = self.relu(density_scale) 217 | 218 | return density_scale 219 | 220 | class WeightNet(nn.Module): 221 | 222 | def __init__(self, in_channel, out_channel, hidden_unit = [8, 8]): 223 | super(WeightNet, self).__init__() 224 | 225 | self.mlp_convs = nn.ModuleList() 226 | self.mlp_bns = nn.ModuleList() 227 | self.relu = nn.ReLU() 228 | if hidden_unit is None or len(hidden_unit) == 0: 229 | self.mlp_convs.append(nn.Conv(in_channel, out_channel, 1)) 230 | self.mlp_bns.append(nn.BatchNorm(out_channel)) 231 | else: 232 | self.mlp_convs.append(nn.Conv(in_channel, hidden_unit[0], 1)) 233 | self.mlp_bns.append(nn.BatchNorm(hidden_unit[0])) 234 | for i in range(1, len(hidden_unit)): 235 | self.mlp_convs.append(nn.Conv(hidden_unit[i - 1], hidden_unit[i], 1)) 236 | self.mlp_bns.append(nn.BatchNorm(hidden_unit[i])) 237 | self.mlp_convs.append(nn.Conv(hidden_unit[-1], out_channel, 1)) 238 | self.mlp_bns.append(nn.BatchNorm(out_channel)) 239 | 240 | def execute(self, localized_xyz): 241 | #xyz : BxCxKxN 242 | 243 | weights = localized_xyz 244 | 245 | for i in range (len(self.mlp_convs)): 246 | conv = self.mlp_convs[i] 247 | bn = self.mlp_bns[i] 248 | weights = self.relu(bn(conv(weights))) 249 | 250 | return weights 251 | 252 | 253 | class PointConvDensitySetInterpolation(nn.Module): 254 | def __init__(self, nsample, in_channel, mlp, bandwidth): 255 | super(PointConvDensitySetInterpolation, self).__init__() 256 | self.bandwidth = bandwidth 257 | self.nsample = nsample 258 | self.in_channel = in_channel 259 | self.mlp_convs = nn.ModuleList() 260 | self.mlp_bns = nn.ModuleList() 261 | self.relu = nn.ReLU() 262 | last_channel = in_channel 263 | self.weightnet = WeightNet(3, 16) 264 | self.densitynet = DensityNet() 265 | 266 | for out_channel in mlp: 267 | self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1)) 268 | self.mlp_bns.append(nn.BatchNorm2d(out_channel)) 269 | last_channel = out_channel 270 | 271 | self.linear = nn.Linear(16 * mlp[-1], mlp[-1]) 272 | self.bn_linear = nn.BatchNorm1d(mlp[-1]) 273 | 274 | def execute(self, xyz1, xyz2, points1, points2): 275 | """ 276 | Input: 277 | xyz1: input points position data, [B, C, N] 278 | xyz2: sampled input points position data, [B, C, S] 279 | points1: input points data, [B, D, N] 280 | points2: input points data, [B, D, S] 281 | Return: 282 | new_points: upsampled points data, [B, D', N] 283 | """ 284 | # print ('xyz1.shape, xyz2.shape') 285 | # print (xyz1.shape, xyz2.shape, points1.shape, points2.shape) 286 | 287 | xyz1 = xyz1.permute(0, 2, 1) 288 | xyz2 = xyz2.permute(0, 2, 1) 289 | points1 = points1.permute(0, 2, 1) 290 | points2 = points2.permute(0, 2, 1) 291 | B, N, C = xyz1.shape 292 | _, S, _ = xyz2.shape 293 | 294 | # points2 = points2.permute(0, 2, 1) 295 | # print (xyz1.shape, xyz2.shape) 296 | dists = square_distance(xyz1, xyz2) 297 | idx, dists = jt.argsort(dists, dim=-1) 298 | dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3] 299 | 300 | dist_recip = 1.0 / (dists + 1e-8) 301 | norm = jt.sum(dist_recip, dim=2, keepdims=True) 302 | weight = dist_recip / norm 303 | interpolated_points = jt.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2) 304 | 305 | # print ('interpolated_points shape', interpolated_points.shape) 306 | 307 | xyz_density = compute_density(xyz1, self.bandwidth) 308 | density_scale = self.densitynet(xyz_density) 309 | 310 | new_xyz, new_points, grouped_xyz_norm, _, grouped_density = sample_and_group(N, self.nsample, xyz1, interpolated_points, density_scale.reshape(B, N, 1)) 311 | 312 | new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint] 313 | 314 | for i in range(len(self.mlp_convs)): 315 | conv = self.mlp_convs[i] 316 | bn = self.mlp_bns[i] 317 | # print ('new new new point shape', new_points.shape) 318 | new_points = self.relu(bn(conv(new_points))) 319 | 320 | grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1) 321 | weights = self.weightnet(grouped_xyz) 322 | new_points = new_points * grouped_density.permute(0, 3, 2, 1) 323 | new_points = jt.matmul(new_points.permute(0, 3, 1, 2), weights.permute(0, 3, 2, 1)).reshape(B, N, -1) 324 | new_points = self.linear(new_points) 325 | new_points = self.bn_linear(new_points.permute(0, 2, 1)) 326 | new_points = self.relu(new_points) 327 | new_xyz = new_xyz.permute(0, 2, 1) 328 | 329 | return new_points 330 | 331 | # new_points = new_points.permute(0, 2, 1) 332 | # # l = len(self.mlp_convs) 333 | # for i, conv in self.mlp_convs.layers.items(): 334 | # # conv = self.mlp_convs[i] 335 | # bn = self.mlp_bns[i] 336 | # new_points = self.relu(bn(conv(new_points))) 337 | # return new_points.permute(0, 2, 1) 338 | 339 | 340 | class PointConvDensitySetAbstraction(nn.Module): 341 | def __init__(self, npoint, nsample, in_channel, mlp, bandwidth, group_all): 342 | super(PointConvDensitySetAbstraction, self).__init__() 343 | self.npoint = npoint 344 | self.nsample = nsample 345 | self.mlp_convs = nn.ModuleList() 346 | self.mlp_bns = nn.ModuleList() 347 | last_channel = in_channel 348 | for out_channel in mlp: 349 | self.mlp_convs.append(nn.Conv(last_channel, out_channel, 1)) 350 | self.mlp_bns.append(nn.BatchNorm(out_channel)) 351 | last_channel = out_channel 352 | 353 | self.weightnet = WeightNet(3, 16) 354 | self.densitynet = DensityNet() 355 | 356 | self.linear = nn.Linear(16 * mlp[-1], mlp[-1]) 357 | self.bn_linear = nn.BatchNorm1d(mlp[-1]) 358 | self.group_all = group_all 359 | self.bandwidth = bandwidth 360 | self.relu = nn.ReLU() 361 | def execute(self, xyz, points): 362 | """ 363 | Input: 364 | xyz: input points position data, [B, C, N] 365 | points: input points data, [B, D, N] 366 | Return: 367 | new_xyz: sampled points position data, [B, C, S] 368 | new_points_concat: sample points feature data, [B, D', S] 369 | """ 370 | B = xyz.shape[0] 371 | N = xyz.shape[2] 372 | xyz = xyz.permute(0, 2, 1) 373 | if points is not None: 374 | points = points.permute(0, 2, 1) 375 | 376 | xyz_density = compute_density(xyz, self.bandwidth) 377 | density_scale = self.densitynet(xyz_density) 378 | 379 | if self.group_all: 380 | new_xyz, new_points, grouped_xyz_norm, grouped_density = sample_and_group_all(xyz, points, density_scale.reshape(B, N, 1)) 381 | else: 382 | new_xyz, new_points, grouped_xyz_norm, _, grouped_density = sample_and_group(self.npoint, self.nsample, xyz, points, density_scale.reshape(B, N, 1)) 383 | 384 | new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint] 385 | for i in range(len(self.mlp_convs)): 386 | # print ('new_point shape', new_points.shape) 387 | conv = self.mlp_convs[i] 388 | bn = self.mlp_bns[i] 389 | new_points = self.relu(bn(conv(new_points))) 390 | 391 | grouped_xyz = grouped_xyz_norm.permute(0, 3, 2, 1) 392 | weights = self.weightnet(grouped_xyz) 393 | new_points = new_points * grouped_density.permute(0, 3, 2, 1) 394 | new_points = jt.matmul(new_points.permute(0, 3, 1, 2), weights.permute(0, 3, 2, 1)).reshape(B, self.npoint, -1) 395 | new_points = self.linear(new_points) 396 | new_points = self.bn_linear(new_points.permute(0, 2, 1)) 397 | new_points = self.relu(new_points) 398 | new_xyz = new_xyz.permute(0, 2, 1) 399 | 400 | return new_xyz, new_points 401 | 402 | 403 | -------------------------------------------------------------------------------- /misc/utils.py: -------------------------------------------------------------------------------- 1 | from jittor import nn 2 | import jittor as jt 3 | from sklearn.neighbors import NearestNeighbors 4 | import math 5 | import numpy as np 6 | 7 | 8 | class LRScheduler: 9 | def __init__(self, optimizer, base_lr): 10 | self.optimizer = optimizer 11 | 12 | self.basic_lr = base_lr 13 | self.lr_decay = 0.6 14 | self.decay_step = 15000 15 | 16 | def step(self, step): 17 | lr_decay = self.lr_decay ** int(step / self.decay_step) 18 | lr_decay = max(lr_decay, 2e-5) 19 | self.optimizer.lr = lr_decay * self.basic_lr 20 | 21 | 22 | def knn_indices_func_cpu(rep_pts, # (N, pts, dim) 23 | pts, # (N, x, dim) 24 | K : int, 25 | D : int): 26 | """ 27 | CPU-based Indexing function based on K-Nearest Neighbors search. 28 | :param rep_pts: Representative points. 29 | :param pts: Point cloud to get indices from. 30 | :param K: Number of nearest neighbors to collect. 31 | :param D: "Spread" of neighboring points. 32 | :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:] 33 | is the set k-nearest neighbors for the representative points in pts[n]. 34 | """ 35 | rep_pts = rep_pts.data 36 | pts = pts.data 37 | region_idx = [] 38 | 39 | for n, p in enumerate(rep_pts): 40 | P_particular = pts[n] 41 | nbrs = NearestNeighbors(D*K + 1, algorithm = "ball_tree").fit(P_particular) 42 | indices = nbrs.kneighbors(p)[1] 43 | region_idx.append(indices[:,1::D]) 44 | 45 | region_idx = jt.array(np.stack(region_idx, axis = 0)) 46 | return region_idx 47 | 48 | def knn_indices_func_gpu(rep_pts, # (N, pts, dim) 49 | pts, # (N, x, dim) 50 | k : int, d : int ): # (N, pts, K) 51 | """ 52 | GPU-based Indexing function based on K-Nearest Neighbors search. 53 | Very memory intensive, and thus unoptimal for large numbers of points. 54 | :param rep_pts: Representative points. 55 | :param pts: Point cloud to get indices from. 56 | :param K: Number of nearest neighbors to collect. 57 | :param D: "Spread" of neighboring points. 58 | :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:] 59 | is the set k-nearest neighbors for the representative points in pts[n]. 60 | """ 61 | region_idx = [] 62 | batch_size = rep_pts.shape[0] 63 | for idx in range (batch_size): 64 | qry = rep_pts[idx] 65 | ref = pts[idx] 66 | n, d = ref.shape 67 | m, d = qry.shape 68 | mref = ref.view(1, n, d).repeat(m, 1, 1) 69 | mqry = qry.view(m, 1, d).repeat(1, n, 1) 70 | 71 | dist2 = jt.sum((mqry - mref)**2, 2) # pytorch has squeeze 72 | _, inds = topk(dist2, k*d + 1, dim = 1, largest = False) 73 | 74 | region_idx.append(inds[:,1::d]) 75 | 76 | region_idx = jt.stack(region_idx, dim = 0) 77 | 78 | return region_idx 79 | 80 | 81 | 82 | def expand(x,shape): 83 | r''' 84 | Returns a new view of the self tensor with singleton dimensions expanded to a larger size. 85 | Tensor can be also expanded to a larger number of dimensions, and the new ones will be appended at the front. 86 | Args: 87 | x-the input tensor. 88 | shape-the shape of expanded tensor. 89 | ''' 90 | x_shape = x.shape 91 | x_l = len(x_shape) 92 | rest_shape=shape[:-x_l] 93 | expand_shape = shape[-x_l:] 94 | indexs=[] 95 | ii = len(rest_shape) 96 | for i,j in zip(expand_shape,x_shape): 97 | if i!=j: 98 | assert j==1 99 | indexs.append(f'i{ii}' if j>1 else f'0') 100 | ii+=1 101 | return x.reindex(shape,indexs) 102 | 103 | 104 | def topk(input, k, dim=None, largest=True, sorted=True): 105 | if dim is None: 106 | dim = -1 107 | if dim<0: 108 | dim+=input.ndim 109 | 110 | transpose_dims = [i for i in range(input.ndim)] 111 | transpose_dims[0] = dim 112 | transpose_dims[dim] = 0 113 | input = input.transpose(transpose_dims) 114 | index,values = jt.argsort(input,dim=0,descending=largest) 115 | indices = index[:k] 116 | values = values[:k] 117 | indices = indices.transpose(transpose_dims) 118 | values = values.transpose(transpose_dims) 119 | return [values,indices] 120 | -------------------------------------------------------------------------------- /networks/cls/blocks.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0=================================0 4 | # | Kernel Point Convolutions | 5 | # 0=================================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # Define network blocks 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 06/03/2020 15 | # 16 | 17 | 18 | import time 19 | import math 20 | import jittor as jt 21 | import jittor.nn as nn 22 | from jittor.init import kaiming_uniform_ 23 | from networks.cls.kernel_points import load_kernels 24 | 25 | # ---------------------------------------------------------------------------------------------------------------------- 26 | # 27 | # Simple functions 28 | # \**********************/ 29 | # 30 | 31 | 32 | def gather(x, idx, method=2): 33 | """ 34 | implementation of a custom gather operation for faster backwards. 35 | :param x: input with shape [N, D_1, ... D_d] 36 | :param idx: indexing with shape [n_1, ..., n_m] 37 | :param method: Choice of the method 38 | :return: x[idx] with shape [n_1, ..., n_m, D_1, ... D_d] 39 | """ 40 | 41 | if method == 0: 42 | return x[idx] 43 | elif method == 1: 44 | x = x.unsqueeze(1) 45 | x = x.expand((-1, idx.shape[-1], -1)) 46 | idx = idx.unsqueeze(2) 47 | idx = idx.expand((-1, -1, x.shape[-1])) 48 | return x.gather(0, idx) 49 | elif method == 2: 50 | for i, ni in enumerate(idx.size()[1:]): 51 | x = x.unsqueeze(i+1) 52 | new_s = list(x.size()) 53 | new_s[i+1] = ni 54 | x = x.expand(new_s) 55 | n = len(idx.size()) 56 | # print("forloop size 2:", len(x.size()[n:])) 57 | for i, di in enumerate(x.size()[n:]): 58 | idx = idx.unsqueeze(i+n) 59 | new_s = list(idx.size()) 60 | new_s[i+n] = di 61 | idx = idx.expand(new_s) 62 | return x.gather(0, idx) 63 | else: 64 | raise ValueError('Unkown method') 65 | 66 | 67 | def radius_gaussian(sq_r, sig, eps=1e-9): 68 | """ 69 | Compute a radius gaussian (gaussian of distance) 70 | :param sq_r: input radiuses [dn, ..., d1, d0] 71 | :param sig: extents of gaussians [d1, d0] or [d0] or float 72 | :return: gaussian of sq_r [dn, ..., d1, d0] 73 | """ 74 | return jt.exp(-sq_r / (2 * sig**2 + eps)) 75 | 76 | def max_pool(x, inds): 77 | """ 78 | Pools features with the maximum values. 79 | :param x: [n1, d] features matrix 80 | :param inds: [n2, max_num] pooling indices 81 | :return: [n2, d] pooled features matrix 82 | """ 83 | 84 | # Add a last row with minimum features for shadow pools 85 | x = jt.concat((x, jt.zeros_like(x[:1, :])), 0) 86 | 87 | # Get all features for each pooling location [n2, max_num, d] 88 | pool_features = gather(x, inds) 89 | 90 | # Pool the maximum [n2, d] 91 | _, max_features= jt.argmax(pool_features, 1) 92 | return max_features 93 | 94 | 95 | def global_average(x, batch_lengths): 96 | """ 97 | Block performing a global average over batch pooling 98 | :param x: [N, D] input features 99 | :param batch_lengths: [B] list of batch lengths 100 | :return: [B, D] averaged features 101 | """ 102 | 103 | # Loop over the clouds of the batch 104 | averaged_features = [] 105 | i0 = 0 106 | for b_i, length in enumerate(batch_lengths): 107 | # length.shape: [1] 108 | # Average features for each batch cloud 109 | averaged_features.append(jt.mean(x[i0:i0 + length.data[0]], dim=0)) 110 | 111 | # Increment for next cloud 112 | i0 += length.data[0] 113 | 114 | # Average features in each batch 115 | return jt.stack(averaged_features) 116 | 117 | 118 | # ---------------------------------------------------------------------------------------------------------------------- 119 | # 120 | # KPConv class 121 | # \******************/ 122 | # 123 | 124 | 125 | class KPConv(nn.Module): 126 | 127 | def __init__(self, kernel_size, p_dim, in_channels, out_channels, KP_extent, radius, 128 | fixed_kernel_points='center', KP_influence='linear', aggregation_mode='sum', 129 | deformable=False, modulated=False): 130 | """ 131 | Initialize parameters for KPConvDeformable. 132 | :param kernel_size: Number of kernel points. 133 | :param p_dim: dimension of the point space. 134 | :param in_channels: dimension of input features. 135 | :param out_channels: dimension of output features. 136 | :param KP_extent: influence radius of each kernel point. 137 | :param radius: radius used for kernel point init. Even for deformable, use the config.conv_radius 138 | :param fixed_kernel_points: fix position of certain kernel points ('none', 'center' or 'verticals'). 139 | :param KP_influence: influence function of the kernel points ('constant', 'linear', 'gaussian'). 140 | :param aggregation_mode: choose to sum influences, or only keep the closest ('closest', 'sum'). 141 | :param deformable: choose deformable or not 142 | :param modulated: choose if kernel weights are modulated in addition to deformed 143 | """ 144 | super(KPConv, self).__init__() 145 | # print(kernel_size, p_dim, in_channels, out_channels, KP_extent, radius, 146 | # fixed_kernel_points, KP_influence, aggregation_mode, 147 | # deformable, modulated) 148 | # Save parameters 149 | self.K = kernel_size 150 | self.p_dim = p_dim 151 | self.in_channels = in_channels 152 | self.out_channels = out_channels 153 | self.radius = radius 154 | self.KP_extent = KP_extent 155 | self.fixed_kernel_points = fixed_kernel_points 156 | self.KP_influence = KP_influence 157 | self.aggregation_mode = aggregation_mode 158 | self.deformable = deformable 159 | self.modulated = modulated 160 | 161 | # Running variable containing deformed KP distance to input points. (used in regularization loss) 162 | self.min_d2 = None 163 | self.deformed_KP = None 164 | self.offset_features = None 165 | 166 | # Initialize weights 167 | self.weights = jt.zeros((self.K, in_channels, out_channels), dtype=jt.float32) # Parameter 168 | 169 | # Initiate weights for offsets 170 | if deformable: 171 | if modulated: 172 | self.offset_dim = (self.p_dim + 1) * self.K 173 | else: 174 | self.offset_dim = self.p_dim * self.K 175 | self.offset_conv = KPConv(self.K, 176 | self.p_dim, 177 | self.in_channels, 178 | self.offset_dim, 179 | KP_extent, 180 | radius, 181 | fixed_kernel_points=fixed_kernel_points, 182 | KP_influence=KP_influence, 183 | aggregation_mode=aggregation_mode) 184 | self.offset_bias = jt.zeros(self.offset_dim, dtype=jt.float32) # Parameter 185 | 186 | else: 187 | self.offset_dim = None 188 | self.offset_conv = None 189 | self.offset_bias = None 190 | 191 | # Reset parameters 192 | self.reset_parameters() 193 | 194 | # Initialize kernel points 195 | self.kernel_points = self.init_KP() 196 | return 197 | 198 | def reset_parameters(self): 199 | kaiming_uniform_(self.weights, a=math.sqrt(5)) 200 | if self.deformable: 201 | nn.init.zeros_(self.offset_bias) 202 | return 203 | 204 | def init_KP(self): 205 | """ 206 | Initialize the kernel point positions in a sphere 207 | :return: the tensor of kernel points 208 | """ 209 | 210 | # Create one kernel disposition (as numpy array). Choose the KP distance to center thanks to the KP extent 211 | K_points_numpy = load_kernels(self.radius, 212 | self.K, 213 | dimension=self.p_dim, 214 | fixed=self.fixed_kernel_points) 215 | 216 | return jt.array(K_points_numpy, dtype=jt.float32).stop_grad() # Parameter 217 | 218 | def execute(self, q_pts, s_pts, neighb_inds, x): 219 | # print("in: ", q_pts.shape, s_pts.shape, neighb_inds.shape, x.shape) 220 | ################### 221 | # Offset generation 222 | ################### 223 | 224 | if self.deformable: 225 | 226 | # Get offsets with a KPConv that only takes part of the features 227 | self.offset_features = self.offset_conv(q_pts, s_pts, neighb_inds, x) + self.offset_bias 228 | 229 | if self.modulated: 230 | 231 | # Get offset (in normalized scale) from features 232 | unscaled_offsets = self.offset_features[:, :self.p_dim * self.K] 233 | unscaled_offsets = unscaled_offsets.view(-1, self.K, self.p_dim) 234 | 235 | # Get modulations 236 | modulations = 2 * jt.sigmoid(self.offset_features[:, self.p_dim * self.K:]) 237 | 238 | else: 239 | 240 | # Get offset (in normalized scale) from features 241 | unscaled_offsets = self.offset_features.view(-1, self.K, self.p_dim) 242 | 243 | # No modulations 244 | modulations = None 245 | 246 | # Rescale offset for this layer 247 | offsets = unscaled_offsets * self.KP_extent 248 | 249 | else: 250 | offsets = None 251 | modulations = None 252 | 253 | ###################### 254 | # Deformed convolution 255 | ###################### 256 | 257 | # Add a fake point in the last row for shadow neighbors 258 | s_pts = jt.concat((s_pts, jt.zeros_like(s_pts[:1, :]) + 1e6), 0) 259 | # Get neighbor points [n_points, n_neighbors, dim] 260 | neighbors = s_pts[neighb_inds, :] 261 | # Center every neighborhood 262 | neighbors = neighbors - q_pts.unsqueeze(1) 263 | 264 | # Apply offsets to kernel points [n_points, n_kpoints, dim] 265 | if self.deformable: 266 | self.deformed_KP = offsets + self.kernel_points 267 | deformed_K_points = self.deformed_KP.unsqueeze(1) 268 | else: 269 | deformed_K_points = self.kernel_points 270 | 271 | # Get all difference matrices [n_points, n_neighbors, n_kpoints, dim] 272 | neighbors = neighbors.unsqueeze(2) 273 | differences = neighbors - deformed_K_points 274 | 275 | # Get the square distances [n_points, n_neighbors, n_kpoints] 276 | sq_distances = jt.sum(differences ** 2, dim=3) 277 | 278 | # Optimization by ignoring points outside a deformed KP range 279 | if self.deformable: 280 | 281 | # Save distances for loss 282 | self.min_d2, _ = jt.min(sq_distances, dim=1) 283 | 284 | # Boolean of the neighbors in range of a kernel point [n_points, n_neighbors] 285 | in_range = jt.any(sq_distances < self.KP_extent ** 2, dim=2).type(jt.int32) 286 | 287 | # New value of max neighbors 288 | new_max_neighb = jt.max(jt.sum(in_range, dim=1)) 289 | 290 | # For each row of neighbors, indices of the ones that are in range [n_points, new_max_neighb] 291 | neighb_row_bool, neighb_row_inds = jt.topk(in_range, new_max_neighb.item(), dim=1) 292 | 293 | # Gather new neighbor indices [n_points, new_max_neighb] 294 | new_neighb_inds = neighb_inds.gather(1, neighb_row_inds, sparse_grad=False) 295 | 296 | # Gather new distances to KP [n_points, new_max_neighb, n_kpoints] 297 | neighb_row_inds = neighb_row_inds.unsqueeze(2) 298 | neighb_row_inds = neighb_row_inds.expand(-1, -1, self.K) 299 | sq_distances = sq_distances.gather(1, neighb_row_inds, sparse_grad=False) 300 | 301 | # New shadow neighbors have to point to the last shadow point 302 | new_neighb_inds *= neighb_row_bool 303 | new_neighb_inds -= (neighb_row_bool.type(jt.int64) - 1) * int(s_pts.shape[0] - 1) 304 | else: 305 | new_neighb_inds = neighb_inds 306 | 307 | # Get Kernel point influences [n_points, n_kpoints, n_neighbors] 308 | if self.KP_influence == 'constant': 309 | # Every point get an influence of 1. 310 | all_weights = jt.ones_like(sq_distances) 311 | all_weights = jt.transpose(all_weights, 1, 2) 312 | 313 | elif self.KP_influence == 'linear': 314 | # Influence decrease linearly with the distance, and get to zero when d = KP_extent. 315 | all_weights = jt.clamp(1 - jt.sqrt(sq_distances) / self.KP_extent, min_v=0.0) 316 | all_weights = jt.transpose(all_weights, 1, 2) 317 | 318 | elif self.KP_influence == 'gaussian': 319 | # Influence in gaussian of the distance. 320 | sigma = self.KP_extent * 0.3 321 | all_weights = radius_gaussian(sq_distances, sigma) 322 | all_weights = jt.transpose(all_weights, 1, 2) 323 | else: 324 | raise ValueError('Unknown influence function type (config.KP_influence)') 325 | 326 | # In case of closest mode, only the closest KP can influence each point 327 | if self.aggregation_mode == 'closest': 328 | neighbors_1nn = jt.argmin(sq_distances, dim=2) 329 | all_weights *= jt.transpose(nn.functional.one_hot(neighbors_1nn, self.K), 1, 2) 330 | 331 | elif self.aggregation_mode != 'sum': 332 | raise ValueError("Unknown convolution mode. Should be 'closest' or 'sum'") 333 | 334 | # Add a zero feature for shadow neighbors 335 | x = jt.concat((x, jt.zeros_like(x[:1, :])), 0) 336 | 337 | # Get the features of each neighborhood [n_points, n_neighbors, in_fdim] 338 | neighb_x = gather(x, new_neighb_inds) 339 | 340 | # Apply distance weights [n_points, n_kpoints, in_fdim] 341 | weighted_features = jt.matmul(all_weights, neighb_x) 342 | 343 | # Apply modulations 344 | if self.deformable and self.modulated: 345 | weighted_features *= modulations.unsqueeze(2) 346 | 347 | # Apply network weights [n_kpoints, n_points, out_fdim] 348 | weighted_features = weighted_features.permute((1, 0, 2)) 349 | kernel_outputs = jt.matmul(weighted_features, self.weights) 350 | 351 | # Convolution sum [n_points, out_fdim] 352 | out = jt.sum(kernel_outputs, dim=0) 353 | # print("out:", out.shape) 354 | return out 355 | 356 | def __repr__(self): 357 | return 'KPConv(radius: {:.2f}, in_feat: {:d}, out_feat: {:d})'.format(self.radius, 358 | self.in_channels, 359 | self.out_channels) 360 | 361 | # ---------------------------------------------------------------------------------------------------------------------- 362 | # 363 | # Complex blocks 364 | # \********************/ 365 | # 366 | 367 | def block_decider(block_name, 368 | radius, 369 | in_dim, 370 | out_dim, 371 | layer_ind, 372 | config): 373 | 374 | if block_name == 'unary': 375 | return UnaryBlock(in_dim, out_dim, config.use_batch_norm, config.batch_norm_momentum) 376 | 377 | elif block_name in ['simple', 378 | 'simple_deformable', 379 | 'simple_invariant', 380 | 'simple_equivariant', 381 | 'simple_strided', 382 | 'simple_deformable_strided', 383 | 'simple_invariant_strided', 384 | 'simple_equivariant_strided']: 385 | return SimpleBlock(block_name, in_dim, out_dim, radius, layer_ind, config) 386 | 387 | elif block_name in ['resnetb', 388 | 'resnetb_invariant', 389 | 'resnetb_equivariant', 390 | 'resnetb_deformable', 391 | 'resnetb_strided', 392 | 'resnetb_deformable_strided', 393 | 'resnetb_equivariant_strided', 394 | 'resnetb_invariant_strided']: 395 | return ResnetBottleneckBlock(block_name, in_dim, out_dim, radius, layer_ind, config) 396 | 397 | elif block_name == 'global_average': 398 | return GlobalAverageBlock() 399 | 400 | else: 401 | raise ValueError('Unknown block name in the architecture definition : ' + block_name) 402 | 403 | 404 | class BatchNormBlock(nn.Module): 405 | 406 | def __init__(self, in_dim, use_bn, bn_momentum): 407 | """ 408 | Initialize a batch normalization block. If network does not use batch normalization, replace with biases. 409 | :param in_dim: dimension input features 410 | :param use_bn: boolean indicating if we use Batch Norm 411 | :param bn_momentum: Batch norm momentum 412 | """ 413 | super(BatchNormBlock, self).__init__() 414 | self.bn_momentum = bn_momentum 415 | self.use_bn = use_bn 416 | self.in_dim = in_dim 417 | if self.use_bn: 418 | self.batch_norm = nn.BatchNorm1d(in_dim, momentum=bn_momentum) 419 | #self.batch_norm = nn.InstanceNorm1d(in_dim, momentum=bn_momentum) 420 | else: 421 | self.bias = jt.zeros(in_dim, dtype=jt.float32) # Parameter 422 | return 423 | 424 | def reset_parameters(self): 425 | nn.init.zeros_(self.bias) 426 | 427 | def execute(self, x): 428 | if self.use_bn: 429 | x = x.unsqueeze(2) 430 | x = x.transpose(0, 2) 431 | x = self.batch_norm(x) 432 | x = x.transpose(0, 2) 433 | return x.squeeze(2) 434 | else: 435 | return x + self.bias 436 | 437 | def __repr__(self): 438 | return 'BatchNormBlock(in_feat: {:d}, momentum: {:.3f}, only_bias: {:s})'.format(self.in_dim, 439 | self.bn_momentum, 440 | str(not self.use_bn)) 441 | 442 | 443 | class UnaryBlock(nn.Module): 444 | 445 | def __init__(self, in_dim, out_dim, use_bn, bn_momentum, no_relu=False): 446 | """ 447 | Initialize a standard unary block with its ReLU and BatchNorm. 448 | :param in_dim: dimension input features 449 | :param out_dim: dimension input features 450 | :param use_bn: boolean indicating if we use Batch Norm 451 | :param bn_momentum: Batch norm momentum 452 | """ 453 | 454 | super(UnaryBlock, self).__init__() 455 | self.bn_momentum = bn_momentum 456 | self.use_bn = use_bn 457 | self.no_relu = no_relu 458 | self.in_dim = in_dim 459 | self.out_dim = out_dim 460 | print("Linear, in_dim", in_dim, "out_dim", out_dim) 461 | self.mlp = nn.Linear(in_dim, out_dim, bias=False) 462 | self.batch_norm = BatchNormBlock(out_dim, self.use_bn, self.bn_momentum) 463 | if not no_relu: 464 | self.leaky_relu = nn.LeakyReLU(0.1) 465 | return 466 | 467 | def execute(self, x, batch=None): 468 | x = self.mlp(x) 469 | x = self.batch_norm(x) 470 | if not self.no_relu: 471 | x = self.leaky_relu(x) 472 | return x 473 | 474 | def __repr__(self): 475 | return 'UnaryBlock(in_feat: {:d}, out_feat: {:d}, BN: {:s}, ReLU: {:s})'.format(self.in_dim, 476 | self.out_dim, 477 | str(self.use_bn), 478 | str(not self.no_relu)) 479 | 480 | 481 | class SimpleBlock(nn.Module): 482 | 483 | def __init__(self, block_name, in_dim, out_dim, radius, layer_ind, config): 484 | """ 485 | Initialize a simple convolution block with its ReLU and BatchNorm. 486 | :param in_dim: dimension input features 487 | :param out_dim: dimension input features 488 | :param radius: current radius of convolution 489 | :param config: parameters 490 | """ 491 | super(SimpleBlock, self).__init__() 492 | 493 | # get KP_extent from current radius 494 | current_extent = radius * config.KP_extent / config.conv_radius 495 | 496 | # Get other parameters 497 | self.bn_momentum = config.batch_norm_momentum 498 | self.use_bn = config.use_batch_norm 499 | self.layer_ind = layer_ind 500 | self.block_name = block_name 501 | self.in_dim = in_dim 502 | self.out_dim = out_dim 503 | 504 | # Define the KPConv class 505 | self.KPConv = KPConv(config.num_kernel_points, 506 | config.in_points_dim, 507 | in_dim, 508 | out_dim // 2, 509 | current_extent, 510 | radius, 511 | fixed_kernel_points=config.fixed_kernel_points, 512 | KP_influence=config.KP_influence, 513 | aggregation_mode=config.aggregation_mode, 514 | deformable='deform' in block_name, 515 | modulated=config.modulated) 516 | 517 | # Other opperations 518 | self.batch_norm = BatchNormBlock(out_dim // 2, self.use_bn, self.bn_momentum) 519 | self.leaky_relu = nn.LeakyReLU(0.1) 520 | 521 | return 522 | 523 | def execute(self, x, batch): 524 | 525 | if 'strided' in self.block_name: 526 | q_pts = batch.points[self.layer_ind + 1] 527 | s_pts = batch.points[self.layer_ind] 528 | neighb_inds = batch.pools[self.layer_ind] 529 | else: 530 | q_pts = batch.points[self.layer_ind] 531 | s_pts = batch.points[self.layer_ind] 532 | neighb_inds = batch.neighbors[self.layer_ind] 533 | 534 | x = self.KPConv(q_pts, s_pts, neighb_inds, x) 535 | return self.leaky_relu(self.batch_norm(x)) 536 | 537 | 538 | class ResnetBottleneckBlock(nn.Module): 539 | 540 | def __init__(self, block_name, in_dim, out_dim, radius, layer_ind, config): 541 | """ 542 | Initialize a resnet bottleneck block. 543 | :param in_dim: dimension input features 544 | :param out_dim: dimension input features 545 | :param radius: current radius of convolution 546 | :param config: parameters 547 | """ 548 | super(ResnetBottleneckBlock, self).__init__() 549 | 550 | # get KP_extent from current radius 551 | current_extent = radius * config.KP_extent / config.conv_radius 552 | 553 | # Get other parameters 554 | self.bn_momentum = config.batch_norm_momentum 555 | self.use_bn = config.use_batch_norm 556 | self.block_name = block_name 557 | self.layer_ind = layer_ind 558 | self.in_dim = in_dim 559 | self.out_dim = out_dim 560 | 561 | # First downscaling mlp 562 | if in_dim != out_dim // 4: 563 | self.unary1 = UnaryBlock(in_dim, out_dim // 4, self.use_bn, self.bn_momentum) 564 | else: 565 | self.unary1 = nn.Identity() 566 | 567 | # KPConv block 568 | self.KPConv = KPConv(config.num_kernel_points, 569 | config.in_points_dim, 570 | out_dim // 4, 571 | out_dim // 4, 572 | current_extent, 573 | radius, 574 | fixed_kernel_points=config.fixed_kernel_points, 575 | KP_influence=config.KP_influence, 576 | aggregation_mode=config.aggregation_mode, 577 | deformable='deform' in block_name, 578 | modulated=config.modulated) 579 | self.batch_norm_conv = BatchNormBlock(out_dim // 4, self.use_bn, self.bn_momentum) 580 | 581 | # Second upscaling mlp 582 | self.unary2 = UnaryBlock(out_dim // 4, out_dim, self.use_bn, self.bn_momentum, no_relu=True) 583 | 584 | # Shortcut optional mpl 585 | if in_dim != out_dim: 586 | self.unary_shortcut = UnaryBlock(in_dim, out_dim, self.use_bn, self.bn_momentum, no_relu=True) 587 | else: 588 | self.unary_shortcut = nn.Identity() 589 | 590 | # Other operations 591 | self.leaky_relu = nn.LeakyReLU(0.1) 592 | 593 | return 594 | 595 | def execute(self, features, batch): 596 | 597 | if 'strided' in self.block_name: 598 | q_pts = batch.points[self.layer_ind + 1] 599 | s_pts = batch.points[self.layer_ind] 600 | neighb_inds = batch.pools[self.layer_ind] 601 | else: 602 | q_pts = batch.points[self.layer_ind] 603 | s_pts = batch.points[self.layer_ind] 604 | neighb_inds = batch.neighbors[self.layer_ind] 605 | 606 | # First downscaling mlp 607 | x = self.unary1(features) 608 | 609 | # Convolution 610 | x = self.KPConv(q_pts, s_pts, neighb_inds, x) 611 | x = self.leaky_relu(self.batch_norm_conv(x)) 612 | 613 | # Second upscaling mlp 614 | x = self.unary2(x) 615 | 616 | # Shortcut 617 | if 'strided' in self.block_name: 618 | shortcut = max_pool(features, neighb_inds) 619 | else: 620 | shortcut = features 621 | shortcut = self.unary_shortcut(shortcut) 622 | 623 | return self.leaky_relu(x + shortcut) 624 | 625 | 626 | class GlobalAverageBlock(nn.Module): 627 | 628 | def __init__(self): 629 | """ 630 | Initialize a global average block with its ReLU and BatchNorm. 631 | """ 632 | super(GlobalAverageBlock, self).__init__() 633 | return 634 | 635 | def execute(self, x, batch): 636 | # print(x.shape) 637 | # print(batch.lengths[-1]) 638 | return global_average(x, batch.lengths[-1]) 639 | -------------------------------------------------------------------------------- /networks/cls/data.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/networks/cls/data.txt -------------------------------------------------------------------------------- /networks/cls/dgcnn.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import jittor as jt 4 | from jittor import nn 5 | from jittor.contrib import concat 6 | 7 | from misc.ops import KNN 8 | 9 | import time 10 | 11 | def topk(input, k, dim=None, largest=True, sorted=True): 12 | if dim is None: 13 | dim = -1 14 | if dim<0: 15 | dim+=input.ndim 16 | 17 | transpose_dims = [i for i in range(input.ndim)] 18 | transpose_dims[0] = dim 19 | transpose_dims[dim] = 0 20 | input = input.transpose(transpose_dims) 21 | index,values = jt.argsort(input,dim=0,descending=largest) 22 | indices = index[:k] 23 | values = values[:k] 24 | indices = indices.transpose(transpose_dims) 25 | values = values.transpose(transpose_dims) 26 | return [values,indices] 27 | 28 | 29 | def get_graph_feature(x, knn=None, k=None, idx=None): 30 | batch_size = x.shape[0] 31 | num_points = x.shape[2] 32 | x = x.reshape(batch_size, -1, num_points) 33 | if idx is None: 34 | idx = knn(x,x) # (batch_size, num_points, k) 35 | idx = idx.permute(0, 2, 1) 36 | idx_base = jt.array(np.arange(0, batch_size)).reshape(-1, 1, 1)*num_points 37 | 38 | idx = idx + idx_base 39 | 40 | idx = idx.reshape(-1) 41 | 42 | _, num_dims, _ = x.shape 43 | 44 | x = x.transpose(0, 2, 1) # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 45 | feature = x.reshape(batch_size*num_points, -1)[idx, :] 46 | feature = feature.reshape(batch_size, num_points, k, num_dims) 47 | x = x.reshape(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 48 | 49 | feature = concat((feature-x, x), dim=3).transpose(0, 3, 1, 2) 50 | return feature 51 | 52 | def knn (x, k): 53 | inner = -2 * jt.nn.bmm(x.transpose(0, 2, 1), x) 54 | xx = jt.sum(x ** 2, dim = 1, keepdims=True) 55 | distance = -xx - inner - xx.transpose(0, 2, 1) 56 | idx = topk(distance ,k=k, dim=-1)[1] 57 | return idx 58 | 59 | 60 | 61 | class DGCNN(nn.Module): 62 | def __init__(self, n_classes=40): 63 | super(DGCNN, self).__init__() 64 | self.k = 20 65 | self.knn = KNN(self.k) 66 | self.bn1 = nn.BatchNorm(64) 67 | self.bn2 = nn.BatchNorm(64) 68 | self.bn3 = nn.BatchNorm(128) 69 | self.bn4 = nn.BatchNorm(256) 70 | self.bn5 = nn.BatchNorm1d(1024) 71 | 72 | self.conv1 = nn.Sequential(nn.Conv(6, 64, kernel_size=1, bias=False), 73 | self.bn1, 74 | nn.LeakyReLU(scale=0.2)) 75 | self.conv2 = nn.Sequential(nn.Conv(64*2, 64, kernel_size=1, bias=False), 76 | self.bn2, 77 | nn.LeakyReLU(scale=0.2)) 78 | self.conv3 = nn.Sequential(nn.Conv(64*2, 128, kernel_size=1, bias=False), 79 | self.bn3, 80 | nn.LeakyReLU(scale=0.2)) 81 | self.conv4 = nn.Sequential(nn.Conv(128*2, 256, kernel_size=1, bias=False), 82 | self.bn4, 83 | nn.LeakyReLU(scale=0.2)) 84 | self.conv5 = nn.Sequential(nn.Conv1d(512, 1024, kernel_size=1, bias=False), 85 | self.bn5, 86 | nn.LeakyReLU(scale=0.2)) 87 | self.linear1 = nn.Linear(1024*2, 512, bias=False) 88 | self.bn6 = nn.BatchNorm1d(512) 89 | self.dp1 = nn.Dropout(p=0.5) 90 | self.linear2 = nn.Linear(512, 256) 91 | self.bn7 = nn.BatchNorm1d(256) 92 | self.dp2 = nn.Dropout(p=0.5) 93 | self.linear3 = nn.Linear(256, n_classes) 94 | 95 | def execute(self, x): 96 | #print (x.shape) 97 | # x = x.transpose(0, 2, 1) # b, n, c -> b, c, n 98 | batch_size = x.shape[0] 99 | 100 | x = get_graph_feature(x, knn=self.knn, k=self.k) 101 | x = self.conv1(x) 102 | x1 = x.max(dim=-1, keepdims=False) 103 | x = get_graph_feature(x1, knn=self.knn, k=self.k) 104 | x = self.conv2(x) 105 | x2 = x.max(dim=-1, keepdims=False) 106 | x = get_graph_feature(x2, knn=self.knn, k=self.k) 107 | x = self.conv3(x) 108 | x3 = x.max(dim=-1, keepdims=False) 109 | x = get_graph_feature(x3, knn=self.knn, k=self.k) 110 | x = self.conv4(x) 111 | x4 = x.max(dim=-1, keepdims=False) 112 | x = concat((x1, x2, x3, x4), dim=1) 113 | x = self.conv5(x) # ############ this line 114 | x1 = x.max(dim=2).reshape(batch_size, -1) 115 | x2 = x.mean(dim=2).reshape(batch_size, -1) 116 | x = concat((x1, x2), 1) 117 | x = nn.leaky_relu(self.bn6(self.linear1(x)), scale=0.2) 118 | x = self.dp1(x) 119 | x = nn.leaky_relu(self.bn7(self.linear2(x)), scale=0.2) 120 | x = self.dp2(x) 121 | x = self.linear3(x) 122 | return x 123 | 124 | 125 | def main(): 126 | net = DGCNN() 127 | x = jt.array(np.random.random((32, 3, 256))) 128 | print (x.shape) 129 | out = net(x) 130 | print (out.shape) 131 | if __name__ == '__main__' : 132 | main() 133 | -------------------------------------------------------------------------------- /networks/cls/kernels/dispositions/k_015_center_3D.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jittor/PointCloudLib/e1c37bab3d158d8a26ca3a360ccafd7bd90d80ad/networks/cls/kernels/dispositions/k_015_center_3D.ply -------------------------------------------------------------------------------- /networks/cls/kpconv.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import jittor as jt 4 | # from jittor import nn 5 | from jittor.contrib import concat 6 | from jittor import init 7 | from networks.cls.blocks import * 8 | import pickle 9 | from jittor_utils import auto_diff 10 | from datasets.ModelNet40 import ModelNet40CustomBatch 11 | 12 | def p2p_fitting_regularizer(net): 13 | 14 | fitting_loss = 0 15 | repulsive_loss = 0 16 | 17 | for m in net.modules(): 18 | 19 | if isinstance(m, KPConv) and m.deformable: 20 | 21 | ############## 22 | # Fitting loss 23 | ############## 24 | 25 | # Get the distance to closest input point and normalize to be independant from layers 26 | KP_min_d2 = m.min_d2 / (m.KP_extent ** 2) 27 | 28 | # Loss will be the square distance to closest input point. We use L1 because dist is already squared 29 | fitting_loss += net.l1(KP_min_d2, jt.zeros_like(KP_min_d2)) 30 | 31 | ################ 32 | # Repulsive loss 33 | ################ 34 | 35 | # Normalized KP locations 36 | KP_locs = m.deformed_KP / m.KP_extent 37 | 38 | # Point should not be close to each other 39 | for i in range(net.K): 40 | other_KP = jt.concat([KP_locs[:, :i, :], KP_locs[:, i + 1:, :]], dim=1) 41 | distances = jt.sqrt(jt.sum((other_KP - KP_locs[:, i:i + 1, :]) ** 2, dim=2)) 42 | rep_loss = jt.sum(jt.clamp(distances - net.repulse_extent, max_v=0.0) ** 2, dim=1) 43 | repulsive_loss += net.l1(rep_loss, jt.zeros_like(rep_loss)) / net.K 44 | 45 | return net.deform_fitting_power * (2 * fitting_loss + repulsive_loss) 46 | 47 | 48 | class KPCNN(nn.Module): 49 | """ 50 | Class defining KPCNN 51 | """ 52 | 53 | def __init__(self, config): 54 | super(KPCNN, self).__init__() 55 | 56 | ##################### 57 | # Network opperations 58 | ##################### 59 | 60 | # Current radius of convolution and feature dimension 61 | layer = 0 62 | r = config.first_subsampling_dl * config.conv_radius 63 | in_dim = config.in_features_dim 64 | out_dim = config.first_features_dim 65 | self.K = config.num_kernel_points 66 | 67 | # Save all block operations in a list of modules 68 | self.block_ops = nn.ModuleList() 69 | 70 | # Loop over consecutive blocks 71 | block_in_layer = 0 72 | for block_i, block in enumerate(config.architecture): 73 | 74 | # Check equivariance 75 | if ('equivariant' in block) and (not out_dim % 3 == 0): 76 | raise ValueError('Equivariant block but features dimension is not a factor of 3') 77 | 78 | # Detect upsampling block to stop 79 | if 'upsample' in block: 80 | break 81 | 82 | # Apply the good block function defining tf ops 83 | self.block_ops.append(block_decider(block, 84 | r, 85 | in_dim, 86 | out_dim, 87 | layer, 88 | config)) 89 | 90 | 91 | # Index of block in this layer 92 | block_in_layer += 1 93 | 94 | # Update dimension of input from output 95 | if 'simple' in block: 96 | in_dim = out_dim // 2 97 | else: 98 | in_dim = out_dim 99 | 100 | 101 | # Detect change to a subsampled layer 102 | if 'pool' in block or 'strided' in block: 103 | # Update radius and feature dimension for next layer 104 | layer += 1 105 | r *= 2 106 | out_dim *= 2 107 | block_in_layer = 0 108 | 109 | self.head_mlp = UnaryBlock(out_dim, 1024, False, 0) 110 | self.head_softmax = UnaryBlock(1024, config.num_classes, False, 0, no_relu=True) 111 | 112 | ################ 113 | # Network Losses 114 | ################ 115 | 116 | self.criterion = nn.CrossEntropyLoss() 117 | self.deform_fitting_mode = config.deform_fitting_mode 118 | self.deform_fitting_power = config.deform_fitting_power 119 | self.deform_lr_factor = config.deform_lr_factor 120 | self.repulse_extent = config.repulse_extent 121 | self.output_loss = 0 122 | self.reg_loss = 0 123 | self.l1 = nn.L1Loss() 124 | 125 | return 126 | 127 | def execute(self, batch): 128 | batch = ModelNet40CustomBatch([batch]) 129 | # Save all block operations in a list of modules 130 | x = batch.features.clone().detach() 131 | 132 | # Loop over consecutive blocks 133 | for i, block_op in enumerate(self.block_ops): 134 | # x.sync() 135 | x = block_op(x, batch) 136 | 137 | # Head of network 138 | x = self.head_mlp(x, batch) 139 | x = self.head_softmax(x, batch) 140 | 141 | return x 142 | 143 | def loss(self, outputs, labels): 144 | """ 145 | Runs the loss on outputs of the model 146 | :param outputs: logits 147 | :param labels: labels 148 | :return: loss 149 | """ 150 | 151 | # Cross entropy loss 152 | self.output_loss = self.criterion(outputs, labels) 153 | 154 | # Regularization of deformable offsets 155 | if self.deform_fitting_mode == 'point2point': 156 | self.reg_loss = p2p_fitting_regularizer(self) 157 | elif self.deform_fitting_mode == 'point2plane': 158 | raise ValueError('point2plane fitting mode not implemented yet.') 159 | else: 160 | raise ValueError('Unknown fitting mode: ' + self.deform_fitting_mode) 161 | 162 | # Combined loss 163 | return self.output_loss + self.reg_loss 164 | 165 | @staticmethod 166 | def accuracy(outputs, labels): 167 | """ 168 | Computes accuracy of the current batch 169 | :param outputs: logits predicted by the network 170 | :param labels: labels 171 | :return: accuracy value 172 | """ 173 | 174 | predicted = jt.argmax(outputs.data, dim=1)[0] 175 | total = labels.shape[0] 176 | correct = (predicted == labels).sum().item() 177 | 178 | return correct / total 179 | 180 | if __name__ == "__main__": 181 | f = open("./data.txt", 'rb') 182 | data = pickle.load(f) 183 | f.close() 184 | batch = ModelNet40CustomBatch([data]) 185 | model = KPCNN(cfg) 186 | hook = auto_diff.Hook("KPCNN") 187 | hook.hook_module(model) 188 | outputs = model(batch, cfg) 189 | -------------------------------------------------------------------------------- /networks/cls/pointcnn.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jittor as jt 3 | from jittor import nn 4 | from jittor.contrib import concat 5 | from jittor import init 6 | 7 | from misc.layers import Dense_Conv1d, Dense_Conv2d, RandPointCNN 8 | 9 | import time 10 | 11 | 12 | jt.flags.use_cuda = 1 13 | 14 | 15 | 16 | # C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2 17 | # (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep) 18 | # Abbreviated PointCNN constructor. 19 | 20 | AbbPointCNN = lambda a, b, c, d, e: RandPointCNN(a, b, 3, c, d, e) 21 | 22 | 23 | class PointCNNcls(nn.Module): 24 | def __init__(self, n_classes=40): 25 | super(PointCNNcls, self).__init__() 26 | self.pcnn1 = AbbPointCNN(3, 48, 8, 1, -1) 27 | self.pcnn2 = nn.Sequential( 28 | AbbPointCNN(48, 96, 12, 2, 384), 29 | AbbPointCNN(96, 192, 16, 2, 128), 30 | AbbPointCNN(192, 384, 16, 3, 128), 31 | ) 32 | 33 | self.fcn = nn.Sequential( 34 | Dense_Conv1d(384, 192), 35 | Dense_Conv1d(192, 128, drop_rate=0.5), 36 | Dense_Conv1d(128, n_classes, with_bn=False, activation=None) 37 | ) 38 | 39 | 40 | def execute(self, x, normal=None): 41 | if normal is None: 42 | x = (x, x) 43 | else : 44 | x = (x, normal) 45 | 46 | x = self.pcnn1(x) 47 | x = self.pcnn2(x)[1] # grab features 48 | 49 | x = x.permute(0, 2, 1) # b, dim, n 50 | logits = self.fcn(x) 51 | logits = jt.mean(logits, dim=2) 52 | return logits 53 | 54 | 55 | def main(): 56 | x_input = init.invariant_uniform([16, 1024, 3], dtype='float') 57 | x_ = x_input 58 | x = (x_, x_input) 59 | model = PointCNN() 60 | y = model(x) 61 | _ = y.data 62 | print (y.shape) 63 | 64 | if __name__ == '__main__': 65 | main() 66 | -------------------------------------------------------------------------------- /networks/cls/pointconv.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jittor as jt 3 | from jittor import nn 4 | from jittor.contrib import concat 5 | 6 | from misc.pointconv_utils import PointConvDensitySetAbstraction 7 | 8 | class PointConvDensityClsSsg(nn.Module): 9 | def __init__(self, n_classes = 40): 10 | super(PointConvDensityClsSsg, self).__init__() 11 | self.sa1 = PointConvDensitySetAbstraction(npoint=512, nsample=32, in_channel=3, mlp=[64, 64, 128], bandwidth = 0.1, group_all=False) 12 | self.sa2 = PointConvDensitySetAbstraction(npoint=128, nsample=64, in_channel=128 + 3, mlp=[128, 128, 256], bandwidth = 0.2, group_all=False) 13 | self.sa3 = PointConvDensitySetAbstraction(npoint=1, nsample=None, in_channel=256 + 3, mlp=[256, 512, 1024], bandwidth = 0.4, group_all=True) 14 | self.fc1 = nn.Linear(1024, 512) 15 | self.bn1 = nn.BatchNorm1d(512) 16 | self.drop1 = nn.Dropout(0.4) 17 | self.fc2 = nn.Linear(512, 256) 18 | self.bn2 = nn.BatchNorm1d(256) 19 | self.drop2 = nn.Dropout(0.4) 20 | self.fc3 = nn.Linear(256, n_classes) 21 | self.relu = nn.ReLU() 22 | 23 | def execute(self, xyz): 24 | xyz = xyz.permute(0, 2, 1) 25 | B, _, _ = xyz.shape 26 | 27 | l1_xyz, l1_points = self.sa1(xyz, None) 28 | l2_xyz, l2_points = self.sa2(l1_xyz, l1_points) 29 | l3_xyz, l3_points = self.sa3(l2_xyz, l2_points) 30 | x = l3_points.reshape(B, 1024) # to reshape 31 | x = self.drop1(self.relu(self.bn1(self.fc1(x)))) 32 | x = self.drop2(self.relu(self.bn2(self.fc2(x)))) 33 | x = self.fc3(x) 34 | return x 35 | 36 | 37 | -------------------------------------------------------------------------------- /networks/cls/pointnet.py: -------------------------------------------------------------------------------- 1 | import jittor as jt 2 | from jittor import nn 3 | from jittor import init 4 | from jittor.contrib import concat 5 | import numpy as np 6 | 7 | 8 | 9 | class PointNet(nn.Module): 10 | def __init__(self, output_channels=40): 11 | super(PointNet, self).__init__() 12 | self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False) 13 | self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 14 | self.conv3 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 15 | self.conv4 = nn.Conv1d(64, 128, kernel_size=1, bias=False) 16 | self.conv5 = nn.Conv1d(128, 1024, kernel_size=1, bias=False) 17 | self.bn1 = nn.BatchNorm1d(64) 18 | self.bn2 = nn.BatchNorm1d(64) 19 | self.bn3 = nn.BatchNorm1d(64) 20 | self.bn4 = nn.BatchNorm1d(128) 21 | self.bn5 = nn.BatchNorm1d(1024) 22 | self.linear1 = nn.Linear(1024, 512, bias=False) 23 | self.bn6 = nn.BatchNorm1d(512) 24 | self.dp1 = nn.Dropout(0.5) 25 | self.linear2 = nn.Linear(512, output_channels) 26 | self.relu = nn.ReLU() 27 | 28 | def execute(self, x): 29 | b, c, n = x.size() 30 | x = self.relu(self.bn1(self.conv1(x))) 31 | x = self.relu(self.bn2(self.conv2(x))) 32 | x = self.relu(self.bn3(self.conv3(x))) 33 | x = self.relu(self.bn4(self.conv4(x))) 34 | x = self.relu(self.bn5(self.conv5(x))) 35 | x = jt.max(x, 2) 36 | x = x.reshape(b, -1) 37 | x = self.relu(self.bn6(self.linear1(x))) 38 | x = self.dp1(x) 39 | x = self.linear2(x) 40 | return x 41 | -------------------------------------------------------------------------------- /networks/cls/pointnet2.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | 3 | import jittor as jt 4 | import jittor.nn as nn 5 | from jittor import init 6 | 7 | from misc.ops import FurthestPointSampler 8 | from misc.ops import BallQueryGrouper 9 | from misc.ops import GroupAll 10 | 11 | class PointNetModuleBase(nn.Module): 12 | def __init__(self): 13 | self.n_points = None 14 | self.sampler = None 15 | self.groupers = None 16 | self.mlps = None 17 | 18 | def build_mlps(self, mlp_spec: List[int], use_xyz: bool=True, 19 | bn: bool = True) -> nn.Sequential: 20 | layers = [] 21 | 22 | if use_xyz: 23 | mlp_spec[0] += 3 24 | 25 | for i in range(1, len(mlp_spec)): 26 | layers.append(nn.Conv(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn)) 27 | if bn: 28 | layers.append(nn.BatchNorm(mlp_spec[i])) 29 | layers.append(nn.ReLU()) 30 | 31 | return nn.Sequential(*layers) 32 | 33 | def execute(self, xyz: jt.Var, feature: Optional[jt.Var]) -> Tuple[jt.Var, jt.Var]: 34 | ''' 35 | Parameters 36 | ---------- 37 | xyz: jt.Var, (B, N, 3) 38 | feature: jt.Var, (B, N, C) 39 | 40 | Returns 41 | ------- 42 | new_xyz: jt.Var, (B, n_points, 3) 43 | new_feature: jt.Var, (B, n_points, C') 44 | ''' 45 | new_xyz = self.sampler(xyz) if self.n_points is not None else None 46 | 47 | new_feature_list = [] 48 | l = len (self.groupers) 49 | for i in range (l): 50 | grouper = self.groupers[i] 51 | new_feature = grouper(new_xyz, xyz, feature) 52 | # [B, n_points, n_samples, C] -> [B, C, n_points, n_samples] 53 | new_feature = new_feature.transpose(0, 3, 1, 2) 54 | new_feature = self.mlps[i](new_feature) 55 | # [B, C, n_points, n_samples] -> [B, n_points, n_samples, C] 56 | new_feature = new_feature.transpose(0, 2, 3, 1) 57 | new_feature = new_feature.argmax(dim=2)[1] 58 | 59 | new_feature_list.append(new_feature) 60 | 61 | new_feature = jt.contrib.concat(new_feature_list, dim=-1) 62 | return new_xyz, new_feature 63 | 64 | 65 | class PointnetModule(PointNetModuleBase): 66 | def __init__(self, mlp: List[int], n_points=None, radius=None, 67 | n_samples=None, bn=True, use_xyz=True): 68 | super().__init__() 69 | 70 | self.n_points = n_points 71 | 72 | self.groupers = nn.ModuleList() 73 | if self.n_points is not None: 74 | self.sampler = FurthestPointSampler(n_points) 75 | self.groupers.append(BallQueryGrouper(radius, n_samples, use_xyz)) 76 | else: 77 | self.groupers.append(GroupAll(use_xyz)) 78 | 79 | self.mlps = nn.ModuleList() 80 | self.mlps.append(self.build_mlps(mlp, use_xyz)) 81 | 82 | 83 | class PointnetModuleMSG(PointNetModuleBase): 84 | def __init__(self, n_points: int, radius: List[float], n_samples: List[int], 85 | mlps: List[List[int]], bn=True, use_xyz=True): 86 | super().__init__() 87 | 88 | self.n_points = n_points 89 | self.sampler = FurthestPointSampler(n_points) 90 | 91 | self.groupers = nn.ModuleList() 92 | for r, s in zip(radius, n_samples): 93 | self.groupers.append(BallQueryGrouper(r, s, use_xyz)) 94 | 95 | self.mlps = nn.ModuleList() 96 | for mlp in mlps.layers.items(): 97 | self.mlps.append(self.build_mlps(mlp, use_xyz)) 98 | 99 | 100 | class PointNet2_cls(nn.Module): 101 | def __init__(self, n_classes=40, use_xyz=True): 102 | super().__init__() 103 | 104 | self.n_classes = n_classes 105 | self.use_xyz = use_xyz 106 | self.build_model() 107 | 108 | def build_model(self): 109 | self.pointnet_modules = nn.ModuleList() 110 | 111 | self.pointnet_modules.append( 112 | PointnetModule( 113 | n_points=512, 114 | radius=0.2, 115 | n_samples=64, 116 | mlp=[3, 64, 64, 128], 117 | use_xyz=self.use_xyz, 118 | ) 119 | ) 120 | 121 | self.pointnet_modules.append( 122 | PointnetModule( 123 | n_points=128, 124 | radius=0.4, 125 | n_samples=64, 126 | mlp=[128, 128, 128, 256], 127 | use_xyz=self.use_xyz, 128 | ) 129 | ) 130 | 131 | self.pointnet_modules.append( 132 | PointnetModule( 133 | mlp=[256, 256, 512, 1024], 134 | use_xyz=self.use_xyz, 135 | ) 136 | ) 137 | 138 | self.fc_layer = nn.Sequential( 139 | nn.Linear(1024, 512, bias=False), 140 | nn.BatchNorm1d(512), 141 | nn.ReLU(), 142 | nn.Linear(512, 256, bias=False), 143 | nn.BatchNorm1d(256), 144 | nn.ReLU(), 145 | nn.Dropout(0.5), 146 | nn.Linear(256, self.n_classes), 147 | ) 148 | 149 | def execute(self, xyz, feature): 150 | l = len(self.pointnet_modules) 151 | for i in range (l): 152 | module = self.pointnet_modules[i] 153 | xyz, feature = module(xyz, feature) 154 | # for module in self.pointnet_modules: 155 | # xyz, feature = module(xyz, feature) 156 | 157 | feature = feature.squeeze(dim=1) 158 | return self.fc_layer(feature) 159 | 160 | 161 | class PointNetMSG(PointNet2_cls): 162 | def build_model(self): 163 | super().build_model() 164 | 165 | self.pointnet_modules = nn.ModuleList() 166 | self.pointnet_modules.append( 167 | PointnetModuleMSG( 168 | n_points=512, 169 | radius=[0.1, 0.2, 0.4], 170 | n_samples=[16, 32, 128], 171 | mlps=[[3, 32, 32, 64], [3, 64, 64, 128], [3, 64, 96, 128]], 172 | use_xyz=self.use_xyz, 173 | ) 174 | ) 175 | 176 | input_channels = 64 + 128 + 128 177 | self.pointnet_modules.append( 178 | PointnetModuleMSG( 179 | n_points=128, 180 | radius=[0.2, 0.4, 0.8], 181 | n_samples=[32, 64, 128], 182 | mlps=[ 183 | [input_channels, 64, 64, 128], 184 | [input_channels, 128, 128, 256], 185 | [input_channels, 128, 128, 256], 186 | ], 187 | use_xyz=self.use_xyz, 188 | ) 189 | ) 190 | 191 | self.pointnet_modules.append( 192 | PointnetModule( 193 | mlp=[128 + 256 + 256, 256, 512, 1024], 194 | use_xyz=self.use_xyz, 195 | ) 196 | ) 197 | 198 | def main(): 199 | model = PointNet(n_classes=40) 200 | input_point = init.gauss([2, 1024, 3], 'float', mean=0.0) 201 | input_feature = init.gauss([2, 1024, 3], 'float', mean=0.0) 202 | print (input_point.shape) 203 | print (input_feature.shape) 204 | outputs = model(input_point, input_feature) 205 | print (outputs.shape) 206 | 207 | 208 | if __name__ == '__main__' : 209 | main() 210 | -------------------------------------------------------------------------------- /networks/seg/dgcnn_partseg.py: -------------------------------------------------------------------------------- 1 | import jittor as jt 2 | from jittor import nn 3 | from jittor.contrib import concat 4 | 5 | import numpy as np 6 | import math 7 | 8 | from misc.ops import KNN 9 | 10 | 11 | def get_graph_feature(x, knn=None, k=None, idx=None): 12 | batch_size = x.shape[0] 13 | num_points = x.shape[2] 14 | x = x.reshape(batch_size, -1, num_points) 15 | if idx is None: 16 | idx = knn(x,x) # (batch_size, num_points, k) 17 | idx = idx.permute(0, 2, 1) 18 | idx_base = jt.array(np.arange(0, batch_size)).reshape(-1, 1, 1)*num_points 19 | 20 | idx = idx + idx_base 21 | 22 | idx = idx.reshape(-1) 23 | 24 | _, num_dims, _ = x.shape 25 | 26 | x = x.transpose(0, 2, 1) # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 27 | feature = x.reshape(batch_size*num_points, -1)[idx, :] 28 | feature = feature.reshape(batch_size, num_points, k, num_dims) 29 | x = x.reshape(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 30 | 31 | feature = concat((feature-x, x), dim=3).transpose(0, 3, 1, 2) 32 | return feature 33 | 34 | 35 | class DGCNN_partseg(nn.Module): 36 | def __init__(self, part_num): 37 | super(DGCNN_partseg, self).__init__() 38 | self.seg_num_all = part_num 39 | self.k = 40 40 | self.knn = KNN(self.k) 41 | self.bn1 = nn.BatchNorm2d(64) 42 | self.bn2 = nn.BatchNorm2d(64) 43 | self.bn3 = nn.BatchNorm2d(64) 44 | self.bn4 = nn.BatchNorm2d(64) 45 | self.bn5 = nn.BatchNorm2d(64) 46 | self.bn6 = nn.BatchNorm1d(1024) 47 | self.bn7 = nn.BatchNorm1d(64) 48 | self.bn8 = nn.BatchNorm1d(256) 49 | self.bn9 = nn.BatchNorm1d(256) 50 | self.bn10 = nn.BatchNorm1d(128) 51 | 52 | self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), 53 | self.bn1, 54 | nn.LeakyReLU(scale=0.2)) 55 | self.conv2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False), 56 | self.bn2, 57 | nn.LeakyReLU(scale=0.2)) 58 | self.conv3 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), 59 | self.bn3, 60 | nn.LeakyReLU(scale=0.2)) 61 | self.conv4 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False), 62 | self.bn4, 63 | nn.LeakyReLU(scale=0.2)) 64 | self.conv5 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), 65 | self.bn5, 66 | nn.LeakyReLU(scale=0.2)) 67 | self.conv6 = nn.Sequential(nn.Conv1d(192, 1024, kernel_size=1, bias=False), 68 | self.bn6, 69 | nn.LeakyReLU(scale=0.2)) 70 | self.conv7 = nn.Sequential(nn.Conv1d(16, 64, kernel_size=1, bias=False), 71 | self.bn7, 72 | nn.LeakyReLU(scale=0.2)) 73 | self.conv8 = nn.Sequential(nn.Conv1d(1280, 256, kernel_size=1, bias=False), 74 | self.bn8, 75 | nn.LeakyReLU(scale=0.2)) 76 | self.dp1 = nn.Dropout(p=0.5) 77 | self.conv9 = nn.Sequential(nn.Conv1d(256, 256, kernel_size=1, bias=False), 78 | self.bn9, 79 | nn.LeakyReLU(scale=0.2)) 80 | self.dp2 = nn.Dropout(p=0.5) 81 | self.conv10 = nn.Sequential(nn.Conv1d(256, 128, kernel_size=1, bias=False), 82 | self.bn10, 83 | nn.LeakyReLU(scale=0.2)) 84 | self.conv11 = nn.Conv1d(128, self.seg_num_all, kernel_size=1, bias=False) 85 | 86 | 87 | def execute(self, x, l): 88 | batch_size = x.size(0) 89 | num_points = x.size(2) 90 | # print (x.size()) 91 | 92 | # x0 = get_graph_feature(x, k=self.k) # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k) 93 | 94 | x = get_graph_feature(x, knn=self.knn, k=self.k) 95 | x = self.conv1(x) # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k) 96 | x = self.conv2(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k) 97 | x1 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) 98 | 99 | x = get_graph_feature(x1, knn=self.knn, k=self.k) 100 | x = self.conv3(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k) 101 | x = self.conv4(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k) 102 | x2 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) 103 | 104 | x = get_graph_feature(x2, knn=self.knn, k=self.k) 105 | x = self.conv5(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k) 106 | x3 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) 107 | 108 | x = concat((x1, x2, x3), dim=1) # (batch_size, 64*3, num_points) 109 | 110 | x = self.conv6(x) # (batch_size, 64*3, num_points) -> (batch_size, emb_dims, num_points) 111 | x = x.max(dim=-1, keepdims=True) # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims, 1) 112 | 113 | l = l.view(batch_size, -1, 1) # (batch_size, num_categoties, 1) 114 | l = self.conv7(l) # (batch_size, num_categoties, 1) -> (batch_size, 64, 1) 115 | 116 | x = concat((x, l), dim=1) # (batch_size, 1088, 1) 117 | x = x.repeat(1, 1, num_points) # (batch_size, 1088, num_points) 118 | 119 | x = concat((x, x1, x2, x3), dim=1) # (batch_size, 1088+64*3, num_points) 120 | 121 | x = self.conv8(x) # (batch_size, 1088+64*3, num_points) -> (batch_size, 256, num_points) 122 | x = self.dp1(x) 123 | x = self.conv9(x) # (batch_size, 256, num_points) -> (batch_size, 256, num_points) 124 | x = self.dp2(x) 125 | x = self.conv10(x) # (batch_size, 256, num_points) -> (batch_size, 128, num_points) 126 | x = self.conv11(x) # (batch_size, 256, num_points) -> (batch_size, seg_num_all, num_points) 127 | 128 | return x 129 | -------------------------------------------------------------------------------- /networks/seg/pointcnn_partseg.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jittor as jt 3 | from jittor import nn 4 | from jittor.contrib import concat 5 | from jittor import init 6 | from misc.layers import Dense_Conv1d, Dense_Conv2d, RandPointCNN, RandPointCNN_Decoder 7 | import time 8 | jt.flags.use_cuda = 1 9 | 10 | 11 | # (C_in : int, C_out : int, dims : int, K : int, D : int, P : int) -> None: 12 | 13 | EncoderCNN = lambda a, b, c, d, e: RandPointCNN(a, b, 3, c, d, e) 14 | 15 | DecoderCNN = lambda a, b, last_c, c, d, e: RandPointCNN_Decoder(a, b, last_c, 3, c, d, e) 16 | 17 | class PointCNN_partseg(nn.Module): 18 | def __init__(self, part_num=50): 19 | super(PointCNN_partseg, self).__init__() 20 | self.encoder_0 = EncoderCNN(3, 256, 8, 1, -1) 21 | self.encoder_1 = EncoderCNN(256, 256, 12, 1, 768) 22 | self.encoder_2 = EncoderCNN(256, 512, 16, 1, 384) 23 | self.encoder_3 = EncoderCNN(512, 1024, 16, 1, 128) 24 | 25 | 26 | self.decoder_0 = DecoderCNN(1024, 1024, 1024, 16, 1, 128) 27 | self.decoder_1 = DecoderCNN(1024, 512, 512, 16, 1, 385) 28 | self.decoder_2 = DecoderCNN(512, 256, 256, 12, 1, 768) 29 | self.decoder_3 = DecoderCNN(256, part_num, 256, 8, 1, 2048) 30 | # self.decoder_4 = DecoderCNN(256, part_num, 8, 4, 2048) 31 | 32 | 33 | 34 | def execute(self, x, normal=None): 35 | x = (x, x) 36 | # jt.sync_all(True) 37 | # start_time = time.time() 38 | x_0 = self.encoder_0(x) 39 | x_1 = self.encoder_1(x_0) 40 | x_2 = self.encoder_2(x_1) 41 | x_3 = self.encoder_3(x_2) 42 | x_3 = self.decoder_0(x_3, x_3) 43 | x_2 = self.decoder_1(x_3, x_2) 44 | x_1 = self.decoder_2(x_2, x_1) 45 | x_0 = self.decoder_3(x_1, x_0) 46 | 47 | return x_0[1].permute(0, 2, 1) 48 | 49 | 50 | def main(): 51 | x_input = init.invariant_uniform([16, 1024, 3], dtype='float') 52 | x_ = x_input 53 | x = (x_, x_input) 54 | model = PointCNN() 55 | y = model(x) 56 | _ = y.data 57 | print (y.shape) 58 | 59 | if __name__ == '__main__': 60 | main() 61 | -------------------------------------------------------------------------------- /networks/seg/pointconv_partseg.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import jittor as jt 4 | from jittor import nn 5 | from jittor.contrib import concat 6 | 7 | from misc.pointconv_utils import PointConvDensitySetAbstraction, PointConvDensitySetInterpolation 8 | 9 | class PointConvDensity_partseg(nn.Module): 10 | def __init__(self, part_num=50): 11 | super(PointConvDensity_partseg, self).__init__() 12 | self.part_num = part_num 13 | 14 | self.sa0 = PointConvDensitySetAbstraction(npoint=1024, nsample=32, in_channel=3, mlp=[32,32,64], bandwidth = 0.1, group_all=False) 15 | self.sa1 = PointConvDensitySetAbstraction(npoint=256, nsample=32, in_channel=64 + 3, mlp=[64,64,128], bandwidth = 0.2, group_all=False) 16 | self.sa2 = PointConvDensitySetAbstraction(npoint=64, nsample=32, in_channel=128 + 3, mlp=[128,128,256], bandwidth = 0.4, group_all=False) 17 | self.sa3 = PointConvDensitySetAbstraction(npoint=36, nsample=32, in_channel=256 + 3, mlp=[256,256,512], bandwidth = 0.8, group_all=False) 18 | 19 | 20 | # TODO upsample 21 | # upsampling 22 | # def __init__(self, nsample, in_channel, mlp, bandwidth): 23 | 24 | self.in0 = PointConvDensitySetInterpolation(nsample=16, in_channel=512 + 3, mlp=[512,512], bandwidth=0.8) 25 | self.in1 = PointConvDensitySetInterpolation(nsample=16, in_channel=512 + 3, mlp=[256,256], bandwidth=0.4) 26 | self.in2 = PointConvDensitySetInterpolation(nsample=16, in_channel=256 + 3, mlp=[128,128], bandwidth=0.2) 27 | self.in3 = PointConvDensitySetInterpolation(nsample=16, in_channel=128 + 3, mlp=[128,128, 128], bandwidth=0.1) 28 | 29 | # self.fp0 = PointConvDensitySetAbstraction(npoint=1024, nsample=32, in_channel=3, mlp=[32,32,64], bandwidth = 0.1, group_all=False) 30 | # self.fp1 = PointConvDensitySetAbstraction(npoint=256, nsample=32, in_channel=64 + 3, mlp=[64,64,128], bandwidth = 0.2, group_all=False) 31 | # self.fp2 = PointConvDensitySetAbstraction(npoint=64, nsample=32, in_channel=128 + 3, mlp=[128,128,256], bandwidth = 0.4, group_all=False) 32 | # self.fp3 = PointConvDensitySetAbstraction(npoint=36, nsample=32, in_channel=256 + 3, mlp=[256,256,512], bandwidth = 0.8, group_all=False) 33 | 34 | self.fc1 = nn.Conv1d(128, 128, 1) 35 | self.bn1 = nn.BatchNorm1d(128) 36 | self.drop1 = nn.Dropout(0.4) 37 | self.fc3 = nn.Conv1d(128, self.part_num, 1) 38 | self.relu = nn.ReLU() 39 | 40 | def execute(self, xyz, cls_label): 41 | xyz = xyz.permute(0, 2, 1) 42 | B, _, _ = xyz.shape 43 | 44 | l1_xyz, l1_points = self.sa0(xyz, None) 45 | l2_xyz, l2_points = self.sa1(l1_xyz, l1_points) 46 | l3_xyz, l3_points = self.sa2(l2_xyz, l2_points) 47 | l4_xyz, l4_points = self.sa3(l3_xyz, l3_points) 48 | # print ('after encoder shape =',l4_xyz.shape, l4_points.shape) 49 | # def execute(self, xyz1, xyz2, points1, points2): 50 | 51 | l3_points = self.in0(l3_xyz, l4_xyz, l3_points, l4_points) 52 | l2_points = self.in1(l2_xyz, l3_xyz, l2_points, l3_points) 53 | l1_points = self.in2(l1_xyz, l2_xyz, l1_points, l2_points) 54 | l0_points = self.in3(xyz, l1_xyz, xyz, l1_points) 55 | 56 | # print ('after decoder shape =', l0_points.shape) 57 | 58 | x = self.drop1(self.relu(self.bn1(self.fc1(l0_points)))) 59 | x = self.fc3(x) 60 | x = x.permute(0, 2, 1) 61 | return x 62 | 63 | 64 | -------------------------------------------------------------------------------- /networks/seg/pointnet2_partseg.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | 3 | import jittor as jt 4 | import jittor.nn as nn 5 | from jittor import init 6 | from jittor.contrib import concat 7 | 8 | 9 | from misc.ops import FurthestPointSampler 10 | from misc.ops import BallQueryGrouper 11 | from misc.ops import GroupAll 12 | from misc.ops import PointNetFeaturePropagation 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | class PointNetModuleBase(nn.Module): 21 | def __init__(self): 22 | self.n_points = None 23 | self.sampler = None 24 | self.groupers = None 25 | self.mlps = None 26 | 27 | def build_mlps(self, mlp_spec: List[int], use_xyz: bool=True, 28 | bn: bool = True) -> nn.Sequential: 29 | layers = [] 30 | 31 | if use_xyz: 32 | mlp_spec[0] += 3 33 | 34 | for i in range(1, len(mlp_spec)): 35 | layers.append(nn.Conv(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn)) 36 | if bn: 37 | layers.append(nn.BatchNorm(mlp_spec[i])) 38 | layers.append(nn.ReLU()) 39 | 40 | return nn.Sequential(*layers) 41 | 42 | def execute(self, xyz: jt.Var, feature: Optional[jt.Var]) -> Tuple[jt.Var, jt.Var]: 43 | ''' 44 | Parameters 45 | ---------- 46 | xyz: jt.Var, (B, N, 3) 47 | feature: jt.Var, (B, N, C) 48 | 49 | Returns 50 | ------- 51 | new_xyz: jt.Var, (B, n_points, 3) 52 | new_feature: jt.Var, (B, n_points, C') 53 | ''' 54 | B, _, _ = xyz.shape 55 | new_xyz = self.sampler(xyz) if self.n_points is not None else jt.zeros((B, 1, 3)) 56 | 57 | new_feature_list = [] 58 | # print (self.groupers) 59 | # lens = 60 | for i, grouper in self.groupers.layers.items(): 61 | new_feature = grouper(new_xyz, xyz, feature) 62 | # [B, n_points, n_samples, C] -> [B, C, n_points, n_samples] 63 | new_feature = new_feature.transpose(0, 3, 1, 2) 64 | new_feature = self.mlps[i](new_feature) 65 | # [B, C, n_points, n_samples] -> [B, n_points, n_samples, C] 66 | new_feature = new_feature.transpose(0, 2, 3, 1) 67 | new_feature = new_feature.argmax(dim=2)[1] 68 | 69 | new_feature_list.append(new_feature) 70 | 71 | new_feature = jt.contrib.concat(new_feature_list, dim=-1) 72 | return new_xyz, new_feature 73 | 74 | 75 | class PointnetModule(PointNetModuleBase): 76 | def __init__(self, mlp: List[int], n_points=None, radius=None, 77 | n_samples=None, bn=True, use_xyz=True): 78 | super().__init__() 79 | 80 | self.n_points = n_points 81 | 82 | self.groupers = nn.ModuleList() 83 | if self.n_points is not None: 84 | self.sampler = FurthestPointSampler(n_points) 85 | self.groupers.append(BallQueryGrouper(radius, n_samples, use_xyz)) 86 | else: 87 | self.groupers.append(GroupAll(use_xyz)) 88 | 89 | self.mlps = nn.ModuleList() 90 | self.mlps.append(self.build_mlps(mlp, use_xyz)) 91 | 92 | 93 | class PointnetModuleMSG(PointNetModuleBase): 94 | def __init__(self, n_points: int, radius: List[float], n_samples: List[int], 95 | mlps: List[List[int]], bn=True, use_xyz=True): 96 | super().__init__() 97 | 98 | self.n_points = n_points 99 | self.sampler = FurthestPointSampler(n_points) 100 | 101 | self.groupers = nn.ModuleList() 102 | for r, s in zip(radius, n_samples): 103 | self.groupers.append(BallQueryGrouper(r, s, use_xyz)) 104 | 105 | self.mlps = nn.ModuleList() 106 | for mlp in mlps: 107 | self.mlps.append(self.build_mlps(mlp, use_xyz)) 108 | 109 | 110 | class PointNet2_partseg(nn.Module): 111 | def __init__(self, part_num=50, use_xyz=True): 112 | super().__init__() 113 | self.part_num = part_num 114 | self.use_xyz = use_xyz 115 | self.build_model() 116 | 117 | def build_model(self): 118 | self.pointnet_modules = nn.ModuleList() 119 | self.pointnet_modules.append( 120 | PointnetModule( 121 | n_points=512, 122 | radius=0.2, 123 | n_samples=64, 124 | mlp=[3, 64, 64, 128], 125 | use_xyz=self.use_xyz, 126 | ) 127 | ) 128 | 129 | self.pointnet_modules.append( 130 | PointnetModule( 131 | n_points=128, 132 | radius=0.4, 133 | n_samples=64, 134 | mlp=[128, 128, 128, 256], 135 | use_xyz=self.use_xyz, 136 | ) 137 | ) 138 | 139 | self.pointnet_modules.append( 140 | PointnetModule( 141 | mlp=[256, 256, 512, 1024], 142 | use_xyz=self.use_xyz, 143 | ) 144 | ) 145 | 146 | self.fp3 = PointNetFeaturePropagation(in_channel=1280, mlp=[256, 256]) 147 | self.fp2 = PointNetFeaturePropagation(in_channel=384, mlp=[256, 128]) 148 | self.fp1 = PointNetFeaturePropagation(in_channel=128+16+6, mlp=[128, 128, 128]) 149 | 150 | 151 | self.fc_layer = nn.Sequential( 152 | nn.Conv1d(128, 128, 1), 153 | nn.BatchNorm1d(128), 154 | nn.Dropout(0.5), 155 | nn.Conv1d(128, self.part_num, 1) 156 | ) 157 | 158 | def execute(self, xyz, feature, cls_label): 159 | # for module in self.pointnet_modules: 160 | # xyz, feature = module(xyz, feature) 161 | 162 | B, N, _ = xyz.shape 163 | l1_xyz, l1_feature = self.pointnet_modules[0](xyz, feature) 164 | l2_xyz, l2_feature = self.pointnet_modules[1](l1_xyz, l1_feature) 165 | l3_xyz, l3_feature = self.pointnet_modules[2](l2_xyz, l2_feature) 166 | # print ('before interpolate shape') 167 | # print (l2_xyz.shape, l2_feature.shape, l3_xyz.shape, l3_feature.shape) 168 | l2_feature = self.fp3(l2_xyz, l3_xyz, l2_feature, l3_feature) 169 | l1_feature = self.fp2(l1_xyz, l2_xyz, l1_feature, l2_feature) 170 | cls_label_one_hot = cls_label.view(B,16,1).repeat(1,1,N).permute(0, 2, 1) 171 | # print ('before concat size ') 172 | # print (cls_label_one_hot.size(),xyz.size(),feature.size()) 173 | feature = self.fp1(xyz, l1_xyz, concat([cls_label_one_hot,xyz,feature],2), l1_feature) 174 | feature = feature.permute(0, 2, 1) 175 | # print (feature.shape) 176 | return self.fc_layer(feature) 177 | 178 | 179 | class PointNetMSG(PointNet2_partseg): 180 | def build_model(self): 181 | super().build_model() 182 | 183 | self.pointnet_modules = nn.ModuleList() 184 | self.pointnet_modules.append( 185 | PointnetModuleMSG( 186 | n_points=512, 187 | radius=[0.1, 0.2, 0.4], 188 | n_samples=[16, 32, 128], 189 | mlps=[[3, 32, 32, 64], [3, 64, 64, 128], [3, 64, 96, 128]], 190 | use_xyz=self.use_xyz, 191 | ) 192 | ) 193 | 194 | input_channels = 64 + 128 + 128 195 | self.pointnet_modules.append( 196 | PointnetModuleMSG( 197 | n_points=128, 198 | radius=[0.2, 0.4, 0.8], 199 | n_samples=[32, 64, 128], 200 | mlps=[ 201 | [input_channels, 64, 64, 128], 202 | [input_channels, 128, 128, 256], 203 | [input_channels, 128, 128, 256], 204 | ], 205 | use_xyz=self.use_xyz, 206 | ) 207 | ) 208 | 209 | self.pointnet_modules.append( 210 | PointnetModule( 211 | mlp=[128 + 256 + 256, 256, 512, 1024], 212 | use_xyz=self.use_xyz, 213 | ) 214 | ) 215 | 216 | def main(): 217 | model = PointNet2_partseg() 218 | input_point = init.gauss([2, 1024, 3], 'float', mean=0.0) 219 | input_feature = init.gauss([2, 1024, 3], 'float', mean=0.0) 220 | cls_label = init.gauss([2, 16], 'float', mean=0.0) 221 | 222 | print (input_point.shape) 223 | print (input_feature.shape) 224 | print (cls_label.shape) 225 | outputs = model(input_point, input_feature, cls_label) 226 | print (outputs.shape) 227 | 228 | 229 | if __name__ == '__main__' : 230 | main() 231 | -------------------------------------------------------------------------------- /networks/seg/pointnet_partseg.py: -------------------------------------------------------------------------------- 1 | import jittor as jt 2 | from jittor import nn 3 | from jittor.contrib import concat 4 | 5 | import numpy as np 6 | import math 7 | 8 | from misc.layers import STN3d, STNkd 9 | 10 | 11 | 12 | 13 | 14 | class PointNet_partseg(nn.Module): 15 | def __init__(self, part_num=50): 16 | super(PointNet_partseg, self).__init__() 17 | self.part_num = part_num 18 | self.stn = STN3d() 19 | self.conv1 = nn.Conv1d(3, 64, 1) 20 | self.conv2 = nn.Conv1d(64, 128, 1) 21 | self.conv3 = nn.Conv1d(128, 128, 1) 22 | self.conv4 = nn.Conv1d(128, 512, 1) 23 | self.conv5 = nn.Conv1d(512, 2048, 1) 24 | self.bn1 = nn.BatchNorm1d(64) 25 | self.bn2 = nn.BatchNorm1d(128) 26 | self.bn3 = nn.BatchNorm1d(128) 27 | self.bn4 = nn.BatchNorm1d(512) 28 | self.bn5 = nn.BatchNorm1d(2048) 29 | self.fstn = STNkd(k=128) 30 | self.convs1 = nn.Conv1d(4944, 256, 1) 31 | self.convs2 = nn.Conv1d(256, 256, 1) 32 | self.convs3 = nn.Conv1d(256, 128, 1) 33 | self.convs4 = nn.Conv1d(128, part_num, 1) 34 | self.bns1 = nn.BatchNorm1d(256) 35 | self.bns2 = nn.BatchNorm1d(256) 36 | self.bns3 = nn.BatchNorm1d(128) 37 | self.relu = nn.ReLU() 38 | def execute(self, point_cloud, label): 39 | B, D, N = point_cloud.size() 40 | trans = self.stn(point_cloud) 41 | point_cloud = point_cloud.transpose(0, 2, 1) 42 | point_cloud = nn.bmm(point_cloud, trans) 43 | 44 | point_cloud = point_cloud.transpose(0, 2, 1) 45 | 46 | out1 = self.relu(self.bn1(self.conv1(point_cloud))) 47 | out2 = self.relu(self.bn2(self.conv2(out1))) 48 | out3 = self.relu(self.bn3(self.conv3(out2))) 49 | 50 | trans_feat = self.fstn(out3) 51 | x = out3.transpose(0, 2, 1) 52 | net_transformed = nn.bmm(x, trans_feat) 53 | net_transformed = net_transformed.transpose(0, 2, 1) 54 | 55 | out4 = self.relu(self.bn4(self.conv4(net_transformed))) 56 | out5 = self.bn5(self.conv5(out4)) 57 | out_max = jt.argmax(out5, 2, keepdims=True)[1] 58 | out_max = out_max.view(-1, 2048) 59 | 60 | out_max = concat((out_max, label),1) 61 | expand = out_max.view(-1, 2048+16, 1).repeat(1, 1, N) 62 | concat_feature = concat([expand, out1, out2, out3, out4, out5], 1) 63 | net = self.relu(self.bns1(self.convs1(concat_feature))) 64 | net = self.relu(self.bns2(self.convs2(net))) 65 | net = self.relu(self.bns3(self.convs3(net))) 66 | net = self.convs4(net) 67 | return net 68 | 69 | if __name__ == '__main__': 70 | from jittor import init 71 | x_input = init.invariant_uniform([16, 3, 1024], dtype='float') 72 | cls_input = init.invariant_uniform([16, 16], dtype='float') 73 | model = PointNet_partseg() 74 | out = model(x_input, cls_input) 75 | print (out.size()) 76 | 77 | -------------------------------------------------------------------------------- /run_cls.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train_cls.py --model pointnet 2 | -------------------------------------------------------------------------------- /run_partseg.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train_partseg.py --model pointnet 2 | -------------------------------------------------------------------------------- /train_cls.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tqdm import tqdm 3 | import jittor as jt 4 | from jittor import nn 5 | from jittor.contrib import concat 6 | import torch 7 | jt.flags.use_cuda = 1 8 | 9 | from networks.cls.pointnet2 import PointNet2_cls 10 | from networks.cls.pointnet import PointNet as PointNet_cls 11 | from networks.cls.dgcnn import DGCNN 12 | from networks.cls.pointcnn import PointCNNcls 13 | from networks.cls.pointconv import PointConvDensityClsSsg 14 | from networks.cls.kpconv import KPCNN 15 | import math 16 | from os.path import exists, join 17 | from data_utils.modelnet40_loader import ModelNet40 18 | from data_utils.kpconv_loader import KPConvLoader 19 | from misc.utils import LRScheduler 20 | import argparse 21 | from datasets.ModelNet40 import ModelNet40Dataset, ModelNet40Sampler, ModelNet40CustomBatch, Modelnet40Config 22 | 23 | import time 24 | import pickle 25 | from jittor_utils import auto_diff 26 | 27 | def freeze_random_seed(): 28 | np.random.seed(0) 29 | 30 | 31 | def soft_cross_entropy_loss(output, target, smoothing=True): 32 | ''' Calculate cross entropy loss, apply label smoothing if needed. ''' 33 | 34 | target = target.view(-1) 35 | softmax = nn.Softmax(dim=1) 36 | if smoothing: 37 | eps = 0.2 38 | b, n_class = output.shape 39 | 40 | one_hot = jt.zeros(output.shape) 41 | for i in range (b): 42 | one_hot[i, target[i].data] = 1 43 | 44 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) 45 | # print (one_hot[0].data) 46 | log_prb = jt.log(softmax(output)) 47 | loss = -(one_hot * log_prb).sum(dim=1).mean() 48 | else: 49 | loss = nn.cross_entropy_loss(output, target) 50 | 51 | return loss 52 | 53 | 54 | def train(net, optimizer, epoch, dataloader, args): 55 | net.train() 56 | 57 | pbar = tqdm(dataloader, desc=f'Epoch {epoch} [TRAIN]') 58 | for pts, normals, labels in pbar: 59 | 60 | # #output = net(pts, normals) 61 | 62 | 63 | if args.model == 'pointnet' or args.model == 'dgcnn' : 64 | pts = pts.transpose(0, 2, 1) 65 | 66 | if args.model == 'pointnet2': 67 | output = net(pts, normals) 68 | else : 69 | output = net(pts) 70 | 71 | loss = soft_cross_entropy_loss(output, labels) 72 | optimizer.step(loss) 73 | pred = np.argmax(output.data, axis=1) 74 | acc = np.mean(pred == labels.data) * 100 75 | pbar.set_description(f'Epoch {epoch} [TRAIN] loss = {loss.data[0]:.2f}, acc = {acc:.2f}') 76 | 77 | def train_kpconv(net, optimizer, epoch, dataloader: KPConvLoader): 78 | net.train() 79 | 80 | pbar = tqdm(dataloader, desc=f'Epoch {epoch} [TRAIN]') 81 | jt.sync_all(True) 82 | for input_list in pbar: 83 | L = (len(input_list) - 5) // 4 84 | labels = jt.array(input_list[4 * L + 1]).squeeze(0) 85 | output = net(input_list) 86 | loss = soft_cross_entropy_loss(output, labels) 87 | optimizer.step(loss) 88 | pred = np.argmax(output.data, axis=1) 89 | acc = np.mean(pred == labels.data) * 100 90 | pbar.set_description(f'Epoch {epoch} [TRAIN] loss = {loss.data[0]:.2f}, acc = {acc:.2f}') 91 | # jt.display_memory_info() 92 | 93 | def evaluate(net, epoch, dataloader, args): 94 | total_acc = 0 95 | total_num = 0 96 | 97 | net.eval() 98 | total_time = 0.0 99 | for pts, normals, labels in tqdm(dataloader, desc=f'Epoch {epoch} [Val]'): 100 | # pts = jt.float32(pts.numpy()) 101 | # normals = jt.float32(normals.numpy()) 102 | # labels = jt.int32(labels.numpy()) 103 | # feature = concat((pts, normals), 2) 104 | if args.model == 'pointnet' or args.model == 'dgcnn' : 105 | pts = pts.transpose(0, 2, 1) 106 | 107 | # pts = pts.transpose(0, 2, 1) # for pointnet DGCNN 108 | 109 | # output = net(pts, feature) 110 | if args.model == 'pointnet2': 111 | output = net(pts, normals) 112 | else : 113 | output = net(pts) 114 | # output = net() 115 | pred = np.argmax(output.data, axis=1) 116 | acc = np.sum(pred == labels.data) 117 | total_acc += acc 118 | total_num += labels.shape[0] 119 | acc = 0.0 120 | acc = total_acc / total_num 121 | return acc 122 | 123 | def evaluate_kpconv(net, epoch, dataloader: KPConvLoader): 124 | total_acc = 0 125 | total_num = 0 126 | 127 | net.eval() 128 | total_time = 0.0 129 | for input_list in tqdm(dataloader, desc=f'Epoch {epoch} [Val]'): 130 | L = (len(input_list) - 5) // 4 131 | labels = jt.array(input_list[4 * L + 1]).squeeze(0) 132 | output = net(input_list) 133 | 134 | pred = np.argmax(output.data, axis=1) 135 | acc = np.sum(pred == labels.data) 136 | total_acc += acc 137 | total_num += labels.shape[0] 138 | acc = 0.0 139 | acc = total_acc / total_num 140 | return acc 141 | 142 | 143 | def fast_confusion(true, pred, label_values=None): # used only by kpconv test 144 | """ 145 | Fast confusion matrix (100x faster than Scikit learn). But only works if labels are la 146 | :param true: 147 | :param false: 148 | :param num_classes: 149 | :return: 150 | """ 151 | 152 | # Ensure data is in the right format 153 | true = np.squeeze(true) 154 | pred = np.squeeze(pred) 155 | if len(true.shape) != 1: 156 | raise ValueError('Truth values are stored in a {:d}D array instead of 1D array'. format(len(true.shape))) 157 | if len(pred.shape) != 1: 158 | raise ValueError('Prediction values are stored in a {:d}D array instead of 1D array'. format(len(pred.shape))) 159 | if true.dtype not in [np.int32, np.int64]: 160 | raise ValueError('Truth values are {:s} instead of int32 or int64'.format(true.dtype)) 161 | if pred.dtype not in [np.int32, np.int64]: 162 | raise ValueError('Prediction values are {:s} instead of int32 or int64'.format(pred.dtype)) 163 | true = true.astype(np.int32) 164 | pred = pred.astype(np.int32) 165 | 166 | # Get the label values 167 | if label_values is None: 168 | # From data if they are not given 169 | label_values = np.unique(np.hstack((true, pred))) 170 | else: 171 | # Ensure they are good if given 172 | if label_values.dtype not in [np.int32, np.int64]: 173 | raise ValueError('label values are {:s} instead of int32 or int64'.format(label_values.dtype)) 174 | if len(np.unique(label_values)) < len(label_values): 175 | raise ValueError('Given labels are not unique') 176 | 177 | # Sort labels 178 | label_values = np.sort(label_values) 179 | 180 | # Get the number of classes 181 | num_classes = len(label_values) 182 | 183 | #print(num_classes) 184 | #print(label_values) 185 | #print(np.max(true)) 186 | #print(np.max(pred)) 187 | #print(np.max(true * num_classes + pred)) 188 | 189 | # Start confusion computations 190 | if label_values[0] == 0 and label_values[-1] == num_classes - 1: 191 | 192 | # Vectorized confusion 193 | vec_conf = np.bincount(true * num_classes + pred) 194 | 195 | # Add possible missing values due to classes not being in pred or true 196 | #print(vec_conf.shape) 197 | if vec_conf.shape[0] < num_classes ** 2: 198 | vec_conf = np.pad(vec_conf, (0, num_classes ** 2 - vec_conf.shape[0]), 'constant') 199 | #print(vec_conf.shape) 200 | 201 | # Reshape confusion in a matrix 202 | return vec_conf.reshape((num_classes, num_classes)) 203 | 204 | 205 | else: 206 | 207 | # Ensure no negative classes 208 | if label_values[0] < 0: 209 | raise ValueError('Unsupported negative classes') 210 | 211 | # Get the data in [0,num_classes[ 212 | label_map = np.zeros((label_values[-1] + 1,), dtype=np.int32) 213 | for k, v in enumerate(label_values): 214 | label_map[v] = k 215 | 216 | pred = label_map[pred] 217 | true = label_map[true] 218 | 219 | # Vectorized confusion 220 | vec_conf = np.bincount(true * num_classes + pred) 221 | 222 | # Add possible missing values due to classes not being in pred or true 223 | if vec_conf.shape[0] < num_classes ** 2: 224 | vec_conf = np.pad(vec_conf, (0, num_classes ** 2 - vec_conf.shape[0]), 'constant') 225 | 226 | # Reshape confusion in a matrix 227 | return vec_conf.reshape((num_classes, num_classes)) 228 | 229 | 230 | def classification_test(net, test_loader: KPConvLoader, config, num_votes=100): # used only by kpconv test 231 | print("validation size:", config.validation_size, "batch_num:", config.batch_num) 232 | ############ 233 | # Initialize 234 | ############ 235 | # Choose test smoothing parameter (0 for no smothing, 0.99 for big smoothing) 236 | softmax = jt.nn.Softmax(1) 237 | 238 | # Number of classes including ignored labels 239 | nc_tot = test_loader.num_classes 240 | 241 | # Number of classes predicted by the model 242 | nc_model = config.num_classes 243 | 244 | # Initiate global prediction over test clouds 245 | test_probs = np.zeros((test_loader.num_models, nc_model)) 246 | test_counts = np.zeros((test_loader.num_models, nc_model)) 247 | print("probs shape:", test_probs.shape) 248 | 249 | t = [time.time()] 250 | mean_dt = np.zeros(1) 251 | last_display = time.time() 252 | while np.min(test_counts) < num_votes: 253 | 254 | # Run model on all test examples 255 | # ****************************** 256 | 257 | # Initiate result containers 258 | probs = [] 259 | targets = [] 260 | obj_inds = [] 261 | idx = 0 262 | # Start validation loop 263 | test_loader.prepare_batch_indices() 264 | for input_list in test_loader: 265 | # print("test", idx) 266 | idx += 1 267 | # batch = ModelNet40CustomBatch([input_list]) 268 | # labels, model_inds = batch.labels, batch.model_inds 269 | L = (len(input_list) - 5) // 4 270 | labels = jt.array(input_list[4 * L + 1]).squeeze(0) 271 | model_inds = jt.array(input_list[4 * L + 4]).squeeze(0) 272 | # print(model_inds) 273 | # New time 274 | t = t[-1:] 275 | t += [time.time()] 276 | 277 | # Forward pass 278 | outputs = net(input_list) 279 | 280 | # Get probs and labels 281 | probs += [softmax(outputs).numpy()] 282 | targets += [labels.numpy()] 283 | obj_inds += [model_inds.numpy()] 284 | # print("probs: ", probs) 285 | # print("targets: ", targets) 286 | # print("obj_inds: ", obj_inds) 287 | 288 | # Average timing 289 | t += [time.time()] 290 | mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) 291 | 292 | # Display 293 | if (t[-1] - last_display) > 5.0: 294 | last_display = t[-1] 295 | message = 'Test vote {:.0f} : {:.1f}% (timings : {:4.2f} {:4.2f})' 296 | print(message.format(np.min(test_counts), 297 | 100 * len(obj_inds) / config.validation_size, 298 | 1000 * (mean_dt[0]), 299 | 1000 * (mean_dt[1]))) 300 | # Stack all validation predictions 301 | probs = np.vstack(probs) 302 | targets = np.hstack(targets) 303 | obj_inds = np.hstack(obj_inds) 304 | # print(obj_inds.shape) 305 | 306 | if np.any(test_loader.input_labels[obj_inds] != targets): 307 | raise ValueError('wrong object indices') 308 | 309 | # Compute incremental average (predictions are always ordered) 310 | test_counts[obj_inds] += 1 311 | # print(test_counts.shape) 312 | # print(test_counts) 313 | test_probs[obj_inds] += (probs - test_probs[obj_inds]) / (test_counts[obj_inds]) 314 | 315 | # Save/Display temporary results 316 | # ****************************** 317 | 318 | test_labels = np.array(test_loader.label_values) 319 | 320 | # Compute classification results 321 | C1 = fast_confusion(test_loader.input_labels, 322 | np.argmax(test_probs, axis=1), 323 | test_labels) 324 | 325 | ACC = 100 * np.sum(np.diag(C1)) / (np.sum(C1) + 1e-6) 326 | print('Test Accuracy = {:.1f}%'.format(ACC), flush=True) 327 | 328 | return 329 | 330 | def hook(): 331 | cfg = Modelnet40Config() 332 | net = KPCNN(cfg) 333 | chkp_path = "/mnt/disk1/chentuo/PointNet/KPConv-PyTorch/results/Log_2022-08-04_15-17-48/checkpoints/current_chkp.tar" 334 | checkpoint = torch.load(chkp_path) 335 | net.load_state_dict(checkpoint['model_state_dict']) 336 | # optimizer不控制 337 | print("Model and training state restored.") 338 | other_params = [v for k, v in net.named_parameters() if 'offset' not in k and 'running' not in k] 339 | optimizer = nn.SGD(other_params, lr = cfg.learning_rate, momentum = cfg.momentum, weight_decay=cfg.weight_decay) 340 | f = open('/mnt/disk1/chentuo/PointNet/PointCloudLib/networks/cls/data.txt', 'rb') 341 | data = pickle.load(f) 342 | f.close() 343 | batch = ModelNet40CustomBatch([data]) 344 | hook = auto_diff.Hook("KPCNN") 345 | hook.hook_module(net) 346 | # hook = auto_diff.Hook("KPCNN_optim") 347 | # hook.hook_optimizer(optimizer) 348 | outputs = net(batch) 349 | loss = net.loss(outputs, batch.labels) 350 | acc = net.accuracy(outputs, batch.labels) 351 | # jt.display_memory_info() 352 | # Backward + optimize 353 | # idx = 0 354 | # for k, v in net.named_parameters(): 355 | # if 'offset' not in k and 'running' not in k and 'output_loss' not in k: 356 | # print(" [", idx, "] ", k, " ###### ", jt.grad(loss, v)) 357 | # idx += 1 358 | optimizer.step(loss) 359 | outputs = net(batch) 360 | # jt.display_memory_info() 361 | exit(0) 362 | 363 | if __name__ == '__main__': 364 | 365 | # hook() 366 | freeze_random_seed() 367 | parser = argparse.ArgumentParser(description='Point Cloud Recognition') 368 | parser.add_argument('--eval', action='store_true', default=False) # only used by kpconv 369 | parser.add_argument('--model', type=str, default='[pointnet]', metavar='N', 370 | choices=['pointnet', 'pointnet2', 'pointcnn', 'dgcnn', 'pointconv', 'kpconv'], 371 | help='Model to use') 372 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 373 | help='Size of batch)') 374 | parser.add_argument('--lr', type=float, default=0.02, metavar='LR', 375 | help='learning rate (default: 0.02)') 376 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 377 | help='SGD momentum (default: 0.9)') 378 | parser.add_argument('--num_points', type=int, default=1024, 379 | help='num of points to use') 380 | parser.add_argument('--epochs', type=int, default=300, metavar='N', 381 | help='number of episode to train ') 382 | 383 | args = parser.parse_args() 384 | 385 | 386 | if args.model == 'pointnet': 387 | net = PointNet_cls() 388 | elif args.model == 'pointnet2': 389 | net = PointNet2_cls() 390 | elif args.model == 'pointcnn': 391 | net = PointCNNcls() 392 | elif args.model == 'dgcnn': 393 | net = DGCNN() 394 | elif args.model == 'pointconv': 395 | net = PointConvDensityClsSsg() 396 | elif args.model == 'kpconv': 397 | cfg = Modelnet40Config() 398 | net = KPCNN(cfg) 399 | else: 400 | raise Exception("Not implemented") 401 | 402 | base_lr = args.lr 403 | if args.model != 'kpconv': 404 | optimizer = nn.SGD(net.parameters(), lr = base_lr, momentum = args.momentum) 405 | else: 406 | other_params = [v for k, v in net.named_parameters() if 'offset' not in k and 'running' not in k] 407 | # idx = 0 408 | # for k, v in net.named_parameters(): 409 | # if 'offset' not in k and 'running' not in k: 410 | # print(" [", idx, "] ", k, " ###### ", v.shape) 411 | # idx += 1 412 | # exit(0) 413 | optimizer = nn.SGD(other_params, lr = cfg.learning_rate, momentum = cfg.momentum, weight_decay=cfg.weight_decay) 414 | lr_scheduler = LRScheduler(optimizer, base_lr) 415 | 416 | batch_size = args.batch_size 417 | n_points = args.num_points 418 | if args.model != 'kpconv': 419 | train_dataloader = ModelNet40(n_points=n_points, batch_size=batch_size, train=True, shuffle=True) 420 | val_dataloader = ModelNet40(n_points=n_points, batch_size=batch_size, train=False, shuffle=False) 421 | else: 422 | if not args.eval: 423 | train_dataloader = KPConvLoader(cfg, train=True, num_workers=0) # you can change num_workers to speed up 424 | cfg.validation_size = 250 425 | cfg.val_batch_num = 10 426 | val_dataloader = KPConvLoader(cfg, train=False, num_workers=4) 427 | #### load model #### 428 | if args.eval: 429 | chkp_path = "/mnt/disk1/chentuo/PointNet/PointCloudLib/checkpoints/kpconv/best_chkp.tar" 430 | checkpoint = jt.load(chkp_path) 431 | net.load_state_dict(checkpoint['model_state_dict']) 432 | # optimizer.load_state_dict(checkpoint['optimizer_state_dict']) 433 | # optimizer.load_state_dict({"defaults": checkpoint['optimizer_state_dict']['param_groups'][0]}) 434 | net.eval() 435 | print("Model and training state restored.") 436 | # print(optimizer.state_dict()) 437 | # print("#####") 438 | # print(checkpoint['optimizer_state_dict']['param_groups'][0]) 439 | ##### 440 | step = 0 441 | best_acc = 0 442 | 443 | 444 | for epoch in range(args.epochs): 445 | if args.model == 'kpconv': 446 | if not args.eval: 447 | train_kpconv(net, optimizer, epoch, train_dataloader) 448 | acc = evaluate_kpconv(net, epoch, val_dataloader) 449 | train_dataloader.prepare_batch_indices() 450 | val_dataloader.prepare_batch_indices() 451 | if epoch in cfg.lr_decays: 452 | optimizer.lr *= cfg.lr_decays[epoch] 453 | if cfg.saving: 454 | # Get current state dict 455 | checkpoint_directory = 'checkpoints/kpconv' 456 | save_dict = {'epoch': epoch, 457 | 'model_state_dict': net.state_dict(), 458 | 'optimizer_state_dict': optimizer.state_dict(), 459 | 'saving_path': cfg.saving_path} 460 | checkpoint_path = join(checkpoint_directory, 'current_chkp.tar') 461 | if acc > best_acc: 462 | checkpoint_path = join(checkpoint_directory, 'best_chkp.tar') 463 | 464 | # Save current state of the network (for restoring purposes) 465 | jt.save(save_dict, checkpoint_path) 466 | 467 | # Save checkpoints occasionally 468 | if (epoch + 1) % cfg.checkpoint_gap == 0: 469 | checkpoint_path = join(checkpoint_directory, 'chkp_{:04d}.tar'.format(epoch + 1)) 470 | jt.save(save_dict, checkpoint_path) 471 | else: # kpconv eval 472 | classification_test(net, val_dataloader, cfg) 473 | exit(0) 474 | else: 475 | lr_scheduler.step(len(train_dataloader) * batch_size) 476 | train(net, optimizer, epoch, train_dataloader, args) 477 | acc = evaluate(net, epoch, val_dataloader, args) 478 | 479 | best_acc = max(best_acc, acc) 480 | print(f'[Epoch {epoch}] val acc={acc:.4f}, best={best_acc:.4f}') 481 | -------------------------------------------------------------------------------- /train_partseg.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import os 4 | import sys 5 | from data_utils.shapenet_loader import ShapeNetPart 6 | import numpy as np 7 | import sklearn.metrics as metrics 8 | from misc.utils import LRScheduler 9 | from networks.seg.pointnet_partseg import PointNet_partseg 10 | from networks.seg.pointnet2_partseg import PointNet2_partseg 11 | from networks.seg.pointconv_partseg import PointConvDensity_partseg 12 | from networks.seg.dgcnn_partseg import DGCNN_partseg 13 | from networks.seg.pointcnn_partseg import PointCNN_partseg 14 | 15 | import time 16 | 17 | # jittor related 18 | import jittor as jt 19 | from jittor import nn 20 | 21 | import argparse 22 | 23 | jt.flags.use_cuda = 1 24 | 25 | seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 26 | index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 27 | 28 | def calculate_shape_IoU(pred_np, seg_np, label, class_choice): 29 | # label = label.squeeze(-1) 30 | shape_ious = [] 31 | for shape_idx in range(seg_np.shape[0]): 32 | if not class_choice: 33 | # print (label[shape_idx][0]) 34 | idx = label[shape_idx][0] 35 | start_index = index_start[idx] 36 | num = seg_num[idx] 37 | parts = range(start_index, start_index + num) 38 | else: 39 | parts = range(seg_num[label[0]]) 40 | part_ious = [] 41 | for part in parts: 42 | I = np.sum(np.logical_and(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 43 | U = np.sum(np.logical_or(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 44 | if U == 0: 45 | iou = 1 # If the union of groundtruth and prediction points is empty, then count part IoU as 1 46 | else: 47 | iou = I / float(U) 48 | # print ('iou ', part, iou) 49 | part_ious.append(iou) 50 | shape_ious.append(np.mean(part_ious)) 51 | # cal average class iou 52 | id2cat = ['airplane', 'bag', 'cap', 'car', 'chair', 53 | 'earphone', 'guitar', 'knife', 'lamp', 'laptop', 54 | 'motor', 'mug', 'pistol', 'rocket', 'skateboard','table'] 55 | 56 | # for cat in range (len(id2cat)): 57 | # class_avg_iou = [] 58 | # for idx, iou in enumerate(shape_ious): 59 | # if label[idx] == cat: 60 | # class_avg_iou.append(iou) 61 | # print (id2cat[cat], 'iou =', np.mean(class_avg_iou)) 62 | 63 | return shape_ious 64 | 65 | 66 | def train(model, args): 67 | batch_size = 16 68 | train_loader = ShapeNetPart(partition='trainval', num_points=2048, class_choice=None, batch_size=batch_size, shuffle=True) 69 | test_loader = ShapeNetPart(partition='test', num_points=2048, class_choice=None, batch_size=batch_size, shuffle=False) 70 | 71 | seg_num_all = 50 72 | seg_start_index = 0 73 | 74 | print(str(model)) 75 | base_lr = 0.01 76 | optimizer = nn.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 77 | lr_scheduler = LRScheduler(optimizer, base_lr) 78 | 79 | # criterion = nn.cross_entropy_loss() # here 80 | 81 | best_test_iou = 0 82 | for epoch in range(200): 83 | #################### 84 | # Train 85 | #################### 86 | lr_scheduler.step(len(train_loader) * batch_size) 87 | train_loss = 0.0 88 | count = 0.0 89 | model.train() 90 | train_true_cls = [] 91 | train_pred_cls = [] 92 | train_true_seg = [] 93 | train_pred_seg = [] 94 | train_label_seg = [] 95 | 96 | # debug = 0 97 | for data, label, seg in train_loader: 98 | # with jt.profile_scope() as report: 99 | 100 | seg = seg - seg_start_index 101 | label_one_hot = np.zeros((label.shape[0], 16)) 102 | # print (label.size()) 103 | for idx in range(label.shape[0]): 104 | label_one_hot[idx, label.numpy()[idx,0]] = 1 105 | label_one_hot = jt.array(label_one_hot.astype(np.float32)) 106 | if args.model == 'pointnet' or args.model == 'dgcnn': 107 | data = data.permute(0, 2, 1) # for pointnet it should not be committed 108 | batch_size = data.size()[0] 109 | if args.model == 'pointnet2': 110 | seg_pred = model(data, data, label_one_hot) 111 | else : 112 | seg_pred = model(data, label_one_hot) 113 | seg_pred = seg_pred.permute(0, 2, 1) 114 | # print (seg_pred.size()) 115 | # print (seg_pred.size(), seg.size()) 116 | loss = nn.cross_entropy_loss(seg_pred.view(-1, seg_num_all), seg.view(-1)) 117 | # print (loss.data) 118 | optimizer.step(loss) 119 | 120 | pred = jt.argmax(seg_pred, dim=2)[0] # (batch_size, num_points) 121 | # print ('pred size =', pred.size(), seg.size()) 122 | count += batch_size 123 | train_loss += loss.numpy() * batch_size 124 | seg_np = seg.numpy() # (batch_size, num_points) 125 | pred_np = pred.numpy() # (batch_size, num_points) 126 | # print (type(label)) 127 | 128 | label = label.numpy() # added 129 | 130 | train_true_cls.append(seg_np.reshape(-1)) # (batch_size * num_points) 131 | train_pred_cls.append(pred_np.reshape(-1)) # (batch_size * num_points) 132 | train_true_seg.append(seg_np) 133 | train_pred_seg.append(pred_np) 134 | temp_label = label.reshape(-1, 1) 135 | 136 | train_label_seg.append(temp_label) 137 | 138 | # print(report) 139 | train_true_cls = np.concatenate(train_true_cls) 140 | train_pred_cls = np.concatenate(train_pred_cls) 141 | # print (train_true_cls.shape ,train_pred_cls.shape) 142 | train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls) 143 | 144 | avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls.data, train_pred_cls.data) 145 | # print ('train acc =',train_acc, 'avg_per_class_acc', avg_per_class_acc) 146 | train_true_seg = np.concatenate(train_true_seg, axis=0) 147 | # print (len(train_pred_seg), train_pred_seg[0].shape) 148 | train_pred_seg = np.concatenate(train_pred_seg, axis=0) 149 | # print (len(train_label_seg), train_label_seg[0].size()) 150 | # print (train_label_seg[0]) 151 | train_label_seg = np.concatenate(train_label_seg, axis=0) 152 | # print (train_pred_seg.shape, train_true_seg.shape, train_label_seg.shape) 153 | train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg, None) 154 | outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % (epoch, 155 | train_loss*1.0/count, 156 | train_acc, 157 | avg_per_class_acc, 158 | np.mean(train_ious)) 159 | # io.cprint(outstr) 160 | print (outstr) 161 | #################### 162 | # Test 163 | #################### 164 | test_loss = 0.0 165 | count = 0.0 166 | model.eval() 167 | test_true_cls = [] 168 | test_pred_cls = [] 169 | test_true_seg = [] 170 | test_pred_seg = [] 171 | test_label_seg = [] 172 | for data, label, seg in test_loader: 173 | seg = seg - seg_start_index 174 | label_one_hot = np.zeros((label.shape[0], 16)) 175 | for idx in range(label.shape[0]): 176 | label_one_hot[idx, label.numpy()[idx,0]] = 1 177 | label_one_hot = jt.array(label_one_hot.astype(np.float32)) 178 | if args.model == 'pointnet' or args.model == 'dgcnn': 179 | data = data.permute(0, 2, 1) # for pointnet it should not be committed 180 | batch_size = data.size()[0] 181 | if args.model == 'pointnet2': 182 | seg_pred = model(data, data, label_one_hot) 183 | else : 184 | seg_pred = model(data, label_one_hot) 185 | seg_pred = seg_pred.permute(0, 2, 1) 186 | loss = nn.cross_entropy_loss(seg_pred.view(-1, seg_num_all), seg.view(-1,1).squeeze(-1)) 187 | pred = jt.argmax(seg_pred, dim=2)[0] 188 | count += batch_size 189 | test_loss += loss.numpy() * batch_size 190 | seg_np = seg.numpy() 191 | pred_np = pred.numpy() 192 | label = label.numpy() # added 193 | 194 | test_true_cls.append(seg_np.reshape(-1)) 195 | test_pred_cls.append(pred_np.reshape(-1)) 196 | test_true_seg.append(seg_np) 197 | test_pred_seg.append(pred_np) 198 | test_label_seg.append(label.reshape(-1, 1)) 199 | test_true_cls = np.concatenate(test_true_cls) 200 | test_pred_cls = np.concatenate(test_pred_cls) 201 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 202 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 203 | test_true_seg = np.concatenate(test_true_seg, axis=0) 204 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 205 | test_label_seg = np.concatenate(test_label_seg) 206 | test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg, None) 207 | outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (epoch, 208 | test_loss*1.0/count, 209 | test_acc, 210 | avg_per_class_acc, 211 | np.mean(test_ious)) 212 | print (outstr) 213 | # io.cprint(outstr) 214 | # if np.mean(test_ious) >= best_test_iou: 215 | # best_test_iou = np.mean(test_ious) 216 | # torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name) 217 | 218 | 219 | 220 | if __name__ == "__main__": 221 | # Training settings 222 | parser = argparse.ArgumentParser(description='Point Cloud Recognition') 223 | parser.add_argument('--model', type=str, default='[pointnet]', metavar='N', 224 | choices=['pointnet', 'pointnet2', 'pointcnn', 'dgcnn', 'pointconv'], 225 | help='Model to use') 226 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 227 | help='Size of batch)') 228 | parser.add_argument('--lr', type=float, default=0.02, metavar='LR', 229 | help='learning rate (default: 0.02)') 230 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 231 | help='SGD momentum (default: 0.9)') 232 | parser.add_argument('--num_points', type=int, default=1024, 233 | help='num of points to use') 234 | parser.add_argument('--epochs', type=int, default=300, metavar='N', 235 | help='number of episode to train ') 236 | 237 | args = parser.parse_args() 238 | 239 | 240 | if args.model == 'pointnet': 241 | model = PointNet_partseg(part_num=50) 242 | elif args.model == 'pointnet2': 243 | model = PointNet2_partseg (part_num=50) 244 | elif args.model == 'pointcnn': 245 | model = PointCNN_partseg(part_num=50) 246 | elif args.model == 'dgcnn': 247 | model = DGCNN_partseg(part_num=50) 248 | elif args.model == 'pointconv': 249 | model = PointConvDensity_partseg(part_num=50) 250 | else: 251 | raise Exception("Not implemented") 252 | 253 | train(model, args) 254 | --------------------------------------------------------------------------------