├── figs ├── pcf.jpg ├── mIoU_runtime_log.jpg └── scannet_benchmark.png ├── cpp_wrappers ├── cpp_neighbors │ ├── build.bat │ ├── .DS_Store │ ├── setup.py │ ├── test.py │ ├── neighbors │ │ ├── neighbors.h │ │ └── neighbors.cpp │ └── wrapper.cpp ├── .DS_Store ├── cpp_utils │ ├── .DS_Store │ └── cloud │ │ ├── cloud.cpp │ │ └── cloud.h ├── cpp_subsampling │ ├── .DS_Store │ ├── setup.py │ ├── grid_subsampling │ │ ├── grid_subsampling.h │ │ └── grid_subsampling.cpp │ └── wrapper.cpp ├── compile_wrappers.sh └── cpp_pcf_kernel │ ├── setup.py │ ├── pcf_cuda1.py │ ├── pcf_cuda.cpp │ └── test.py ├── run_distributed.sh ├── compile_wrappers.sh ├── setup.sh ├── CONTRIBUTING.md ├── util ├── cp_batchnorm.py ├── checkpoint.py ├── voxelize.py ├── logger.py ├── common_util.py └── lr.py ├── data_preparation ├── scannetv2_test.txt ├── prepare_data.py └── scannetv2_val.txt ├── LICENSE ├── configs ├── configPCF_10cm.yaml ├── configPCF_10cm_lite.yaml ├── configPCF_5cm_lite.yaml ├── configPCF_5cm.yaml └── configPCF_2cm_PTF2.yaml ├── CODE_OF_CONDUCT.md ├── README.md ├── ACKNOWLEDGEMENTS ├── test_ScanNet_simple.py ├── layer_utils.py ├── transforms.py ├── test_ScanNet_voting.py ├── scannet_data_loader_color_DDP.py └── datasetCommon.py /figs/pcf.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/figs/pcf.jpg -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/build.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | py setup.py build_ext --inplace 3 | 4 | 5 | pause -------------------------------------------------------------------------------- /cpp_wrappers/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/cpp_wrappers/.DS_Store -------------------------------------------------------------------------------- /figs/mIoU_runtime_log.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/figs/mIoU_runtime_log.jpg -------------------------------------------------------------------------------- /figs/scannet_benchmark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/figs/scannet_benchmark.png -------------------------------------------------------------------------------- /run_distributed.sh: -------------------------------------------------------------------------------- 1 | python -m torch.distributed.launch --nproc_per_node=$1 train_ScanNet_DDP_WarmUP.py --config ./$2 2 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/cpp_wrappers/cpp_utils/.DS_Store -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/cpp_wrappers/cpp_neighbors/.DS_Store -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/apple/ml-pointconvformer/HEAD/cpp_wrappers/cpp_subsampling/.DS_Store -------------------------------------------------------------------------------- /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 | cd cpp_neighbors 8 | python3 setup.py build_ext --inplace 9 | cd ../ 10 | cd cpp_pcf_kernel 11 | python3 setup.py install 12 | cd ../ 13 | -------------------------------------------------------------------------------- /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 | cd cpp_neighbors 8 | python3 setup.py build_ext --inplace 9 | cd ../ 10 | cd cpp_pcf_kernel 11 | python3 setup.py install 12 | cd ../ 13 | -------------------------------------------------------------------------------- /setup.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub 4 | add-apt-repository ppa:deadsnakes/ppa 5 | apt update 6 | apt install -y zip libglu1-mesa 7 | apt install python3.8-distutils 8 | 9 | 10 | FILE=./data 11 | if [ ! -d "$FILE" ]; then 12 | echo "mkdir $FILE" 13 | mkdir $FILE 14 | fi 15 | 16 | #unzip ScanNet_withNewNormal.zip 17 | #rm ScanNet_withNewNormal.zip 18 | #cd ../ 19 | 20 | pip install -U ipdb scikit-learn matplotlib open3d easydict 21 | 22 | pip install tensorboard timm termcolor tensorboardX 23 | 24 | cd cpp_wrappers/ 25 | sh compile_wrappers.sh 26 | cd .. 27 | 28 | 29 | 30 | 31 | 32 | 33 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_pcf_kernel/setup.py: -------------------------------------------------------------------------------- 1 | # PCF CUDA Kernel: 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | setup( 9 | name='PCFcuda', 10 | version='1.1', 11 | author='Fuxin Li', 12 | author_email='fli26@apple.com', 13 | description='PointConvFormer CUDA Kernel', 14 | ext_modules=[ 15 | CUDAExtension('pcf_cuda', [ 16 | 'pcf_cuda.cpp', 17 | 'pcf_cuda_kernel.cu', 18 | ], extra_compile_args={'nvcc': ['-L/usr/local/cuda/lib64 -lcudadevrt -lcudart']}) 19 | ], 20 | cmdclass={ 21 | 'build_ext': BuildExtension 22 | }) 23 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contribution Guide 2 | 3 | Thanks for your interest in contributing. This project was released to accompany a research paper for purposes of reproducability, and beyond its publication there are limited plans for future development of the repository. 4 | 5 | While we welcome new pull requests and issues please note that our response may be limited. Forks and out-of-tree improvements are strongly encouraged. 6 | 7 | ## Before you get started 8 | 9 | By submitting a pull request, you represent that you have the right to license your contribution to Apple and the community, and agree by submitting the patch that your contributions are licensed under the [LICENSE](LICENSE). 10 | 11 | We ask that all community members read and observe our [Code of Conduct](CODE_OF_CONDUCT.md). -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | 3 | # Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | from distutils.core import setup, Extension 7 | import numpy.distutils.misc_util 8 | 9 | # Adding OpenCV to project 10 | # ************************ 11 | 12 | # Adding sources of the project 13 | # ***************************** 14 | 15 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 16 | "neighbors/neighbors.cpp", 17 | "wrapper.cpp"] 18 | 19 | module = Extension(name="radius_neighbors", 20 | sources=SOURCES, 21 | extra_compile_args=['-std=c++11', 22 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 23 | 24 | 25 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 26 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | 3 | # Copyright (c) 2019 HuguesTHOMAS 4 | 5 | from distutils.core import setup, Extension 6 | import numpy.distutils.misc_util 7 | 8 | # Adding OpenCV to project 9 | # ************************ 10 | 11 | # Adding sources of the project 12 | # ***************************** 13 | 14 | m_name = "grid_subsampling" 15 | 16 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 17 | "grid_subsampling/grid_subsampling.cpp", 18 | "wrapper.cpp"] 19 | 20 | module = Extension(m_name, 21 | sources=SOURCES, 22 | extra_compile_args=['-std=c++11', 23 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 24 | 25 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 26 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/test.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | 3 | # Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | import radius_neighbors as cpp_neighbors 7 | import numpy as np 8 | import time 9 | from sklearn.neighbors import KDTree 10 | import torch 11 | 12 | 13 | data = torch.load('test_pt.pth') 14 | queries = data[0] 15 | queries = np.single(queries) 16 | for i in range(10): 17 | # queries = np.random.rand(50000,3) 18 | queries = np.single(queries) 19 | t1 = time.time() 20 | indx = cpp_neighbors.batch_kquery(queries, queries, [queries.shape[0]], queries.shape[0], K=int(16)) 21 | t2 = time.time() 22 | print('Time for nanoflann: ', (t2 - t1)*1000, 'ms') 23 | print(indx[0, :]) 24 | print(indx[1, :]) 25 | print(indx[0:10, 3]) 26 | t3 = time.time() 27 | kdt = KDTree(queries) 28 | neighbors_idx = kdt.query(queries, k=16, return_distance=False) 29 | t4 = time.time() 30 | print('Time for SKlearn: ', (t4-t3)*1000, 'ms') 31 | print(neighbors_idx[0, :]) 32 | print(neighbors_idx[1, :]) 33 | print(neighbors_idx[0:10, 3]) 34 | # neighbors_idx = neighbors_idx[:, ::dialated_rate] 35 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /util/cp_batchnorm.py: -------------------------------------------------------------------------------- 1 | # cp_batchnorm.py taken from https://github.com/csrhddlam/pytorch-checkpoint/blob/master/ 2 | # copyright (c) 2018 Huiyu Wang 3 | # MIT License: https://github.com/csrhddlam/pytorch-checkpoint/blob/master/LICENSE 4 | 5 | import torch 6 | from torch.nn import functional as F 7 | 8 | 9 | class CpBatchNorm2d(torch.nn.BatchNorm2d): 10 | def __init__(self, *args, **kwargs): 11 | super(CpBatchNorm2d, self).__init__(*args, **kwargs) 12 | 13 | def forward(self, input): 14 | self._check_input_dim(input) 15 | if self.training: 16 | exponential_average_factor = 0.0 17 | if self.training and self.track_running_stats: 18 | self.num_batches_tracked += 1 19 | if self.momentum is None: # use cumulative moving average 20 | exponential_average_factor = 1.0 / self.num_batches_tracked.item() 21 | else: # use exponential moving average 22 | exponential_average_factor = self.momentum 23 | return F.batch_norm( 24 | input, self.running_mean, self.running_var, self.weight, self.bias, 25 | self.training or not self.track_running_stats, 26 | exponential_average_factor, self.eps) 27 | else: 28 | return F.batch_norm( 29 | input, self.running_mean, self.running_var, self.weight, self.bias, 30 | self.training or not self.track_running_stats, 0.0, self.eps) 31 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_pcf_kernel/pcf_cuda1.py: -------------------------------------------------------------------------------- 1 | # PCF CUDA Kernel: 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | # Example code using PCF CUDA Kernel 7 | 8 | import torch 9 | 10 | import pcf_cuda 11 | 12 | 13 | class PCFFunction(torch.autograd.Function): 14 | # Currently, PCF CUDA kernel is by default not used since it leads to some mysterious training degradations, however, it works a bit faster than the non-CUDA 15 | # implementation at testing time and should be used to achieve the speed benchmark performance 16 | @staticmethod 17 | def forward(ctx, input_feat, neighbor_inds, guidance, weightnet): 18 | # Make sure we are not computing gradient on neighbor_inds 19 | neighbor_inds.requires_grad = False 20 | output = pcf_cuda.forward(input_feat, neighbor_inds, guidance, weightnet) 21 | ctx.save_for_backward({input_feat, neighbor_inds, guidance, weightnet}) 22 | return output 23 | 24 | @staticmethod 25 | def backward(ctx, grad_output): 26 | grad_input, grad_guidance, grad_weight = pcf_cuda.backward(grad_output.contiguous(), *ctx.saved_tensors) 27 | return grad_input, None, grad_guidance, grad_weight 28 | 29 | 30 | class PCF(torch.nn.Module): 31 | def __init__(self): 32 | super(PCF, self).__init__() 33 | 34 | def forward(self, input_features, neighbor_inds, guidance, weightnet): 35 | return PCFFunction.apply(input_features, neighbor_inds, guidance, weightnet) 36 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/neighbors/neighbors.h: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | #include "../../cpp_utils/cloud/cloud.h" 6 | #include "../../cpp_utils/nanoflann/nanoflann.hpp" 7 | 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | 14 | void ordered_neighbors(vector& queries, 15 | vector& supports, 16 | vector& neighbors_indices, 17 | float radius); 18 | 19 | void batch_ordered_neighbors(vector& queries, 20 | vector& supports, 21 | vector& q_batches, 22 | vector& s_batches, 23 | vector& neighbors_indices, 24 | float radius); 25 | 26 | void batch_nanoflann_neighbors(vector& queries, 27 | vector& supports, 28 | vector& q_batches, 29 | vector& s_batches, 30 | vector& neighbors_indices, 31 | float radius); 32 | 33 | void batch_nanoflann_knn(vector& queries, 34 | vector& supports, 35 | vector& q_batches, 36 | vector& s_batches, 37 | vector& neighbors_indices, 38 | size_t K); 39 | -------------------------------------------------------------------------------- /data_preparation/scannetv2_test.txt: -------------------------------------------------------------------------------- 1 | scene0707_00 2 | scene0708_00 3 | scene0709_00 4 | scene0710_00 5 | scene0711_00 6 | scene0712_00 7 | scene0713_00 8 | scene0714_00 9 | scene0715_00 10 | scene0716_00 11 | scene0717_00 12 | scene0718_00 13 | scene0719_00 14 | scene0720_00 15 | scene0721_00 16 | scene0722_00 17 | scene0723_00 18 | scene0724_00 19 | scene0725_00 20 | scene0726_00 21 | scene0727_00 22 | scene0728_00 23 | scene0729_00 24 | scene0730_00 25 | scene0731_00 26 | scene0732_00 27 | scene0733_00 28 | scene0734_00 29 | scene0735_00 30 | scene0736_00 31 | scene0737_00 32 | scene0738_00 33 | scene0739_00 34 | scene0740_00 35 | scene0741_00 36 | scene0742_00 37 | scene0743_00 38 | scene0744_00 39 | scene0745_00 40 | scene0746_00 41 | scene0747_00 42 | scene0748_00 43 | scene0749_00 44 | scene0750_00 45 | scene0751_00 46 | scene0752_00 47 | scene0753_00 48 | scene0754_00 49 | scene0755_00 50 | scene0756_00 51 | scene0757_00 52 | scene0758_00 53 | scene0759_00 54 | scene0760_00 55 | scene0761_00 56 | scene0762_00 57 | scene0763_00 58 | scene0764_00 59 | scene0765_00 60 | scene0766_00 61 | scene0767_00 62 | scene0768_00 63 | scene0769_00 64 | scene0770_00 65 | scene0771_00 66 | scene0772_00 67 | scene0773_00 68 | scene0774_00 69 | scene0775_00 70 | scene0776_00 71 | scene0777_00 72 | scene0778_00 73 | scene0779_00 74 | scene0780_00 75 | scene0781_00 76 | scene0782_00 77 | scene0783_00 78 | scene0784_00 79 | scene0785_00 80 | scene0786_00 81 | scene0787_00 82 | scene0788_00 83 | scene0789_00 84 | scene0790_00 85 | scene0791_00 86 | scene0792_00 87 | scene0793_00 88 | scene0794_00 89 | scene0795_00 90 | scene0796_00 91 | scene0797_00 92 | scene0798_00 93 | scene0799_00 94 | scene0800_00 95 | scene0801_00 96 | scene0802_00 97 | scene0803_00 98 | scene0804_00 99 | scene0805_00 100 | scene0806_00 101 | -------------------------------------------------------------------------------- /util/checkpoint.py: -------------------------------------------------------------------------------- 1 | # checkpoint.py taken from https://github.com/csrhddlam/pytorch-checkpoint/blob/master/checkpoint.py 2 | # copyright (c) 2018 Huiyu Wang 3 | # MIT License: https://github.com/csrhddlam/pytorch-checkpoint/blob/master/LICENSE 4 | 5 | import torch 6 | import warnings 7 | 8 | 9 | def detach_variable(inputs): 10 | if isinstance(inputs, tuple): 11 | out = [] 12 | for inp in inputs: 13 | x = inp.detach() 14 | x.requires_grad = inp.requires_grad 15 | out.append(x) 16 | return tuple(out) 17 | else: 18 | raise RuntimeError( 19 | "Only tuple of tensors is supported. Got Unsupported input type: ", type(inputs).__name__) 20 | 21 | 22 | def check_backward_validity(inputs): 23 | if not any(inp.requires_grad for inp in inputs): 24 | warnings.warn("None of the inputs have requires_grad=True. Gradients will be None") 25 | 26 | 27 | class CheckpointFunction(torch.autograd.Function): 28 | @staticmethod 29 | def forward(ctx, run_function, length, *args): 30 | ctx.run_function = run_function 31 | ctx.input_tensors = list(args[:length]) 32 | ctx.input_params = list(args[length:]) 33 | with torch.no_grad(): 34 | output_tensors = ctx.run_function(*ctx.input_tensors) 35 | return output_tensors 36 | 37 | @staticmethod 38 | def backward(ctx, *output_grads): 39 | for i in range(len(ctx.input_tensors)): 40 | temp = ctx.input_tensors[i] 41 | ctx.input_tensors[i] = temp.detach() 42 | ctx.input_tensors[i].requires_grad = temp.requires_grad 43 | with torch.enable_grad(): 44 | output_tensors = ctx.run_function(*ctx.input_tensors) 45 | input_grads = torch.autograd.grad(output_tensors, ctx.input_tensors + ctx.input_params, output_grads, allow_unused=True) 46 | return (None, None) + input_grads 47 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | #include "../../cpp_utils/cloud/cloud.h" 7 | 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | class SampledData 14 | { 15 | public: 16 | 17 | // Elements 18 | // ******** 19 | 20 | int count; 21 | PointXYZ point; 22 | vector features; 23 | vector> labels; 24 | 25 | 26 | // Methods 27 | // ******* 28 | 29 | // Constructor 30 | SampledData() 31 | { 32 | count = 0; 33 | point = PointXYZ(); 34 | } 35 | 36 | SampledData(const size_t fdim, const size_t ldim) 37 | { 38 | count = 0; 39 | point = PointXYZ(); 40 | features = vector(fdim); 41 | labels = vector>(ldim); 42 | } 43 | 44 | // Method Update 45 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 46 | { 47 | count += 1; 48 | point += p; 49 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 50 | int i = 0; 51 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 52 | { 53 | labels[i][*it] += 1; 54 | i++; 55 | } 56 | return; 57 | } 58 | void update_features(const PointXYZ p, vector::iterator f_begin) 59 | { 60 | count += 1; 61 | point += p; 62 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 63 | return; 64 | } 65 | void update_classes(const PointXYZ p, vector::iterator l_begin) 66 | { 67 | count += 1; 68 | point += p; 69 | int i = 0; 70 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 71 | { 72 | labels[i][*it] += 1; 73 | i++; 74 | } 75 | return; 76 | } 77 | void update_points(const PointXYZ p) 78 | { 79 | count += 1; 80 | point += p; 81 | return; 82 | } 83 | }; 84 | 85 | 86 | 87 | void grid_subsampling(vector& original_points, 88 | vector& subsampled_points, 89 | vector& original_features, 90 | vector& subsampled_features, 91 | vector& original_classes, 92 | vector& subsampled_classes, 93 | float sampleDl, 94 | int verbose); 95 | 96 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 2 | 3 | IMPORTANT: This Apple software is supplied to you by Apple 4 | Inc. ("Apple") in consideration of your agreement to the following 5 | terms, and your use, installation, modification or redistribution of 6 | this Apple software constitutes acceptance of these terms. If you do 7 | not agree with these terms, please do not use, install, modify or 8 | redistribute this Apple software. 9 | 10 | In consideration of your agreement to abide by the following terms, and 11 | subject to these terms, Apple grants you a personal, non-exclusive 12 | license, under Apple's copyrights in this original Apple software (the 13 | "Apple Software"), to use, reproduce, modify and redistribute the Apple 14 | Software, with or without modifications, in source and/or binary forms; 15 | provided that if you redistribute the Apple Software in its entirety and 16 | without modifications, you must retain this notice and the following 17 | text and disclaimers in all such redistributions of the Apple Software. 18 | Neither the name, trademarks, service marks or logos of Apple Inc. may 19 | be used to endorse or promote products derived from the Apple Software 20 | without specific prior written permission from Apple. Except as 21 | expressly stated in this notice, no other rights or licenses, express or 22 | implied, are granted by Apple herein, including but not limited to any 23 | patent rights that may be infringed by your derivative works or by other 24 | works in which the Apple Software may be incorporated. 25 | 26 | The Apple Software is provided by Apple on an "AS IS" basis. APPLE 27 | MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION 28 | THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS 29 | FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND 30 | OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS. 31 | 32 | IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL 33 | OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 34 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 35 | INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION, 36 | MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED 37 | AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE), 38 | STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE 39 | POSSIBILITY OF SUCH DAMAGE. 40 | -------------------------------------------------------------------------------- /configs/configPCF_10cm.yaml: -------------------------------------------------------------------------------- 1 | # See get_default_configs() in train_ScanNet_DDP_Warmup.py 2 | # and get_default_configs() in model_architecture.py for the meanings of those configs 3 | BATCH_NORM: True 4 | BATCH_SIZE: 16 5 | NUM_WORKERS: 24 6 | USE_XYZ: True 7 | USE_MLP: False 8 | USE_WEIGHT: True 9 | USE_PE: True 10 | device: 'cuda' 11 | manual_seed: 1 12 | sync_bn: True 13 | 14 | print_freq: 5 15 | eval_freq: 5 16 | save_freq: 5 17 | 18 | TIME: False 19 | 20 | MAX_POINTS_NUM: 550000 21 | use_ASPP: False 22 | NormalizedXYZ: False 23 | K_forward: [16, 16, 16, 16, 16] 24 | K_propagate: [16, 16, 16, 16, 16] 25 | K_self: [16, 16, 16, 16, 16] 26 | point_dim: 3 27 | total_epoches: 300 28 | 29 | num_level: 5 30 | grid_size: [0.1, 0.2, 0.4, 0.8, 1.6] 31 | base_dim: 64 32 | feat_dim: [64, 128, 192, 256, 384] 33 | mid_dim: [16,16,16,16,16] 34 | mid_dim_back: 1 35 | 36 | guided_level: 0 37 | num_heads: 8 38 | resblocks: [ 0, 2, 4, 6, 6] 39 | resblocks_back: [ 0, 0, 0, 0, 0] 40 | 41 | train_data_path: './data/ScanNet_withNewNormal/train/*.pth' 42 | val_data_path: './data/ScanNet_withNewNormal/val/*.pth' 43 | test_data_path: './data/ScanNet_withNewNormal/test/*.pth' 44 | pretrain: null 45 | optimizer: 'AdamW' 46 | adamw_decay: 0.05 47 | learning_rate: 0.02 48 | gamma: 0.5 49 | label_smoothing: 0.2 50 | scheduler: 'MultiStepWithWarmup' # 'MultiStepWithWarmup' 'CosineAnnealingWarmupRestarts' 51 | milestones: [80, 140, 180, 220, 260] 52 | ft_learning_rate: 0.016 53 | decay_rate: 0.0001 54 | ignore_label: -100 55 | drop_path_rate: 0. 56 | dropout_rate: 0. 57 | dropout_fc: 0. 58 | layer_norm_guidance: False 59 | 60 | scheduler_update: 'step' 61 | warmup: 'linear' 62 | warmup_epochs: 10 63 | warmup_ratio: 0.00001 64 | 65 | use_tensorboard: False 66 | model_name: 'NewPointConvFormer_10cm' 67 | experiment_dir: './guided_experiment_10cm/' 68 | ft_experiment_dir: './ft_guided_experiment_10cm/' 69 | num_classes: 20 70 | ft_model_path: '/mnt/task_runtime/guided_experiment_10cm/hsc/model/model_best.pth' 71 | 72 | eval_path: './evaluation_10cm/' 73 | 74 | classes: ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 75 | 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 76 | 'refridgerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 77 | -------------------------------------------------------------------------------- /configs/configPCF_10cm_lite.yaml: -------------------------------------------------------------------------------- 1 | # See get_default_configs() in train_ScanNet_DDP_Warmup.py 2 | # and get_default_configs() in model_architecture.py for the meanings of those configs 3 | BATCH_NORM: True 4 | BATCH_SIZE: 16 5 | NUM_WORKERS: 24 6 | USE_XYZ: True 7 | USE_MLP: False 8 | USE_WEIGHT: True 9 | USE_PE: True 10 | device: 'cuda' 11 | manual_seed: 1 12 | sync_bn: True 13 | 14 | print_freq: 5 15 | eval_freq: 5 16 | save_freq: 5 17 | 18 | TIME: False 19 | 20 | MAX_POINTS_NUM: 550000 21 | use_ASPP: False 22 | NormalizedXYZ: False 23 | K_forward: [16, 16, 16, 16, 16] 24 | K_propagate: [16, 16, 16, 16, 16] 25 | K_self: [16, 16, 16, 16, 16] 26 | point_dim: 3 27 | total_epoches: 300 28 | 29 | num_level: 5 30 | grid_size: [0.1, 0.2, 0.4, 0.8, 1.6] 31 | base_dim: 64 32 | feat_dim: [64, 128, 192, 256, 384] 33 | mid_dim: [4,4,4,4,4] 34 | mid_dim_back: 1 35 | 36 | guided_level: 0 37 | num_heads: 8 38 | resblocks: [ 0, 3,3,3,3] 39 | resblocks_back: [ 0, 0, 0, 0, 0] 40 | 41 | train_data_path: './data/ScanNet_withNewNormal/train/*.pth' 42 | val_data_path: './data/ScanNet_withNewNormal/val/*.pth' 43 | test_data_path: './data/ScanNet_withNewNormal/test/*.pth' 44 | pretrain: null 45 | optimizer: 'AdamW' 46 | adamw_decay: 0.05 47 | learning_rate: 0.02 48 | gamma: 0.5 49 | label_smoothing: 0.2 50 | scheduler: 'MultiStepWithWarmup' # 'MultiStepWithWarmup' 'CosineAnnealingWarmupRestarts' 51 | milestones: [80, 140, 180, 220, 260] 52 | ft_learning_rate: 0.016 53 | decay_rate: 0.0001 54 | ignore_label: -100 55 | drop_path_rate: 0. 56 | dropout_rate: 0. 57 | dropout_fc: 0. 58 | layer_norm_guidance: False 59 | 60 | scheduler_update: 'step' 61 | warmup: 'linear' 62 | warmup_epochs: 10 63 | warmup_ratio: 0.00001 64 | 65 | use_tensorboard: False 66 | model_name: 'NewPointConvFormer_10cm' 67 | experiment_dir: './guided_experiment_10cm/' 68 | ft_experiment_dir: './ft_guided_experiment_10cm/' 69 | num_classes: 20 70 | ft_model_path: '/mnt/task_runtime/guided_experiment_10cm/hsc/model/model_best.pth' 71 | 72 | eval_path: './evaluation_10cm/' 73 | 74 | classes: ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 75 | 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 76 | 'refridgerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 77 | -------------------------------------------------------------------------------- /configs/configPCF_5cm_lite.yaml: -------------------------------------------------------------------------------- 1 | # See get_default_configs() in train_ScanNet_DDP_Warmup.py 2 | # and get_default_configs() in model_architecture.py for the meanings of those configs 3 | BATCH_NORM: True 4 | BATCH_SIZE: 3 5 | NUM_WORKERS: 24 6 | USE_XYZ: True 7 | USE_MLP: False 8 | USE_WEIGHT: True 9 | USE_PE: True 10 | device: 'cuda' 11 | DDP: True 12 | manual_seed: 1 13 | sync_bn: True 14 | 15 | print_freq: 5 16 | eval_freq: 5 17 | save_freq: 5 18 | 19 | TIME: False 20 | 21 | MAX_POINTS_NUM: 550000 22 | use_ASPP: False 23 | NormalizedXYZ: False 24 | K_forward: [16, 16, 16, 16, 16] 25 | K_propagate: [16, 16, 16, 16, 16] 26 | K_self: [16, 16, 16, 16, 16] 27 | point_dim: 3 28 | total_epoches: 300 29 | 30 | num_level: 5 31 | grid_size: [0.05, 0.1, 0.2, 0.4, 0.8] 32 | dialated_rate: [ 1, 1, 1, 1] 33 | base_dim: 64 34 | feat_dim: [64, 128, 192, 256, 384] 35 | mid_dim: [4,4,4,4,4] 36 | mid_dim_back: 1 37 | 38 | guided_level: 0 39 | num_heads: 8 40 | resblocks: [ 0, 3,3,3,3] 41 | resblocks_back: [ 0, 0, 0, 0, 0] 42 | 43 | train_data_path: './data/ScanNet_withNewNormal/train/*.pth' 44 | val_data_path: './data/ScanNet_withNewNormal/val/*.pth' 45 | test_data_path: './data/ScanNet_withNewNormal/test/*.pth' 46 | pretrain: null 47 | optimizer: 'AdamW' 48 | adamw_decay: 0.05 49 | learning_rate: 0.01 50 | accum_iter: 3 51 | gamma: 0.5 52 | label_smoothing: 0.2 53 | mix3D: True 54 | scheduler: 'MultiStepWithWarmup' # 'MultiStepWithWarmup' 'CosineAnnealingWarmupRestarts' 55 | milestones: [60, 120, 170, 220, 260] 56 | ft_learning_rate: 0.016 57 | decay_rate: 0.0001 58 | multi_gpu: null 59 | ignore_label: -100 60 | drop_path_rate: 0. 61 | dropout_rate: 0. 62 | dropout_fc: 0. 63 | layer_norm_guidance: False 64 | 65 | scheduler_update: 'step' 66 | warmup: 'linear' 67 | warmup_epochs: 10 68 | warmup_ratio: 0.00001 69 | 70 | use_tensorboard: False 71 | model_name: 'NewPointConvFormer_10cm' 72 | experiment_dir: './guided_experiment_10cm/' 73 | ft_experiment_dir: './ft_guided_experiment_10cm/' 74 | num_classes: 20 75 | ft_model_path: '/mnt/task_runtime/guided_experiment_10cm/hsc/model/model_best.pth' 76 | 77 | eval_path: './evaluation_10cm/' 78 | 79 | classes: ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 80 | 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 81 | 'refridgerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 82 | -------------------------------------------------------------------------------- /configs/configPCF_5cm.yaml: -------------------------------------------------------------------------------- 1 | # See get_default_configs() in train_ScanNet_DDP_Warmup.py 2 | # and get_default_configs() in model_architecture.py for the meanings of those configs 3 | BATCH_NORM: True 4 | BATCH_SIZE: 3 5 | NUM_WORKERS: 24 6 | USE_XYZ: True 7 | USE_MLP: False 8 | USE_WEIGHT: True 9 | USE_PE: True 10 | device: 'cuda' 11 | manual_seed: 1 12 | sync_bn: True 13 | 14 | print_freq: 5 15 | eval_freq: 5 16 | save_freq: 5 17 | 18 | TIME: False 19 | 20 | MAX_POINTS_NUM: 550000 21 | use_ASPP: False 22 | NormalizedXYZ: False 23 | K_forward: [16, 16, 16, 16, 16] 24 | K_propagate: [16, 16, 16, 16, 16] 25 | K_self: [16, 16, 16, 16, 16] 26 | point_dim: 3 27 | total_epoches: 300 28 | 29 | num_level: 5 30 | grid_size: [0.05, 0.1, 0.2, 0.4, 0.8] 31 | dialated_rate: [ 1, 1, 1, 1] 32 | base_dim: 64 33 | feat_dim: [64, 128, 192, 256, 384] 34 | mid_dim: [16, 16,16,16,16] 35 | mid_dim_back: 1 36 | 37 | guided_level: 0 38 | num_heads: 8 39 | resblocks: [ 0, 2, 4, 6, 6] 40 | # resblocks: [ 0, 1, 3, 5, 5, 5] 41 | resblocks_back: [ 0, 0, 0, 0, 0] 42 | # resblocks_back: [ 1, 1, 1, 1, 1, 1] 43 | 44 | train_data_path: './data/ScanNet_withNewNormal/train/*.pth' 45 | val_data_path: './data/ScanNet_withNewNormal/val/*.pth' 46 | test_data_path: './data/ScanNet_withNewNormal/test/*.pth' 47 | pretrain: null 48 | optimizer: 'AdamW' 49 | adamw_decay: 0.05 50 | learning_rate: 0.01 51 | accum_iter: 3 52 | gamma: 0.5 53 | label_smoothing: 0.2 54 | mix3D: True 55 | scheduler: 'MultiStepWithWarmup' # 'MultiStepWithWarmup' 'CosineAnnealingWarmupRestarts' 56 | # milestones: [200, 300, 400, 500] 57 | milestones: [60, 120, 170, 220, 260] 58 | ft_learning_rate: 0.016 59 | decay_rate: 0.0001 60 | multi_gpu: null 61 | ignore_label: -100 62 | drop_path_rate: 0. 63 | dropout_rate: 0. 64 | dropout_fc: 0. 65 | layer_norm_guidance: False 66 | 67 | scheduler_update: 'step' 68 | warmup: 'linear' 69 | warmup_epochs: 10 70 | warmup_ratio: 0.00001 71 | 72 | use_tensorboard: False 73 | model_name: 'NewPointConvFormer_10cm' 74 | experiment_dir: './guided_experiment_10cm/' 75 | ft_experiment_dir: './ft_guided_experiment_10cm/' 76 | num_classes: 20 77 | ft_model_path: '/mnt/task_runtime/guided_experiment_10cm/hsc/model/model_best.pth' 78 | 79 | eval_path: './evaluation_10cm/' 80 | 81 | classes: ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 82 | 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 83 | 'refridgerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 84 | -------------------------------------------------------------------------------- /configs/configPCF_2cm_PTF2.yaml: -------------------------------------------------------------------------------- 1 | # See get_default_configs() in train_ScanNet_DDP_Warmup.py 2 | # and get_default_configs() in model_architecture.py for the meanings of those configs 3 | BATCH_NORM: True 4 | BATCH_SIZE: 2 5 | NUM_WORKERS: 24 6 | USE_XYZ: True 7 | USE_MLP: False 8 | USE_PE: True 9 | USE_WEIGHT: True 10 | use_level_1: False 11 | device: 'cuda' 12 | manual_seed: 1 13 | sync_bn: True 14 | 15 | print_freq: 5 16 | eval_freq: 5 17 | save_freq: 5 18 | 19 | TIME: False 20 | 21 | # Use this for saving memory, during testing time, one should reset this to 550000 22 | MAX_POINTS_NUM: 120000 23 | USE_MULTI_GPU: True 24 | use_ASPP: False 25 | NormalizedXYZ: False 26 | K_forward: [16, 16, 16, 16, 16,16] 27 | K_propagate: [16, 16, 16, 16, 16,16] 28 | K_self: [16, 16, 16, 16, 16,16] 29 | point_dim: 3 30 | total_epoches: 300 31 | 32 | num_level: 5 33 | grid_size: [0.02, 0.06, 0.15, 0.375, 0.9375] 34 | dialated_rate: [ 1, 1, 1, 1] 35 | base_dim: 64 36 | feat_dim: [64, 128, 192, 256, 384] 37 | mid_dim: [16,16,16,16,16] 38 | mid_dim_back: 3 39 | label_smoothing: 0.2 40 | drop_path_rate: 0.2 41 | mix3D: True 42 | 43 | guided_level: 0 44 | num_heads: 8 45 | resblocks: [ 0, 2, 4, 6, 6, 2] 46 | resblocks_back: [ 0,0,0,0,0] 47 | 48 | train_data_path: './data/ScanNet_withNewNormal/train/*.pth' 49 | val_data_path: './data/ScanNet_withNewNormal/val/*.pth' 50 | test_data_path: './data/ScanNet_withNewNormal/test/*.pth' 51 | pretrain: null 52 | optimizer: 'AdamW' 53 | adamw_decay: 0.05 54 | learning_rate: 0.01 55 | accum_iter: 5 56 | gamma: 0.5 57 | scheduler: 'MultiStepWithWarmup' # 'MultiStepWithWarmup' 'CosineAnnealingWarmupRestarts' 58 | milestones: [70, 130, 180, 220, 260] 59 | ft_learning_rate: 0.016 60 | decay_rate: 0.0001 61 | multi_gpu: null 62 | ignore_label: -100 63 | dropout_rate: 0. 64 | dropout_fc: 0. 65 | layer_norm_guidance: False 66 | 67 | scheduler_update: 'step' 68 | warmup: 'linear' 69 | warmup_epochs: 10 70 | warmup_ratio: 0.00001 71 | 72 | use_tensorboard: False 73 | model_name: 'NewPointConvFormer_2cm' 74 | experiment_dir: './guided_experiment_2cm/' 75 | ft_experiment_dir: './ft_guided_experiment_2cm/' 76 | num_classes: 20 77 | ft_model_path: '/mnt/task_runtime/guided_experiment_2cm/hsc/model/model_best.pth' 78 | 79 | eval_path: './evaluation_2cm/' 80 | 81 | classes: ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 82 | 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 83 | 'refridgerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 84 | -------------------------------------------------------------------------------- /data_preparation/prepare_data.py: -------------------------------------------------------------------------------- 1 | # Prepare ScanNet Data into the form we want 2 | 3 | import os 4 | import glob 5 | import plyfile 6 | import numpy as np 7 | import multiprocessing as mp 8 | import torch 9 | import open3d as o3d 10 | import argparse 11 | 12 | 13 | def load_obj_with_normals(filepath): 14 | mesh = o3d.io.read_triangle_mesh(str(filepath)) 15 | if not mesh.has_vertex_normals(): 16 | mesh.compute_vertex_normals() 17 | coords = np.asarray(mesh.vertices) 18 | normals = np.asarray(mesh.vertex_normals) 19 | colors = np.asarray(mesh.vertex_colors) 20 | feats = np.hstack((colors, normals)) 21 | 22 | return coords, feats 23 | 24 | 25 | def newf(fn): 26 | fn2 = fn[:-3]+'labels.ply' 27 | coords, feats = load_obj_with_normals(fn) 28 | a = plyfile.PlyData().read(fn2) 29 | labels = remapper[np.array(a.elements[0]['label'])] 30 | torch.save((coords, feats, labels, fn.split('/')[-1]), fn[:-4]+'.pth') 31 | print(fn, fn2) 32 | 33 | 34 | def newf_test(fn): 35 | coords, feats = load_obj_with_normals(fn) 36 | num_points = coords.shape[0] 37 | labels = np.zeros(num_points).astype(np.int32) 38 | torch.save((coords, feats, labels, fn.split('/')[-1]), fn[:-4] + '.pth') 39 | print(fn) 40 | 41 | if __name__ == '__main__': 42 | # Map relevant classes to {0,1,...,19}, and ignored classes to -100 43 | parser = argparse.ArgumentParser(description="Prepare ScanNet Data") 44 | parser.add_argument("basepath", help="base path of the downloaded ScanNet dataset") 45 | parser.add_argument("split", help="The split that you want to parse (train/validation/test). For test, the labels are set to 0") 46 | args = parser.parse_args() 47 | data_path = args.basepath 48 | split = args.split 49 | if split != "test": 50 | scan_path = os.path.join(data_path, "scans") 51 | else: 52 | scan_path = os.path.join(data_path, "scans_test") 53 | 54 | with open("scannetv2_" + split + ".txt") as f: 55 | scans_name = [x.strip() for x in f.readlines()] 56 | 57 | print(len(scans_name)) 58 | 59 | for scan_name in scans_name: 60 | print("copy scene: ", scan_name) 61 | ply_path = os.path.join(scan_path, scan_name, scan_name + "_vh_clean_2.ply") 62 | label_path = os.path.join(scan_path, scan_name, scan_name + "_vh_clean_2.labels.ply") 63 | segsjson_path = os.path.join(scan_path, scan_name, scan_name + "_vh_clean_2.0.010000.segs.json") 64 | aggjson_path = os.path.join(scan_path, scan_name, scan_name + ".aggregation.json") 65 | 66 | os.system('cp %s %s' % (ply_path, split)) 67 | 68 | if split != "test": 69 | os.system('cp %s %s' % (label_path, split)) 70 | os.system('cp %s %s' % (segsjson_path, split)) 71 | os.system('cp %s %s' % (aggjson_path, split)) 72 | 73 | remapper = np.ones(150)*(-100) 74 | for i, x in enumerate([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]): 75 | remapper[x] = i 76 | 77 | files = sorted(glob.glob(os.path.join(data_path, '*/*_vh_clean_2.ply'))) 78 | if split != 'test': 79 | files2 = sorted(glob.glob(os.path.join(data_path, '*/*_vh_clean_2.labels.ply'))) 80 | assert len(files) == len(files2) 81 | p = mp.Pool(processes=mp.cpu_count()) 82 | if split != 'test': 83 | p.map(newf, files) 84 | else: 85 | p.map(newf_test, files) 86 | p.close() 87 | p.join() 88 | -------------------------------------------------------------------------------- /util/voxelize.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | import numpy as np 7 | import torch 8 | 9 | 10 | def fnv_hash_vec(arr): 11 | """ 12 | FNV64-1A 13 | """ 14 | assert arr.ndim == 2 15 | # Floor first for negative coordinates 16 | arr = arr.copy() 17 | arr = arr.astype(np.uint64, copy=False) 18 | hashed_arr = np.uint64(14695981039346656037) * np.ones(arr.shape[0], dtype=np.uint64) 19 | for j in range(arr.shape[1]): 20 | hashed_arr *= np.uint64(1099511628211) 21 | hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j]) 22 | return hashed_arr 23 | 24 | 25 | def ravel_hash_vec(arr): 26 | """ 27 | Ravel the coordinates after subtracting the min coordinates. 28 | """ 29 | assert arr.ndim == 2 30 | arr = arr.copy() 31 | arr -= arr.min(0) 32 | arr = arr.astype(np.uint64, copy=False) 33 | arr_max = arr.max(0).astype(np.uint64) + 1 34 | 35 | keys = np.zeros(arr.shape[0], dtype=np.uint64) 36 | # Fortran style indexing 37 | for j in range(arr.shape[1] - 1): 38 | keys += arr[:, j] 39 | keys *= arr_max[j + 1] 40 | keys += arr[:, -1] 41 | return keys 42 | 43 | 44 | def voxelize(coord, voxel_size=0.05, hash_type='fnv', mode='random'): 45 | ''' 46 | Voxelization of the input coordinates 47 | Parameters: 48 | coord: input coordinates (N x D) 49 | voxel_size: Size of the voxels 50 | hash_type: Type of the hashing function, can be chosen from 'ravel' and 'fnv' 51 | mode: 'random', 'deterministic' or 'multiple' mode. In training mode one selects a random point within the voxel as the representation of the voxel. 52 | In deterministic model right now one always uses the first point. Usually random mode is preferred for training. In 'multiple' mode, we will return 53 | multiple sets of indices, so that each point will be covered in at least one of these sets 54 | Returns: 55 | idx_unique: the indices of the points so that there is at most one point for each voxel 56 | ''' 57 | discrete_coord = np.floor(coord / np.array(voxel_size)) 58 | if hash_type == 'ravel': 59 | key = ravel_hash_vec(discrete_coord) 60 | else: 61 | key = fnv_hash_vec(discrete_coord) 62 | 63 | idx_sort = np.argsort(key) 64 | key_sort = key[idx_sort] 65 | _, count = np.unique(key_sort, return_counts=True) 66 | if mode == 'deterministic': 67 | # idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + torch.randint(count.max(), (count.size,)).numpy() % count 68 | idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.zeros((count.size,), dtype=np.int32) 69 | idx_unique = idx_sort[idx_select] 70 | return idx_unique 71 | elif mode == 'multiple': # mode is 'multiple' 72 | idx_data = [] 73 | for i in range(count.max()): 74 | idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + i % count 75 | idx_part = idx_sort[idx_select] 76 | idx_data.append(idx_part) 77 | return idx_data 78 | else: # mode == 'random' 79 | # idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(), count.size) % count 80 | idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + torch.randint(count.max(), (count.size,)).numpy() % count 81 | idx_unique = idx_sort[idx_select] 82 | return idx_unique 83 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the open source team at [opensource-conduct@group.apple.com](mailto:opensource-conduct@group.apple.com). All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org), version 1.4, 71 | available at [https://www.contributor-covenant.org/version/1/4/code-of-conduct.html](https://www.contributor-covenant.org/version/1/4/code-of-conduct.html) -------------------------------------------------------------------------------- /util/logger.py: -------------------------------------------------------------------------------- 1 | ## Adapted from StratifiedTransformer 2 | ## https://github.com/dvlab-research/Stratified-Transformer 3 | ## Copyright @ 2022 DV Lab 4 | ## MIT License, https://github.com/dvlab-research/Stratified-Transformer/blob/main/LICENSE.md 5 | 6 | import os 7 | import logging 8 | import functools 9 | import sys 10 | from termcolor import colored 11 | 12 | class _ColorfulFormatter(logging.Formatter): 13 | def __init__(self, *args, **kwargs): 14 | self._root_name = kwargs.pop("root_name") + "." 15 | self._abbrev_name = kwargs.pop("abbrev_name", "") 16 | if len(self._abbrev_name): 17 | self._abbrev_name = self._abbrev_name + "." 18 | super(_ColorfulFormatter, self).__init__(*args, **kwargs) 19 | 20 | def formatMessage(self, record): 21 | record.name = record.name.replace(self._root_name, self._abbrev_name) 22 | log = super(_ColorfulFormatter, self).formatMessage(record) 23 | if record.levelno == logging.WARNING: 24 | prefix = colored("WARNING", "red", attrs=["blink"]) 25 | elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL: 26 | prefix = colored("ERROR", "red", attrs=["blink", "underline"]) 27 | else: 28 | return log 29 | return prefix + " " + log 30 | 31 | 32 | # so that calling setup_logger multiple times won't add many handlers 33 | @functools.lru_cache() 34 | def get_logger( 35 | output=None, color=True, name="main-logger", abbrev_name=None 36 | ): 37 | """ 38 | Initialize the detectron2 logger and set its verbosity level to "INFO". 39 | Args: 40 | output (str): a file name or a directory to save log. If None, will not save log file. 41 | If ends with ".txt" or ".log", assumed to be a file name. 42 | Otherwise, logs will be saved to `output/log.txt`. 43 | name (str): the root module name of this logger 44 | Returns: 45 | logging.Logger: a logger 46 | """ 47 | logger = logging.getLogger(name) 48 | logger.setLevel(logging.DEBUG) 49 | logger.propagate = False 50 | 51 | if abbrev_name is None: 52 | abbrev_name = name 53 | 54 | plain_formatter = logging.Formatter( 55 | "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S" 56 | ) 57 | # stdout logging: master only 58 | ch = logging.StreamHandler(stream=sys.stdout) 59 | ch.setLevel(logging.DEBUG) 60 | if color: 61 | formatter = _ColorfulFormatter( 62 | colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s", 63 | datefmt="%m/%d %H:%M:%S", 64 | root_name=name, 65 | abbrev_name=str(abbrev_name), 66 | ) 67 | else: 68 | formatter = plain_formatter 69 | ch.setFormatter(formatter) 70 | logger.addHandler(ch) 71 | 72 | # file logging: also master only 73 | if output is not None: 74 | if output.endswith(".txt") or output.endswith(".log"): 75 | filename = output 76 | else: 77 | filename = os.path.join(output, "log.txt") 78 | os.makedirs(os.path.dirname(filename), exist_ok=True) 79 | 80 | fh = logging.StreamHandler(_cached_log_stream(filename)) 81 | fh.setLevel(logging.DEBUG) 82 | fh.setFormatter(plain_formatter) 83 | logger.addHandler(fh) 84 | 85 | return logger 86 | 87 | # cache the opened file object, so that different calls to `setup_logger` 88 | # with the same file name can safely write to the same file. 89 | @functools.lru_cache(maxsize=None) 90 | def _cached_log_stream(filename): 91 | return open(filename, "a") -------------------------------------------------------------------------------- /cpp_wrappers/cpp_pcf_kernel/pcf_cuda.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // For licensing see accompanying LICENSE file. 3 | // Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | // 5 | #include 6 | 7 | #include 8 | 9 | // TODO: For now, not attempting to fuse the next downsampling linear layer because hand-written matmul 10 | // cannot compare with optimized gemm kernels. May explore using CUTLASS for that at a later time. 11 | torch::Tensor pcf_cuda_forward( 12 | torch::Tensor input, 13 | torch::Tensor neighbor_inds, 14 | torch::Tensor guidance, 15 | torch::Tensor weights 16 | ); 17 | 18 | // Need a version for PointConv too 19 | torch::Tensor pconv_cuda_forward( 20 | torch::Tensor input, 21 | torch::Tensor neighbor_inds, 22 | torch::Tensor weights, 23 | torch::Tensor additional_features); 24 | 25 | std::vector pconv_cuda_backward( 26 | torch::Tensor grad_output, 27 | torch::Tensor input, 28 | torch::Tensor neighbor_inds, 29 | torch::Tensor weights, 30 | torch::Tensor additional_features); 31 | 32 | std::vector pcf_cuda_backward( 33 | torch::Tensor grad_output, 34 | torch::Tensor input, 35 | torch::Tensor neighbor_inds, 36 | torch::Tensor guidance, 37 | torch::Tensor weights); 38 | 39 | #define CHECK_CUDA(x) {TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")} 40 | #define CHECK_CONTIGUOUS(x) {TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")} 41 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 42 | 43 | torch::Tensor pconv_forward( 44 | torch::Tensor input, 45 | torch::Tensor neighbor_inds, 46 | torch::Tensor weights, 47 | torch::Tensor additional_features 48 | ) 49 | { 50 | CHECK_INPUT(input); 51 | CHECK_INPUT(neighbor_inds); 52 | CHECK_INPUT(weights); 53 | CHECK_INPUT(additional_features); 54 | 55 | return pconv_cuda_forward(input, neighbor_inds, weights, additional_features); 56 | } 57 | 58 | std::vector pconv_backward( 59 | torch::Tensor grad_output, 60 | torch::Tensor input, 61 | torch::Tensor neighbor_inds, 62 | torch::Tensor weights, 63 | torch::Tensor additional_features) 64 | { 65 | CHECK_INPUT(grad_output); 66 | CHECK_INPUT(neighbor_inds); 67 | CHECK_INPUT(weights); 68 | CHECK_INPUT(additional_features); 69 | return pconv_cuda_backward(grad_output,input, neighbor_inds, weights,additional_features); 70 | } 71 | 72 | 73 | torch::Tensor pcf_forward( 74 | torch::Tensor input, 75 | torch::Tensor neighbor_inds, 76 | torch::Tensor guidance, 77 | torch::Tensor weights 78 | ) 79 | { 80 | CHECK_INPUT(input); 81 | CHECK_INPUT(neighbor_inds); 82 | CHECK_INPUT(guidance); 83 | CHECK_INPUT(weights); 84 | return pcf_cuda_forward(input, neighbor_inds, guidance, weights); 85 | } 86 | 87 | std::vector pcf_backward( 88 | torch::Tensor grad_output, 89 | torch::Tensor input, 90 | torch::Tensor neighbor_inds, 91 | torch::Tensor guidance, 92 | torch::Tensor weights) 93 | { 94 | CHECK_INPUT(grad_output); 95 | CHECK_INPUT(neighbor_inds); 96 | CHECK_INPUT(guidance); 97 | CHECK_INPUT(weights); 98 | return pcf_cuda_backward(grad_output,input, neighbor_inds, guidance, weights); 99 | } 100 | 101 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 102 | m.def("pcf_forward", &pcf_forward, "PCF forward (CUDA)"); 103 | m.def("pcf_backward", &pcf_backward, "PCF backward (CUDA)"); 104 | m.def("pconv_forward", &pconv_forward, "PointConv forward (CUDA)"); 105 | m.def("pconv_backward", &pconv_backward, "PointConv backward (CUDA)"); 106 | } 107 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | #include "grid_subsampling.h" 7 | 8 | 9 | void grid_subsampling(vector& original_points, 10 | vector& subsampled_points, 11 | vector& original_features, 12 | vector& subsampled_features, 13 | vector& original_classes, 14 | vector& subsampled_classes, 15 | float sampleDl, 16 | int verbose) { 17 | 18 | // Initiate variables 19 | // ****************** 20 | 21 | // Number of points in the cloud 22 | size_t N = original_points.size(); 23 | 24 | // Dimension of the features 25 | size_t fdim = original_features.size() / N; 26 | size_t ldim = original_classes.size() / N; 27 | 28 | // Limits of the cloud 29 | PointXYZ minCorner = min_point(original_points); 30 | PointXYZ maxCorner = max_point(original_points); 31 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 32 | 33 | // Dimensions of the grid 34 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 35 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 36 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 37 | 38 | // Check if features and classes need to be processed 39 | bool use_feature = original_features.size() > 0; 40 | bool use_classes = original_classes.size() > 0; 41 | 42 | 43 | // Create the sampled map 44 | // ********************** 45 | 46 | // Verbose parameters 47 | int i = 0; 48 | int nDisp = N / 100; 49 | 50 | // Initiate variables 51 | size_t iX, iY, iZ, mapIdx; 52 | unordered_map data; 53 | 54 | for (auto& p : original_points) 55 | { 56 | // Position of point in sample map 57 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 58 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 59 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 60 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 61 | 62 | // If not already created, create key 63 | if (data.count(mapIdx) < 1) 64 | data.emplace(mapIdx, SampledData(fdim, ldim)); 65 | 66 | // Fill the sample map 67 | if (use_feature && use_classes) 68 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 69 | else if (use_feature) 70 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 71 | else if (use_classes) 72 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 73 | else 74 | data[mapIdx].update_points(p); 75 | 76 | // Display 77 | i++; 78 | if (verbose > 1 && i%nDisp == 0) 79 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 80 | 81 | } 82 | 83 | // Divide for barycentre and transfer to a vector 84 | subsampled_points.reserve(data.size()); 85 | if (use_feature) 86 | subsampled_features.reserve(data.size() * fdim); 87 | if (use_classes) 88 | subsampled_classes.reserve(data.size() * ldim); 89 | for (auto& v : data) 90 | { 91 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 92 | if (use_feature) 93 | { 94 | float count = (float)v.second.count; 95 | transform(v.second.features.begin(), 96 | v.second.features.end(), 97 | v.second.features.begin(), 98 | [count](float f) { return f / count;}); 99 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 100 | } 101 | if (use_classes) 102 | { 103 | for (int i = 0; i < ldim; i++) 104 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 105 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 106 | } 107 | } 108 | 109 | return; 110 | } 111 | -------------------------------------------------------------------------------- /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 | struct PointCloud 150 | { 151 | 152 | std::vector pts; 153 | 154 | // Must return the number of data points 155 | inline size_t kdtree_get_point_count() const { return pts.size(); } 156 | 157 | // Returns the dim'th component of the idx'th point in the class: 158 | // Since this is inlined and the "dim" argument is typically an immediate value, the 159 | // "if/else's" are actually solved at compile time. 160 | inline float kdtree_get_pt(const size_t idx, const size_t dim) const 161 | { 162 | if (dim == 0) return pts[idx].x; 163 | else if (dim == 1) return pts[idx].y; 164 | else return pts[idx].z; 165 | } 166 | 167 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 168 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 169 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 170 | template 171 | bool kdtree_get_bbox(BBOX& /* bb */) const { return false; } 172 | 173 | }; 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PointConvFormer 2 | *Wenxuan Wu, Li Fuxin, Qi Shan* 3 | 4 | This is the PyTorch implementation of our paper [**PointConvFormer**] 5 |
6 | 7 |
8 | 9 | ## Introduction 10 | 11 | We introduce PointConvFormer, a novel building block for point cloud based deep network architectures. Inspired 12 | by generalization theory, PointConvFormer combines ideas from point convolution, where filter weights are only based 13 | on relative position, and Transformers which utilize feature-based attention. In PointConvFormer, attention computed 14 | from feature difference between points in the neighborhood is used to modify the convolutional weights at each point. 15 | Hence, we preserved the invariances from point convolution, whereas attention helps to select relevant points in the 16 | neighborhood for convolution. PointConvFormer is suitable for multiple tasks that require details at the point level, such 17 | as segmentation and scene flow estimation tasks. Our results show that PointConvFormer offers a better accuracy-speed 18 | tradeoff than classic convolutions, regular transformers, and voxelized sparse convolution approaches. Visualizations 19 | show that PointConvFormer performs similarly to convolution on flat areas, whereas the neighborhood selection effect is stronger on object boundaries, showing that it has got the best of both worlds. 20 | 21 | ## Highlight 22 | 1. We introduce PointConvFormer which modifies convolution by an attention weight computed from the differences of local neighbourhood features. We further extend the PointConvFormer with a multi-head mechanism. 23 | 2. We conduct thorough experiments on semantic segmentation tasks for both indoor and outdoor scenes, as well as scene flow estimation from 3D point clouds on multiple datasets. Extensive ablation studies are conducted to study the properties and design choice of PointConvFormer. 24 | 25 | ## Citation 26 | Please cite it as W. Wu, L. Fuxin, Q. Shan. PointConvFormer: Revenge of the Point-Cloud Convolution. CVPR 2023 27 | 28 | ## Installation 29 | 30 | ### Environment 31 | The minimal GPU requirement is GTX 1050 due to the use of CUDA compute capabilities 6. 32 | 33 | 1. Install dependencies 34 | 35 | One would need a working PyTorch installed. This code was tested to work in pyTorch 1.8 and 1.12, with CUDA 10.1 and 11.3 respectively. 36 | 37 | One possibility is to run: 38 | ``` 39 | conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.1 -c pytorch 40 | ``` 41 | However usually you should already have an environment that takes care of pyTorch. 42 | 43 | Afterwards, running setup.sh will download all the prepared ScanNet dataset as well as installing the required packages. 44 | 45 | ``` 46 | setup.sh 47 | ``` 48 | 49 | ### Data Preparation 50 | 51 | ### Training 52 | 53 | 1. Before training, please setup the `train_data_path` and `val_data_path` in `configPCF_10cm.yaml`; 54 | 55 | 2. You might also want to set the `model_name`, `experiment_dir` accordingly in `configPCF_10cm.yaml`; 56 | 57 | 4. Change other settings in `configFLPCF_10cm.yaml` based on your experiments; 58 | 59 | 5. Make sure you have the same number of CPUs as num_workers in the config file (e.g. configPCF_10cm.yaml). 60 | 61 | 6. Run ```sh run_distributed.sh num_gpus config_file_name``` to train the model with num_gpus GPUs, or python train_ScanNet_DDP_WarmUP.py --config configPCF_10cm.yaml to train the model with a single GPU. 62 | 63 | ### Evaluation 64 | 65 | 66 | 67 | Then, you can evaluate with the following comand: 68 | 69 | ``` 70 | python test_ScanNet_simple.py --config ./configPCF_10cm.yaml --pretrain_path ./pretrain/[model_weights].pth 71 | ``` 72 | which will evaluate the time as well as outputting ply files for each scene. 73 | 74 | ### Configuration Files and Default Models 75 | 76 | We have provided a few sample configuration files for the ScanNet dataset. configPCF_10cm.yaml is a sample 10cm configuration file 77 | and configPCF_10cm_lite.yaml corresponds to the lightweight version in the paper. Similar configuration files were provided for the 5cm 78 | configuration and a configuration file configPCF_2cm_PTF2.yaml is also provided utilizing some of the ideas from PointTransformerv2 79 | (grid size, mix3D augmentation), which is cheaper than the main one reported in the paper yet still achieves a comparable 74.4% mIOU on the ScanNet validation set. 80 | 81 | Besides, in model_architecture.py we have provided a few default models such as PCF_Tiny, PCF_Small, PCF_Normal and PCF_Large, which can serve for different scenarios. 82 | -------------------------------------------------------------------------------- /data_preparation/scannetv2_val.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_pcf_kernel/test.py: -------------------------------------------------------------------------------- 1 | # PCF CUDA Kernel: 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | # Unit tests for the CUDA Kernel, saving the code here to ease future improvements on the unit tests 7 | 8 | import time 9 | import torch 10 | import pcf_cuda 11 | 12 | 13 | def index_points(points, idx): 14 | """ 15 | 16 | Input: 17 | points: input points data, shape [B, N, C] 18 | idx: sample index data, shape [B, S] / [B, S, K] 19 | Return: 20 | new_points:, indexed points data, shape [B, S, C] / [B, S, K, C] 21 | """ 22 | device = points.device 23 | BB = points.shape[0] 24 | view_shape = list(idx.shape) 25 | view_shape[1:] = [1] * (len(view_shape) - 1) 26 | repeat_shape = list(idx.shape) 27 | repeat_shape[0] = 1 28 | batch_indices = torch.arange(BB, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) 29 | new_points = points[batch_indices, idx, :] 30 | return new_points 31 | 32 | 33 | data = torch.load('test_files.pth') 34 | guidance = data['guidance'] 35 | guidance = guidance.permute(0, 3, 1, 2).contiguous() 36 | linear_weight = data['linear_weight'][:, :768].contiguous() 37 | weightnet = data['weightnet_out'].permute(0, 3, 2, 1).contiguous() 38 | input_feat = data['layer1_out'].contiguous().to('cuda:0') 39 | neighbor_inds = data['neighbor_inds'].to('cuda:0') 40 | weightnet = weightnet.to('cuda:0') 41 | linear_weight = linear_weight.to('cuda:0') 42 | guidance = guidance.to('cuda:0') 43 | 44 | forward = 0 45 | backward = 0 46 | forward_pconv = 0 47 | reps = 1000 48 | discard = 50 49 | 50 | print(input_feat.shape) 51 | print(neighbor_inds.shape) 52 | print(guidance.shape) 53 | print(weightnet.shape) 54 | weightnet2 = weightnet.permute(0, 1, 3, 2).contiguous() 55 | gathered_feat = index_points(input_feat, neighbor_inds) 56 | for i in range(reps+discard): 57 | torch.cuda.synchronize() 58 | start = time.time() 59 | output = pcf_cuda.pcf_forward(input_feat, neighbor_inds, guidance, weightnet2) 60 | # ,linear_weight, torch.zeros(1,dtype=torch.float32).to('cuda:0')) 61 | torch.cuda.synchronize() 62 | end = time.time() 63 | if i > discard: 64 | forward += end - start 65 | start = end 66 | grad_input, grad_guidance, grad_weights = pcf_cuda.pcf_backward(torch.ones_like(output), input_feat, neighbor_inds, guidance, weightnet2) 67 | torch.cuda.synchronize() 68 | end = time.time() 69 | if i > discard: 70 | backward += end - start 71 | start = end 72 | output_pconv = pcf_cuda.pconv_forward(input_feat, neighbor_inds, weightnet2, gathered_feat) 73 | torch.cuda.synchronize() 74 | end = time.time() 75 | if i > discard: 76 | forward_pconv += end - start 77 | print(output.shape) 78 | print('CUDA Kernel Forward: ', forward / reps * 1000, 'ms') 79 | print('CUDA Kernel Backward: ', backward / reps * 1000, 'ms') 80 | print('CUDA Kernel Forward Pconv: ', forward_pconv / reps * 1000, 'ms') 81 | 82 | print('Input grad shape: ', grad_input.shape) 83 | print('Guidance grad shape: ', grad_guidance.shape) 84 | print('Weights grad shape: ', grad_weights.shape) 85 | 86 | forward = 0 87 | forward1 = 0 88 | forward2 = 0 89 | backward1 = 0 90 | B = input_feat.shape[0] 91 | K = neighbor_inds.shape[2] 92 | M = neighbor_inds.shape[1] 93 | num_heads = guidance.shape[2] 94 | guidance = guidance.permute(0, 2, 3, 1).contiguous() 95 | # weightnet = weightnet.permute(0,1,3,2).contiguous() 96 | input_feat.requires_grad_(True) 97 | input_feat.retain_grad() 98 | guidance.requires_grad_(True) 99 | weightnet.requires_grad_(True) 100 | guidance.retain_grad() 101 | weightnet.retain_grad() 102 | for i in range(reps+discard): 103 | torch.cuda.synchronize() 104 | start = time.time() 105 | gathered_feat = index_points(input_feat, neighbor_inds) 106 | new_feat = gathered_feat.permute(0, 3, 2, 1) 107 | new_feat = new_feat.view(B, -1, num_heads, K, M) 108 | new_feat_inm = (new_feat * guidance).view(B, -1, K, M).contiguous() 109 | new_feat2 = torch.matmul(input=new_feat_inm.permute(0, 3, 1, 2).contiguous(), other=weightnet) 110 | new_feat2 = new_feat2.view(new_feat2.shape[0], new_feat2.shape[1], -1) 111 | torch.cuda.synchronize() 112 | t1 = time.time() 113 | if i > discard: 114 | forward1 += t1 - start 115 | fin_func = torch.sum(new_feat2) 116 | fin_func.backward() 117 | torch.cuda.synchronize() 118 | t2 = time.time() 119 | gathered_feat2 = torch.cat([gathered_feat, gathered_feat], dim=-1) 120 | new_pconv = torch.matmul(input=gathered_feat2.permute(0, 1, 3, 2).contiguous(), other=weightnet) 121 | new_pconv = new_pconv.view(new_pconv.shape[0], new_pconv.shape[1], -1) 122 | torch.cuda.synchronize() 123 | t3 = time.time() 124 | grad_input2 = input_feat.grad.clone().detach() 125 | grad_guidance2 = guidance.grad.clone().detach() 126 | grad_weights2 = weightnet.grad.clone().detach() 127 | guidance.grad.data.zero_() 128 | weightnet.grad.data.zero_() 129 | input_feat.grad.data.zero_() 130 | if i > discard: 131 | backward1 += t2 - t1 132 | forward2 += t3 - t2 133 | grad_guidance2 = grad_guidance2.permute(0, 3, 1, 2) 134 | # Shape between C_mid and K is different between the 2 versions 135 | grad_weights = grad_weights.permute(0, 1, 3, 2) 136 | print('input grad shape: ', grad_input2.shape) 137 | print('guidance grad shape: ', grad_guidance2.shape) 138 | print('weight grad shape: ', grad_weights2.shape) 139 | 140 | print('Pytorch Forward1: ', forward1 / reps * 1000, 'ms') 141 | print('Pytorch backward1: ', backward1 / reps * 1000, 'ms') 142 | print('Pytorch Forward Pconv matmul only: ', forward2 / reps * 1000, 'ms') 143 | print('diff of forward: ', torch.linalg.norm(new_feat2 - output)) 144 | print('diff of pconv: ', torch.linalg.norm(new_pconv - output_pconv)) 145 | print('diff of input grad: ', torch.linalg.norm(grad_input2 - grad_input)) 146 | print('diff of guidance grad: ', torch.linalg.norm(grad_guidance2 - grad_guidance)) 147 | print('diff of weight grad: ', torch.linalg.norm(grad_weights2 - grad_weights)) 148 | -------------------------------------------------------------------------------- /ACKNOWLEDGEMENTS: -------------------------------------------------------------------------------- 1 | Acknowledgements 2 | Portions of this AutoFocusFormer Software may utilize the following copyrighted 3 | material, the use of which is hereby acknowledged. 4 | 5 | ____________________ 6 | 7 | KPConv-PyTorch 8 | 9 | MIT License 10 | 11 | Copyright (c) 2019 HuguesTHOMAS 12 | 13 | Permission is hereby granted, free of charge, to any person obtaining a copy 14 | of this software and associated documentation files (the "Software"), to deal 15 | in the Software without restriction, including without limitation the rights 16 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 17 | copies of the Software, and to permit persons to whom the Software is 18 | furnished to do so, subject to the following conditions: 19 | 20 | The above copyright notice and this permission notice shall be included in all 21 | copies or substantial portions of the Software. 22 | 23 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 24 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 25 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 26 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 27 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 28 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 29 | SOFTWARE. 30 | 31 | ____________________ 32 | 33 | StratifiedFormer 34 | 35 | MIT License 36 | 37 | Copyright (c) 2022 DV Lab 38 | 39 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 40 | 41 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 42 | 43 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 44 | 45 | ____________________ 46 | 47 | pytorch-checkpoint 48 | 49 | MIT License 50 | 51 | Copyright (c) 2018 Huiyu Wang 52 | 53 | Permission is hereby granted, free of charge, to any person obtaining a copy 54 | of this software and associated documentation files (the "Software"), to deal 55 | in the Software without restriction, including without limitation the rights 56 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 57 | copies of the Software, and to permit persons to whom the Software is 58 | furnished to do so, subject to the following conditions: 59 | 60 | The above copyright notice and this permission notice shall be included in all 61 | copies or substantial portions of the Software. 62 | 63 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 64 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 65 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 66 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 67 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 68 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 69 | SOFTWARE. 70 | 71 | ____________________ 72 | 73 | VI_PointConv 74 | 75 | The Devils in the Point Clouds: Studying the Robustness of Point Cloud 76 | Convolutions. 77 | 78 | Copyright (c) 2021, Deep Machine Vision group, Oregon State University 79 | 80 | The MIT License (MIT) 81 | 82 | Copyright (c) 2020 Xingyi Li 83 | 84 | Permission is hereby granted, free of charge, to any person obtaining a copy 85 | of this software and associated documentation files (the "Software"), to deal 86 | in the Software without restriction, including without limitation the rights 87 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 88 | copies of the Software, and to permit persons to whom the Software is 89 | furnished to do so, subject to the following conditions: 90 | 91 | The above copyright notice and this permission notice shall be included in all 92 | copies or substantial portions of the Software. 93 | 94 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 95 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 96 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 97 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 98 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 99 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 100 | SOFTWARE. 101 | 102 | ____________________ 103 | 104 | nanoflann 105 | 106 | Software License Agreement (BSD License) 107 | 108 | Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved. 109 | Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved. 110 | Copyright 2011 Jose L. Blanco (joseluisblancoc@gmail.com). All rights reserved. 111 | 112 | THE BSD LICENSE 113 | 114 | Redistribution and use in source and binary forms, with or without 115 | modification, are permitted provided that the following conditions 116 | are met: 117 | 118 | 1. Redistributions of source code must retain the above copyright 119 | notice, this list of conditions and the following disclaimer. 120 | 2. Redistributions in binary form must reproduce the above copyright 121 | notice, this list of conditions and the following disclaimer in the 122 | documentation and/or other materials provided with the distribution. 123 | 124 | THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR 125 | IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES 126 | OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 127 | IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 128 | INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT 129 | NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 130 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 131 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 132 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 133 | THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 134 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | #include 7 | #include 8 | #include "grid_subsampling/grid_subsampling.h" 9 | #include 10 | 11 | 12 | 13 | // docstrings for our module 14 | // ************************* 15 | 16 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 17 | 18 | static char compute_docstring[] = "function subsampling a pointcloud"; 19 | 20 | 21 | // Declare the functions 22 | // ********************* 23 | 24 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 25 | 26 | 27 | // Specify the members of the module 28 | // ********************************* 29 | 30 | static PyMethodDef module_methods[] = 31 | { 32 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 33 | {NULL, NULL, 0, NULL} 34 | }; 35 | 36 | 37 | // Initialize the module 38 | // ********************* 39 | 40 | static struct PyModuleDef moduledef = 41 | { 42 | PyModuleDef_HEAD_INIT, 43 | "grid_subsampling", // m_name 44 | module_docstring, // m_doc 45 | -1, // m_size 46 | module_methods, // m_methods 47 | NULL, // m_reload 48 | NULL, // m_traverse 49 | NULL, // m_clear 50 | NULL, // m_free 51 | }; 52 | 53 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 54 | { 55 | import_array(); 56 | return PyModule_Create(&moduledef); 57 | } 58 | 59 | 60 | // Actual wrapper 61 | // ************** 62 | 63 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 64 | { 65 | 66 | // Manage inputs 67 | // ************* 68 | 69 | // Args containers 70 | PyObject *points_obj = NULL; 71 | PyObject *features_obj = NULL; 72 | PyObject *classes_obj = NULL; 73 | 74 | // Keywords containers 75 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 76 | float sampleDl = 0.1; 77 | const char *method_buffer = "barycenters"; 78 | int verbose = 0; 79 | 80 | // Parse the input 81 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 82 | { 83 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 84 | return NULL; 85 | } 86 | 87 | // Get the method argument 88 | string method(method_buffer); 89 | 90 | // Interpret method 91 | if (method.compare("barycenters") && method.compare("voxelcenters")) 92 | { 93 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 94 | return NULL; 95 | } 96 | 97 | // Check if using features or classes 98 | bool use_feature = true, use_classes = true; 99 | if (features_obj == NULL) 100 | use_feature = false; 101 | if (classes_obj == NULL) 102 | use_classes = false; 103 | 104 | // Interpret the input objects as numpy arrays. 105 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 106 | PyObject *features_array = NULL; 107 | PyObject *classes_array = NULL; 108 | if (use_feature) 109 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 110 | if (use_classes) 111 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 112 | 113 | // Verify data was load correctly. 114 | if (points_array == NULL) 115 | { 116 | Py_XDECREF(points_array); 117 | Py_XDECREF(classes_array); 118 | Py_XDECREF(features_array); 119 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 120 | return NULL; 121 | } 122 | if (use_feature && features_array == NULL) 123 | { 124 | Py_XDECREF(points_array); 125 | Py_XDECREF(classes_array); 126 | Py_XDECREF(features_array); 127 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 128 | return NULL; 129 | } 130 | if (use_classes && classes_array == NULL) 131 | { 132 | Py_XDECREF(points_array); 133 | Py_XDECREF(classes_array); 134 | Py_XDECREF(features_array); 135 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 136 | return NULL; 137 | } 138 | 139 | // Check that the input array respect the dims 140 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 141 | { 142 | Py_XDECREF(points_array); 143 | Py_XDECREF(classes_array); 144 | Py_XDECREF(features_array); 145 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 146 | return NULL; 147 | } 148 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 149 | { 150 | Py_XDECREF(points_array); 151 | Py_XDECREF(classes_array); 152 | Py_XDECREF(features_array); 153 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 154 | return NULL; 155 | } 156 | 157 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 158 | { 159 | Py_XDECREF(points_array); 160 | Py_XDECREF(classes_array); 161 | Py_XDECREF(features_array); 162 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 163 | return NULL; 164 | } 165 | 166 | // Number of points 167 | int N = (int)PyArray_DIM(points_array, 0); 168 | 169 | // Dimension of the features 170 | int fdim = 0; 171 | if (use_feature) 172 | fdim = (int)PyArray_DIM(features_array, 1); 173 | 174 | //Dimension of labels 175 | int ldim = 1; 176 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 177 | ldim = (int)PyArray_DIM(classes_array, 1); 178 | 179 | // Check that the input array respect the number of points 180 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 181 | { 182 | Py_XDECREF(points_array); 183 | Py_XDECREF(classes_array); 184 | Py_XDECREF(features_array); 185 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 186 | return NULL; 187 | } 188 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 189 | { 190 | Py_XDECREF(points_array); 191 | Py_XDECREF(classes_array); 192 | Py_XDECREF(features_array); 193 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 194 | return NULL; 195 | } 196 | 197 | 198 | // Call the C++ function 199 | // ********************* 200 | 201 | // Create pyramid 202 | if (verbose > 0) 203 | cout << "Computing cloud pyramid with support points: " << endl; 204 | 205 | 206 | // Convert PyArray to Cloud C++ class 207 | vector original_points; 208 | vector original_features; 209 | vector original_classes; 210 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 211 | if (use_feature) 212 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 213 | if (use_classes) 214 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 215 | 216 | // Subsample 217 | vector subsampled_points; 218 | vector subsampled_features; 219 | vector subsampled_classes; 220 | grid_subsampling(original_points, 221 | subsampled_points, 222 | original_features, 223 | subsampled_features, 224 | original_classes, 225 | subsampled_classes, 226 | sampleDl, 227 | verbose); 228 | 229 | // Check result 230 | if (subsampled_points.size() < 1) 231 | { 232 | PyErr_SetString(PyExc_RuntimeError, "Error"); 233 | return NULL; 234 | } 235 | 236 | // Manage outputs 237 | // ************** 238 | 239 | // Dimension of input containers 240 | npy_intp* point_dims = new npy_intp[2]; 241 | point_dims[0] = subsampled_points.size(); 242 | point_dims[1] = 3; 243 | npy_intp* feature_dims = new npy_intp[2]; 244 | feature_dims[0] = subsampled_points.size(); 245 | feature_dims[1] = fdim; 246 | npy_intp* classes_dims = new npy_intp[2]; 247 | classes_dims[0] = subsampled_points.size(); 248 | classes_dims[1] = ldim; 249 | 250 | // Create output array 251 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 252 | PyObject *res_features_obj = NULL; 253 | PyObject *res_classes_obj = NULL; 254 | PyObject *ret = NULL; 255 | 256 | // Fill output array with values 257 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 258 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 259 | if (use_feature) 260 | { 261 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 262 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 263 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 264 | } 265 | if (use_classes) 266 | { 267 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 268 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 269 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 270 | } 271 | 272 | 273 | // Merge results 274 | if (use_feature && use_classes) 275 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 276 | else if (use_feature) 277 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 278 | else if (use_classes) 279 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 280 | else 281 | ret = Py_BuildValue("N", res_points_obj); 282 | 283 | // Clean up 284 | // ******** 285 | 286 | Py_DECREF(points_array); 287 | Py_XDECREF(features_array); 288 | Py_XDECREF(classes_array); 289 | 290 | return ret; 291 | } -------------------------------------------------------------------------------- /util/common_util.py: -------------------------------------------------------------------------------- 1 | ## Partly adapted from StratifiedTransformer 2 | ## https://github.com/dvlab-research/Stratified-Transformer 3 | ## Copyright @ 2022 DV Lab 4 | ## MIT License, https://github.com/dvlab-research/Stratified-Transformer/blob/main/LICENSE.md 5 | 6 | # For to_device() and replace_batchnorm(): 7 | # For licensing see accompanying LICENSE file. 8 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 9 | ## 10 | 11 | import os 12 | import numpy as np 13 | from PIL import Image 14 | import random 15 | 16 | import torch 17 | from torch import nn 18 | from torch.nn.modules.conv import _ConvNd 19 | from torch.nn.modules.batchnorm import _BatchNorm 20 | import torch.nn.init as initer 21 | import torch.nn.functional as F 22 | 23 | 24 | class AverageMeter(object): 25 | """Computes and stores the average and current value""" 26 | def __init__(self): 27 | self.reset() 28 | 29 | def reset(self): 30 | self.val = 0 31 | self.avg = 0 32 | self.sum = 0 33 | self.count = 0 34 | 35 | def update(self, val, n=1): 36 | self.val = val 37 | self.sum += val * n 38 | self.count += n 39 | self.avg = self.sum / self.count 40 | 41 | 42 | def step_learning_rate(optimizer, base_lr, epoch, step_epoch, multiplier=0.1, clip=1e-6): 43 | """Sets the learning rate to the base LR decayed by 10 every step epochs""" 44 | lr = max(base_lr * (multiplier ** (epoch // step_epoch)), clip) 45 | for param_group in optimizer.param_groups: 46 | param_group['lr'] = lr 47 | 48 | 49 | def poly_learning_rate(optimizer, base_lr, curr_iter, max_iter, power=0.9): 50 | """poly learning rate policy""" 51 | lr = base_lr * (1 - float(curr_iter) / max_iter) ** power 52 | for param_group in optimizer.param_groups: 53 | param_group['lr'] = lr 54 | 55 | 56 | def intersectionAndUnion(output, target, K, ignore_index=255): 57 | # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. 58 | assert (output.ndim in [1, 2, 3]) 59 | assert output.shape == target.shape 60 | # output = output.reshape(output.size()).copy() 61 | # target = target.reshape(target.size()) 62 | output = output.flatten() 63 | target = target.flatten() 64 | output[np.where(target == ignore_index)[0]] = ignore_index 65 | intersection = output[np.where(output == target)[0]] 66 | area_intersection, _ = np.histogram(intersection, bins=np.arange(K+1)) 67 | area_output, _ = np.histogram(output, bins=np.arange(K+1)) 68 | area_target, _ = np.histogram(target, bins=np.arange(K+1)) 69 | area_union = area_output + area_target - area_intersection 70 | return area_intersection, area_union, area_target 71 | 72 | 73 | def intersectionAndUnionGPU(output, target, K, ignore_index=255): 74 | # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. 75 | assert (output.dim() in [1, 2, 3]) 76 | assert output.shape == target.shape 77 | output = output.view(-1) 78 | target = target.view(-1) 79 | output[target == ignore_index] = ignore_index 80 | intersection = output[output == target] 81 | area_intersection = torch.histc(intersection, bins=K, min=0, max=K-1) 82 | area_output = torch.histc(output, bins=K, min=0, max=K-1) 83 | area_target = torch.histc(target, bins=K, min=0, max=K-1) 84 | area_union = area_output + area_target - area_intersection 85 | return area_intersection, area_union, area_target 86 | 87 | 88 | def check_mkdir(dir_name): 89 | if not os.path.exists(dir_name): 90 | os.mkdir(dir_name) 91 | 92 | 93 | def check_makedirs(dir_name): 94 | if not os.path.exists(dir_name): 95 | os.makedirs(dir_name) 96 | 97 | 98 | def init_weights(model, conv='kaiming', batchnorm='normal', linear='kaiming', lstm='kaiming'): 99 | """ 100 | :param model: Pytorch Model which is nn.Module 101 | :param conv: 'kaiming' or 'xavier' 102 | :param batchnorm: 'normal' or 'constant' 103 | :param linear: 'kaiming' or 'xavier' 104 | :param lstm: 'kaiming' or 'xavier' 105 | """ 106 | for m in model.modules(): 107 | if isinstance(m, (_ConvNd)): 108 | if conv == 'kaiming': 109 | initer.kaiming_normal_(m.weight) 110 | elif conv == 'xavier': 111 | initer.xavier_normal_(m.weight) 112 | else: 113 | raise ValueError("init type of conv error.\n") 114 | if m.bias is not None: 115 | initer.constant_(m.bias, 0) 116 | 117 | elif isinstance(m, _BatchNorm): 118 | if batchnorm == 'normal': 119 | initer.normal_(m.weight, 1.0, 0.02) 120 | elif batchnorm == 'constant': 121 | initer.constant_(m.weight, 1.0) 122 | else: 123 | raise ValueError("init type of batchnorm error.\n") 124 | initer.constant_(m.bias, 0.0) 125 | 126 | elif isinstance(m, nn.Linear): 127 | if linear == 'kaiming': 128 | initer.kaiming_normal_(m.weight) 129 | elif linear == 'xavier': 130 | initer.xavier_normal_(m.weight) 131 | else: 132 | raise ValueError("init type of linear error.\n") 133 | if m.bias is not None: 134 | initer.constant_(m.bias, 0) 135 | 136 | elif isinstance(m, nn.LSTM): 137 | for name, param in m.named_parameters(): 138 | if 'weight' in name: 139 | if lstm == 'kaiming': 140 | initer.kaiming_normal_(param) 141 | elif lstm == 'xavier': 142 | initer.xavier_normal_(param) 143 | else: 144 | raise ValueError("init type of lstm error.\n") 145 | elif 'bias' in name: 146 | initer.constant_(param, 0) 147 | 148 | 149 | def convert_to_syncbn(model): 150 | def recursive_set(cur_module, name, module): 151 | if len(name.split('.')) > 1: 152 | recursive_set(getattr(cur_module, name[:name.find('.')]), name[name.find('.')+1:], module) 153 | else: 154 | setattr(cur_module, name, module) 155 | from lib.sync_bn import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d 156 | for name, m in model.named_modules(): 157 | if isinstance(m, nn.BatchNorm1d): 158 | recursive_set(model, name, SynchronizedBatchNorm1d(m.num_features, m.eps, m.momentum, m.affine)) 159 | elif isinstance(m, nn.BatchNorm2d): 160 | recursive_set(model, name, SynchronizedBatchNorm2d(m.num_features, m.eps, m.momentum, m.affine)) 161 | elif isinstance(m, nn.BatchNorm3d): 162 | recursive_set(model, name, SynchronizedBatchNorm3d(m.num_features, m.eps, m.momentum, m.affine)) 163 | 164 | 165 | def colorize(gray, palette): 166 | # gray: numpy array of the label and 1*3N size list palette 167 | color = Image.fromarray(gray.astype(np.uint8)).convert('P') 168 | color.putpalette(palette) 169 | return color 170 | 171 | 172 | def find_free_port(): 173 | import socket 174 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 175 | # Binding to port 0 will cause the OS to find an available port for us 176 | sock.bind(("", 0)) 177 | port = sock.getsockname()[1] 178 | sock.close() 179 | # NOTE: there is still a chance the port could be taken by other processes. 180 | return port 181 | 182 | 183 | def memory_use(): 184 | BYTES_IN_GB = 1024 ** 3 185 | return 'ALLOCATED: {:>6.3f} ({:>6.3f}) CACHED: {:>6.3f} ({:>6.3f})'.format( 186 | torch.cuda.memory_allocated() / BYTES_IN_GB, 187 | torch.cuda.max_memory_allocated() / BYTES_IN_GB, 188 | torch.cuda.memory_reserved() / BYTES_IN_GB, 189 | torch.cuda.max_memory_reserved() / BYTES_IN_GB, 190 | ) 191 | 192 | 193 | def to_device(input_data, device='cuda', non_blocking=False): 194 | ''' 195 | Move input_data to device. If the data is a list, move everything in the list to device 196 | ''' 197 | if isinstance(input_data, list) and len(input_data) > 0: 198 | if isinstance(input_data[0], list): 199 | for idx in range(len(input_data)): 200 | for idx2 in range(len(input_data[idx])): 201 | input_data[idx][idx2] = input_data[idx][idx2].to(device, non_blocking=non_blocking) 202 | else: 203 | for idx in range(len(input_data)): 204 | input_data[idx] = input_data[idx].to(device, non_blocking=non_blocking) 205 | 206 | if isinstance(input_data, torch.Tensor): 207 | input_data = input_data.to(device, non_blocking=non_blocking) 208 | 209 | return input_data 210 | 211 | 212 | def init_seeds(seed=0, cuda_deterministic=True): 213 | random.seed(seed) 214 | np.random.seed(seed) 215 | torch.manual_seed(seed) 216 | torch.cuda.manual_seed(seed) 217 | torch.cuda.manual_seed_all(seed) 218 | 219 | if cuda_deterministic: 220 | torch.backends.cudnn.deterministic = True 221 | torch.backends.cudnn.benchmark = False 222 | else: 223 | torch.backends.cudnn.deterministic = False 224 | torch.backends.cudnn.benchmark = True 225 | 226 | 227 | def smooth_loss(output, target, eps=0.1): 228 | w = torch.zeros_like(output).scatter(1, target.unsqueeze(1), 1) 229 | w = w * (1 - eps) + (1 - w) * eps / (output.shape[1] - 1) 230 | log_prob = F.log_softmax(output, dim=1) 231 | loss = (-w * log_prob).sum(dim=1).mean() 232 | return loss 233 | 234 | 235 | # Inspired by the LeViT repository 236 | # https://github.com/facebookresearch/LeViT 237 | def replace_batchnorm(net): 238 | ''' 239 | Fuse the batch normalization during inference time 240 | ''' 241 | for child_name, child in net.named_children(): 242 | if hasattr(child, 'fuse'): 243 | setattr(net, child_name, child.fuse()) 244 | elif isinstance(child, torch.nn.BatchNorm2d): 245 | setattr(net, child_name, torch.nn.Identity()) 246 | else: 247 | replace_batchnorm(child) 248 | -------------------------------------------------------------------------------- /test_ScanNet_simple.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | # Testing code without voting, can be used to benchmark the inference runtime 7 | # Cannot be used for benchmarking testing set performance 8 | # Usage: python test_ScanNet_simple.py --config config_file --pretrain_path model --split split 9 | # where config_file is the training config file (e.g. configFLPCF_10cm.yaml, model is the .pth file 10 | # stored from training, and split is the split you want to test on (e.g. 'validation')) 11 | 12 | import os 13 | import time 14 | import argparse 15 | import datetime 16 | import yaml 17 | import numpy as np 18 | from easydict import EasyDict as edict 19 | 20 | import torch 21 | import torch.nn.functional as F 22 | import torch.utils.data 23 | 24 | import open3d as o3d 25 | 26 | from util.common_util import AverageMeter, intersectionAndUnion, replace_batchnorm, to_device, init_seeds 27 | from util.logger import get_logger 28 | 29 | from model_architecture import PointConvFormer_Segmentation as Main_Model 30 | from scannet_data_loader_color_DDP import ScanNetDataset 31 | from datasetCommon import getdataLoader 32 | from train_ScanNet_DDP_WarmUP import get_default_configs 33 | 34 | 35 | def get_parser(): 36 | ''' 37 | Get the arguments of the call. 38 | ''' 39 | parser = argparse.ArgumentParser( 40 | description='PointConvFormer for Semantic Segmentation') 41 | parser.add_argument( 42 | '--config', 43 | default='./configFLPCF_10cm.yaml', 44 | type=str, 45 | help='config file') 46 | parser.add_argument( 47 | '--pretrain_path', 48 | default='./', 49 | type=str, 50 | help='the path of pretrain weights') 51 | parser.add_argument( 52 | '--split', 53 | default='validation', 54 | type=str, 55 | help='the dataset split to be tested on') 56 | 57 | args = parser.parse_args() 58 | assert args.config is not None 59 | cfg = edict(yaml.safe_load(open(args.config, 'r'))) 60 | cfg = get_default_configs(cfg) 61 | cfg.pretrain_path = args.pretrain_path 62 | cfg.config = args.config 63 | cfg.split = args.split 64 | return cfg 65 | 66 | 67 | def collect_fns(data_list): 68 | return data_list 69 | 70 | 71 | MAX_NUM_POINTS = 550000 72 | 73 | 74 | def main_vote(): 75 | global args, logger 76 | args = get_parser() 77 | 78 | file_dir = os.path.join(args.eval_path, 'TIME_%s_SemSeg-' % 79 | (args.model_name) + 80 | str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) 81 | args.file_dir = file_dir 82 | ply_dir = os.path.join(args.file_dir, 'ply') 83 | txt_dir = os.path.join(args.file_dir, 'txt') 84 | prob_dir = os.path.join(args.file_dir, 'prob') 85 | 86 | if not os.path.exists(args.eval_path): 87 | os.makedirs(args.eval_path) 88 | 89 | if not os.path.exists(file_dir): 90 | os.makedirs(file_dir) 91 | 92 | if not os.path.exists(ply_dir): 93 | os.makedirs(ply_dir) 94 | 95 | if not os.path.exists(txt_dir): 96 | os.makedirs(txt_dir) 97 | 98 | if not os.path.exists(prob_dir): 99 | os.makedirs(prob_dir) 100 | 101 | logger = get_logger(file_dir) 102 | logger.info(args) 103 | assert args.num_classes > 1 104 | logger.info("=> creating model ...") 105 | logger.info("Classes: {}".format(args.classes)) 106 | 107 | # get model 108 | model = Main_Model(args) 109 | 110 | logger.info(model) 111 | 112 | init_seeds(args.manual_seed) 113 | 114 | if os.path.isfile(args.pretrain_path): 115 | logger.info("=> loading checkpoint '{}'".format(args.pretrain_path)) 116 | checkpoint = torch.load(args.pretrain_path) 117 | state_dict = checkpoint['state_dict'] 118 | model.load_state_dict(state_dict, strict=True) 119 | logger.info( 120 | "=> loaded checkpoint '{}' (epoch {})".format( 121 | args.pretrain_path, 122 | checkpoint['epoch'])) 123 | args.epoch = checkpoint['epoch'] 124 | else: 125 | raise RuntimeError( 126 | "=> no checkpoint found at '{}'".format( 127 | args.pretrain_path)) 128 | 129 | logger.info( 130 | '>>>>>>>>>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>>>>>>>>>>>') 131 | intersection_meter = AverageMeter() 132 | union_meter = AverageMeter() 133 | target_meter = AverageMeter() 134 | 135 | colormap = np.array(create_color_palette40()) / 255.0 136 | mapper = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 137 | 12, 14, 16, 24, 28, 33, 34, 36, 39]) 138 | 139 | model.eval() 140 | # Fuse the linear layer and batch normalization in the model 141 | replace_batchnorm(model) 142 | model = model.to('cuda') 143 | 144 | torch.cuda.empty_cache() 145 | 146 | output_dict = {} 147 | args.BATCH_SIZE = 1 148 | val_data_loader, val_dataset = getdataLoader( 149 | args, ScanNetDataset, args.split, torch.utils.data.SequentialSampler) 150 | 151 | pcd = o3d.geometry.PointCloud() 152 | time_list = [] 153 | 154 | for idx, data in enumerate(val_data_loader): 155 | scene_name = val_dataset.get_scene_name(idx) 156 | features, pointclouds, edges_self, edges_forward, edges_propagate, label, norms = data 157 | features, pointclouds, edges_self, edges_forward, edges_propagate, label, norms = \ 158 | to_device(features), to_device(pointclouds), \ 159 | to_device(edges_self), to_device(edges_forward), \ 160 | to_device(edges_propagate), to_device(label), to_device(norms) 161 | 162 | with torch.no_grad(): 163 | torch.cuda.synchronize() 164 | st = time.time() 165 | pred = model( 166 | features, 167 | pointclouds, 168 | edges_self, 169 | edges_forward, 170 | edges_propagate, 171 | norms) 172 | torch.cuda.synchronize() 173 | et = time.time() 174 | time_list.append(et - st) 175 | pred = pred.contiguous().view(-1, args.num_classes) 176 | pred = F.softmax(pred, dim=-1) 177 | torch.cuda.empty_cache() 178 | logger.info('Test: {}/{}, {}'.format(idx + 1, 179 | len(val_data_loader), et - st)) 180 | np.save(os.path.join(prob_dir, scene_name + '.npy'), 181 | {'pred': pred, 'target': label, 'xyz': pointclouds[0]}) 182 | 183 | cur_dict = { 184 | 'pred': pred, 185 | 'target': label, 186 | 'xyz': pointclouds[0] 187 | } 188 | output_dict[scene_name] = cur_dict 189 | 190 | pcd = o3d.geometry.PointCloud() 191 | for scene_name, output_data in output_dict.items(): 192 | print('scene_name : ', scene_name) 193 | labels = output_data['target'].view(-1, 1)[:, 0] 194 | 195 | pred = output_data['pred'] 196 | xyz = output_data['xyz'].cpu().numpy() 197 | 198 | output = pred.max(1)[1] 199 | labels = labels.cpu().numpy().astype(np.int32) 200 | output = output.cpu().numpy() 201 | 202 | intersection, union, target = intersectionAndUnion( 203 | output, labels, args.num_classes, args.ignore_label) 204 | 205 | intersection_meter.update(intersection) 206 | union_meter.update(union) 207 | target_meter.update(target) 208 | 209 | output = mapper[output.astype(np.int32)] 210 | pcd.points = o3d.utility.Vector3dVector(xyz[0]) 211 | pcd.colors = o3d.utility.Vector3dVector(colormap[output]) 212 | o3d.io.write_point_cloud(str(ply_dir) + '/' + scene_name, pcd) 213 | 214 | fp = open(str(txt_dir) + '/' + scene_name.replace( 215 | "_vh_clean_2.ply", ".txt"), "w") 216 | for l in output: 217 | fp.write(str(int(l)) + '\n') 218 | fp.close() 219 | 220 | iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) 221 | accuracy_class = intersection_meter.sum / (target_meter.sum + 1e-10) 222 | mIoU = np.mean(iou_class) 223 | mAcc = np.mean(accuracy_class) 224 | allAcc = sum(intersection_meter.sum) / (sum(target_meter.sum) + 1e-10) 225 | logger.info( 226 | 'Val result: mIoU/mAcc/allAcc {:.4f}/{:.4f}/{:.4f}.'.format(mIoU, mAcc, allAcc)) 227 | for i in range(args.num_classes): 228 | logger.info('Class_{} Result: iou/accuracy {:.4f}/{:.4f}.'.format(i, 229 | iou_class[i], accuracy_class[i])) 230 | logger.info('<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<') 231 | 232 | logger.info('Average running time per frame: ', np.mean(time_list)) 233 | 234 | 235 | def create_color_palette40(): 236 | return [ 237 | (0, 0, 0), 238 | (174, 199, 232), # wall 239 | (152, 223, 138), # floor 240 | (31, 119, 180), # cabinet 241 | (255, 187, 120), # bed 242 | (188, 189, 34), # chair 243 | (140, 86, 75), # sofa 244 | (255, 152, 150), # table 245 | (214, 39, 40), # door 246 | (197, 176, 213), # window 247 | (148, 103, 189), # bookshelf 248 | (196, 156, 148), # picture 249 | (23, 190, 207), # counter 250 | (178, 76, 76), 251 | (247, 182, 210), # desk 252 | (66, 188, 102), 253 | (219, 219, 141), # curtain 254 | (140, 57, 197), 255 | (202, 185, 52), 256 | (51, 176, 203), 257 | (200, 54, 131), 258 | (92, 193, 61), 259 | (78, 71, 183), 260 | (172, 114, 82), 261 | (255, 127, 14), # refrigerator 262 | (91, 163, 138), 263 | (153, 98, 156), 264 | (140, 153, 101), 265 | (158, 218, 229), # shower curtain 266 | (100, 125, 154), 267 | (178, 127, 135), 268 | (120, 185, 128), 269 | (146, 111, 194), 270 | (44, 160, 44), # toilet 271 | (112, 128, 144), # sink 272 | (96, 207, 209), 273 | (227, 119, 194), # bathtub 274 | (213, 92, 176), 275 | (94, 106, 211), 276 | (82, 84, 163), # otherfurn 277 | (100, 85, 144) 278 | ] 279 | 280 | 281 | if __name__ == '__main__': 282 | main_vote() 283 | -------------------------------------------------------------------------------- /layer_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import pcf_cuda 10 | from util.cp_batchnorm import CpBatchNorm2d 11 | 12 | 13 | def index_points(points, idx): 14 | """ 15 | Input: 16 | points: input points data, shape [B, N, C] 17 | idx: sample index data, shape [B, S] / [B, S, K] 18 | Return: 19 | new_points:, indexed points data, shape [B, S, C] / [B, S, K, C] 20 | """ 21 | device = points.device 22 | B = points.shape[0] 23 | view_shape = list(idx.shape) 24 | view_shape[1:] = [1] * (len(view_shape) - 1) 25 | repeat_shape = list(idx.shape) 26 | repeat_shape[0] = 1 27 | batch_indices = torch.arange(B, dtype=torch.long).to( 28 | device).view(view_shape).repeat(repeat_shape) 29 | new_points = points[batch_indices, idx, :] 30 | return new_points 31 | 32 | 33 | class PCFFunction(torch.autograd.Function): 34 | ''' 35 | Function for the PCF CUDA kernel 36 | ''' 37 | @staticmethod 38 | def forward(ctx, input_feat, neighbor_inds, guidance, weightnet): 39 | # Make sure we are not computing gradient on neighbor_inds 40 | neighbor_inds.requires_grad = False 41 | output = pcf_cuda.pcf_forward( 42 | input_feat, neighbor_inds, guidance, weightnet) 43 | ctx.save_for_backward(input_feat, neighbor_inds, guidance, weightnet) 44 | return output 45 | 46 | @staticmethod 47 | def backward(ctx, grad_output): 48 | grad_input, grad_guidance, grad_weight = pcf_cuda.pcf_backward( 49 | grad_output.contiguous(), *ctx.saved_tensors) 50 | return grad_input, None, grad_guidance, grad_weight 51 | 52 | 53 | class PCF(torch.nn.Module): 54 | ''' 55 | This class uses the CUDA kernel to fuse gather -> matrix multiplication in PCF which improves speed. 56 | Right now, it is numerically correct, but somehow it will mysteriously reduce training accuracy, hence only recommended to use during testing time 57 | ''' 58 | 59 | def __init__(self): 60 | super(PCF, self).__init__() 61 | 62 | @staticmethod 63 | def forward(input_features, neighbor_inds, guidance, weightnet): 64 | return PCFFunction.apply( 65 | input_features, 66 | neighbor_inds, 67 | guidance, 68 | weightnet) 69 | 70 | 71 | class PConvFunction(torch.autograd.Function): 72 | ''' 73 | Function for the PointConv CUDA kernel 74 | ''' 75 | @staticmethod 76 | def forward( 77 | ctx, 78 | input_feat, 79 | neighbor_inds, 80 | weightnet, 81 | additional_features): 82 | # Make sure we are not computing gradient on neighbor_inds 83 | neighbor_inds.requires_grad = False 84 | output = pcf_cuda.pconv_forward( 85 | input_feat, neighbor_inds, weightnet, additional_features) 86 | ctx.save_for_backward( 87 | input_feat, 88 | neighbor_inds, 89 | weightnet, 90 | additional_features) 91 | return output 92 | 93 | @staticmethod 94 | def backward(ctx, grad_output): 95 | grad_input, grad_weight, grad_additional = pcf_cuda.pconv_backward( 96 | grad_output.contiguous(), *ctx.saved_tensors) 97 | return grad_input, None, grad_weight, grad_additional 98 | 99 | 100 | class PConv(torch.nn.Module): 101 | ''' 102 | This class uses the CUDA kernel to fuse gather -> matrix multiplication in PointConv which improves speed. 103 | Right now, it is numerically correct, but somehow it will mysteriously reduce training accuracy, hence only recommended to use during testing time 104 | ''' 105 | 106 | def __init__(self): 107 | super(PConv, self).__init__() 108 | 109 | @staticmethod 110 | def forward(input_features, neighbor_inds, weightnet, additional_features=None): 111 | if additional_features is None: 112 | additional_features = torch.zeros(input_features.shape[0], input_features.shape[1], neighbor_inds.shape[2], 0) 113 | return PConvFunction.apply( 114 | input_features, 115 | neighbor_inds, 116 | weightnet, 117 | additional_features) 118 | 119 | 120 | def VI_coordinate_transform(localized_xyz, gathered_norm, sparse_xyz_norm, K): 121 | """ 122 | Compute the viewpoint-invariance aware relative position encoding in VI_PointConv 123 | From: X. Li et al. Improving the Robustness of Point Convolution on k-Nearest Neighbor Neighborhoods with a Viewpoint-Invariant Coordinate Transform. WACV 2023 124 | Code copyright 2020 Xingyi Li (MIT License) 125 | Input: 126 | dense_xyz: 3D coordinates (note VI only works on 3D) 127 | nei_inds: indices of neighborhood points for each point 128 | dense_xyz_norm: surface normals for each point 129 | sparse_xyz_norm: surface normals for each point in the lower resolution (normally 130 | the same as dense_xyz_norm, except when downsampling) 131 | Return: 132 | VI-transformed point coordinates: a concatenation of rotation+scale invariant dimensions, scale-invariant dimensions and non-invariant dimensions 133 | """ 134 | r_hat = F.normalize(localized_xyz, dim=3) 135 | v_miu = sparse_xyz_norm.unsqueeze( 136 | dim=2) - torch.matmul( 137 | sparse_xyz_norm.unsqueeze( 138 | dim=2), r_hat.permute( 139 | 0, 1, 3, 2)).permute( 140 | 0, 1, 3, 2) * r_hat 141 | v_miu = F.normalize(v_miu, dim=3) 142 | w_miu = torch.cross(r_hat, v_miu, dim=3) 143 | w_miu = F.normalize(w_miu, dim=3) 144 | theta1 = torch.matmul(gathered_norm, sparse_xyz_norm.unsqueeze(dim=3)) 145 | theta2 = torch.matmul(r_hat, sparse_xyz_norm.unsqueeze(dim=3)) 146 | theta3 = torch.sum(r_hat * gathered_norm, dim=3, keepdim=True) 147 | theta4 = torch.matmul(localized_xyz, sparse_xyz_norm.unsqueeze(dim=3)) 148 | theta5 = torch.sum(gathered_norm * r_hat, dim=3, keepdim=True) 149 | theta6 = torch.sum(gathered_norm * v_miu, dim=3, keepdim=True) 150 | theta7 = torch.sum(gathered_norm * w_miu, dim=3, keepdim=True) 151 | theta8 = torch.sum( 152 | localized_xyz * 153 | torch.cross( 154 | gathered_norm, 155 | sparse_xyz_norm.unsqueeze( 156 | dim=2).repeat( 157 | 1, 158 | 1, 159 | K, 160 | 1), 161 | dim=3), 162 | dim=3, 163 | keepdim=True) 164 | theta9 = torch.norm(localized_xyz, dim=3, keepdim=True) 165 | return torch.cat([theta1, 166 | theta2, 167 | theta3, 168 | theta4, 169 | theta5, 170 | theta6, 171 | theta7, 172 | theta8, 173 | theta9, 174 | localized_xyz], 175 | dim=3).contiguous() 176 | 177 | 178 | # We did not like that the pyTorch batch normalization requires C to be the 2nd dimension of the Tensor 179 | # It's hard to deal with it during training time, but we can fuse it during inference time 180 | # This one takes in a 4D tensor of shape BNKC, run a linear layer and a BN layer, and then fuses it during inference time 181 | # Output is BNKC as well 182 | # B is batch size, N is number of points, K is number of neighbors 183 | # one would need to call the fuse function during inference time (see 184 | # utils.replace_bn_layers) 185 | class Linear_BN(torch.nn.Module): 186 | def __init__( 187 | self, 188 | in_dim, 189 | out_dim, 190 | bn_ver='2d', 191 | bn_weight_init=1, 192 | bn_momentum=0.1): 193 | super().__init__() 194 | self.c = torch.nn.Linear(in_dim, out_dim) 195 | self.bn_ver = bn_ver 196 | if bn_ver == '2d': 197 | bn = CpBatchNorm2d(out_dim, momentum=bn_momentum) 198 | else: 199 | bn = torch.nn.BatchNorm1d(out_dim, momentum=bn_momentum) 200 | torch.nn.init.constant_(bn.weight, bn_weight_init) 201 | # torch.nn.init.constant_(bn.bias, 0) 202 | self.bn = bn 203 | 204 | @torch.no_grad() 205 | @torch.jit.ignore() 206 | def fuse(self): 207 | w = self.bn.weight / (self.bn.running_var + self.bn.eps) ** 0.5 208 | w = self.c.weight * w[:, None] 209 | b = self.bn.bias + (self.c.bias - self.bn.running_mean) * self.bn.weight / \ 210 | (self.bn.running_var + self.bn.eps)**0.5 211 | new_layer = torch.nn.Linear(w.size(1), w.size(0)) 212 | new_layer.weight.data.copy_(w) 213 | new_layer.bias.data.copy_(b) 214 | return new_layer 215 | 216 | def forward(self, x): 217 | x = self.c(x) 218 | if self.bn_ver == '2d': 219 | return self.bn(x.permute(0, 3, 2, 1)).permute(0, 3, 2, 1) 220 | else: 221 | return self.bn(x.permute(0, 2, 1)).permute(0, 2, 1) 222 | 223 | 224 | # Linear_BN + Leaky ReLU activation 225 | class UnaryBlock(nn.Module): 226 | def __init__(self, in_dim, out_dim, use_bn, bn_momentum, no_relu=False): 227 | """ 228 | Initialize a standard unary block with its ReLU and BatchNorm. 229 | :param in_dim: dimension input features 230 | :param out_dim: dimension input features 231 | :param use_bn: boolean indicating if we use Batch Norm 232 | :param bn_momentum: Batch norm momentum 233 | """ 234 | 235 | super(UnaryBlock, self).__init__() 236 | self.bn_momentum = bn_momentum 237 | self.use_bn = use_bn 238 | self.no_relu = no_relu 239 | self.in_dim = in_dim 240 | self.out_dim = out_dim 241 | if use_bn: 242 | self.mlp = Linear_BN( 243 | in_dim, 244 | out_dim, 245 | bn_momentum=bn_momentum, 246 | bn_ver='1d') 247 | else: 248 | self.mlp = nn.Linear(in_dim, out_dim) 249 | if not no_relu: 250 | self.leaky_relu = nn.LeakyReLU(0.1) 251 | else: 252 | self.leaky_relu = nn.Identity() 253 | return 254 | 255 | def forward(self, x): 256 | x = self.mlp(x) 257 | if not self.no_relu: 258 | x = self.leaky_relu(x) 259 | return x 260 | 261 | def __repr__(self): 262 | return 'UnaryBlock(in_feat: {:d}, out_feat: {:d}, BN: {:s}, ReLU: {:s})'.format( 263 | self.in_dim, self.out_dim, str(self.use_bn), str(not self.no_relu)) 264 | -------------------------------------------------------------------------------- /util/lr.py: -------------------------------------------------------------------------------- 1 | ## Adapted from StratifiedTransformer 2 | ## https://github.com/dvlab-research/Stratified-Transformer 3 | ## Copyright @ 2022 DV Lab 4 | ## MIT License, https://github.com/dvlab-research/Stratified-Transformer/blob/main/LICENSE.md 5 | 6 | import math 7 | import torch 8 | 9 | from torch.optim.lr_scheduler import LambdaLR, StepLR, OneCycleLR 10 | import torch.optim as optim 11 | from torch.optim.lr_scheduler import _LRScheduler 12 | 13 | 14 | class LambdaStepLR(LambdaLR): 15 | 16 | def __init__(self, optimizer, lr_lambda, last_step=-1): 17 | super(LambdaStepLR, self).__init__(optimizer, lr_lambda, last_step) 18 | 19 | @property 20 | def last_step(self): 21 | """Use last_epoch for the step counter""" 22 | return self.last_epoch 23 | 24 | @last_step.setter 25 | def last_step(self, v): 26 | self.last_epoch = v 27 | 28 | 29 | class PolyLRwithWarmup(LambdaStepLR): 30 | """DeepLab learning rate policy""" 31 | 32 | def __init__(self, optimizer, max_iter, warmup='linear', warmup_iters=1500, warmup_ratio=1e-6, power=1.0, last_step=-1): 33 | 34 | assert warmup == 'linear' 35 | def poly_with_warmup(s): 36 | coeff = (1 - s / (max_iter+1)) ** power 37 | if s <= warmup_iters: 38 | warmup_coeff = 1 - (1 - s / warmup_iters) * (1 - warmup_ratio) 39 | else: 40 | warmup_coeff = 1.0 41 | return coeff * warmup_coeff 42 | 43 | super(PolyLRwithWarmup, self).__init__(optimizer, poly_with_warmup, last_step) 44 | # torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False) 45 | # lr_lambda: A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups. 46 | 47 | 48 | class MultiStepWithWarmup(LambdaStepLR): 49 | def __init__(self, optimizer, milestones, gamma=0.1, warmup='linear', warmup_iters=10, warmup_ratio=1e-6, last_step=-1): 50 | 51 | assert warmup == 'linear' 52 | def multi_step_with_warmup(s): 53 | factor = 1.0 54 | for i in range(len(milestones)): 55 | if s < milestones[i]: 56 | break 57 | factor *= gamma 58 | 59 | if s <= warmup_iters: 60 | warmup_coeff = 1 - (1 - s / warmup_iters) * (1 - warmup_ratio) 61 | else: 62 | warmup_coeff = 1.0 63 | return warmup_coeff * factor 64 | 65 | super(MultiStepWithWarmup, self).__init__(optimizer, multi_step_with_warmup, last_step) 66 | 67 | 68 | class PolyLR(LambdaStepLR): 69 | """DeepLab learning rate policy""" 70 | 71 | def __init__(self, optimizer, max_iter, power=0.9, last_step=-1): 72 | super(PolyLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1))**power, last_step) 73 | # torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False) 74 | # lr_lambda: A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups. 75 | 76 | 77 | class SquaredLR(LambdaStepLR): 78 | """ Used for SGD Lars""" 79 | 80 | def __init__(self, optimizer, max_iter, last_step=-1): 81 | super(SquaredLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1))**2, last_step) 82 | 83 | 84 | class ExpLR(LambdaStepLR): 85 | 86 | def __init__(self, optimizer, step_size, gamma=0.9, last_step=-1): 87 | # (0.9 ** 21.854) = 0.1, (0.95 ** 44.8906) = 0.1 88 | # To get 0.1 every N using gamma 0.9, N * log(0.9)/log(0.1) = 0.04575749 N 89 | # To get 0.1 every N using gamma g, g ** N = 0.1 -> N * log(g) = log(0.1) -> g = np.exp(log(0.1) / N) 90 | super(ExpLR, self).__init__(optimizer, lambda s: gamma**(s / step_size), last_step) 91 | 92 | 93 | def initialize_scheduler(optimizer, config, last_epoch=-1, scheduler_epoch=True, logger=None): 94 | # scheduler_epoch: the step_size are given in epoch num 95 | last_step = -1 if last_epoch < 0 else config.iter_per_epoch_train * (last_epoch + 1) - 1 96 | if scheduler_epoch: 97 | config.step_size = config.iter_per_epoch_train * config.step_size 98 | config.exp_step_size = config.iter_per_epoch_train * config.exp_step_size 99 | 100 | if config.scheduler == 'StepLR': 101 | return StepLR(optimizer, step_size=config.step_size, gamma=config.step_gamma, last_epoch=last_step) 102 | elif config.scheduler == 'PolyLR': 103 | return PolyLR(optimizer, max_iter=config.max_iter, power=config.poly_power, last_step=last_step) 104 | elif config.scheduler == 'PolyLRwithWarmup': 105 | return PolyLRwithWarmup(optimizer, max_iter=config.max_iter, warmup=config.warmup, warmup_iters=config.warmup_iters, warmup_ratio=config.warmup_ratio, power=config.poly_power, last_step=last_step) 106 | elif config.scheduler == 'SquaredLR': 107 | return SquaredLR(optimizer, max_iter=config.max_iter, last_step=last_step) 108 | elif config.scheduler == 'ExpLR': 109 | return ExpLR(optimizer, step_size=config.exp_step_size, gamma=config.exp_gamma, last_step=last_step) 110 | elif config.scheduler == 'OneCycleLR': 111 | return OneCycleLR(optimizer, max_lr=config.oc_max_lr, total_steps=config.max_iter, pct_start=config.oc_pct_start, 112 | anneal_strategy=config.oc_anneal_strategy, div_factor=config.oc_div_factor, 113 | final_div_factor=config.oc_final_div_factor, last_epoch=last_step) 114 | # (optimizer, max_lr, total_steps=None, epochs=None, steps_per_epoch=None, pct_start=0.3, anneal_strategy='cos', cycle_momentum=True, base_momentum=0.85, max_momentum=0.95, div_factor=25.0, final_div_factor=10000.0, last_epoch=-1) 115 | else: 116 | if logger is not None: 117 | logger.info('Scheduler not supported') 118 | else: print('Scheduler not supported') 119 | 120 | 121 | class CosineAnnealingWarmupRestarts(_LRScheduler): 122 | """ 123 | optimizer (Optimizer): Wrapped optimizer. 124 | first_cycle_steps (int): First cycle step size. 125 | cycle_mult(float): Cycle steps magnification. Default: -1. 126 | max_lr(float): First cycle's max learning rate. Default: 0.1. 127 | min_lr(float): Min learning rate. Default: 0.001. 128 | warmup_steps(int): Linear warmup step size. Default: 0. 129 | gamma(float): Decrease rate of max learning rate by cycle. Default: 1. 130 | last_epoch (int): The index of last epoch. Default: -1. 131 | """ 132 | 133 | def __init__(self, 134 | optimizer : torch.optim.Optimizer, 135 | first_cycle_steps : int, 136 | cycle_mult : float = 1., 137 | max_lr : float = 0.1, 138 | min_lr : float = 0.001, 139 | warmup_steps : int = 0, 140 | gamma : float = 1., 141 | last_epoch : int = -1 142 | ): 143 | assert warmup_steps < first_cycle_steps 144 | 145 | self.first_cycle_steps = first_cycle_steps # first cycle step size 146 | self.cycle_mult = cycle_mult # cycle steps magnification 147 | self.base_max_lr = max_lr # first max learning rate 148 | self.max_lr = max_lr # max learning rate in the current cycle 149 | self.min_lr = min_lr # min learning rate 150 | self.warmup_steps = warmup_steps # warmup step size 151 | self.gamma = gamma # decrease rate of max learning rate by cycle 152 | 153 | self.cur_cycle_steps = first_cycle_steps # first cycle step size 154 | self.cycle = 0 # cycle count 155 | self.step_in_cycle = last_epoch # step size of the current cycle 156 | 157 | super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch) 158 | 159 | # set learning rate min_lr 160 | self.init_lr() 161 | 162 | def init_lr(self): 163 | self.base_lrs = [] 164 | for param_group in self.optimizer.param_groups: 165 | param_group['lr'] = self.min_lr 166 | self.base_lrs.append(self.min_lr) 167 | 168 | def get_lr(self): 169 | if self.step_in_cycle == -1: 170 | return self.base_lrs 171 | elif self.step_in_cycle < self.warmup_steps: 172 | return [(self.max_lr - base_lr)*self.step_in_cycle / self.warmup_steps + base_lr for base_lr in self.base_lrs] 173 | else: 174 | return [base_lr + (self.max_lr - base_lr) \ 175 | * (1 + math.cos(math.pi * (self.step_in_cycle-self.warmup_steps) \ 176 | / (self.cur_cycle_steps - self.warmup_steps))) / 2 177 | for base_lr in self.base_lrs] 178 | 179 | def get_last_lr(self): 180 | return self.get_lr() 181 | 182 | def step(self, epoch=None): 183 | if epoch is None: 184 | epoch = self.last_epoch + 1 185 | self.step_in_cycle = self.step_in_cycle + 1 186 | if self.step_in_cycle >= self.cur_cycle_steps: 187 | self.cycle += 1 188 | self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps 189 | self.cur_cycle_steps = int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult) + self.warmup_steps 190 | else: 191 | if epoch >= self.first_cycle_steps: 192 | if self.cycle_mult == 1.: 193 | self.step_in_cycle = epoch % self.first_cycle_steps 194 | self.cycle = epoch // self.first_cycle_steps 195 | else: 196 | n = int(math.log((epoch / self.first_cycle_steps * (self.cycle_mult - 1) + 1), self.cycle_mult)) 197 | self.cycle = n 198 | self.step_in_cycle = epoch - int(self.first_cycle_steps * (self.cycle_mult ** n - 1) / (self.cycle_mult - 1)) 199 | self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (n) 200 | else: 201 | self.cur_cycle_steps = self.first_cycle_steps 202 | self.step_in_cycle = epoch 203 | 204 | self.max_lr = self.base_max_lr * (self.gamma**self.cycle) 205 | self.last_epoch = math.floor(epoch) 206 | for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): 207 | param_group['lr'] = lr 208 | 209 | 210 | 211 | if __name__ == '__main__': 212 | import torchvision.models as models 213 | model = models.vgg16() 214 | optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001) 215 | optimizer.param_groups[0]['initial_lr'] = 0.2 / 25.0 216 | optimizer.param_groups[0]['max_lr'] = 0.2 217 | optimizer.param_groups[0]['min_lr'] = 0.2 / 10000.0 218 | optimizer.param_groups[0]['max_momentum'] = 0.95 219 | optimizer.param_groups[0]['base_momentum'] = 0.85 220 | last_step = 2 221 | max_iter = 100 222 | # scheduler = PolyLR(optimizer, max_iter=max_iter, power=0.9, last_step=last_step) 223 | scheduler = OneCycleLR(optimizer, max_lr=0.2, total_steps=max_iter, pct_start=0.1, anneal_strategy='cos', div_factor=25.0, 224 | final_div_factor=10000.0, last_epoch=last_step) 225 | lr_list = [] 226 | for epoch in range(max(last_step + 1, 0), min(max_iter, 100)): 227 | lrs = ', '.join(['{:.5e}'.format(x) for x in scheduler.get_last_lr()]) 228 | print('epoch {} lrs {}'.format(epoch, lrs)) 229 | lr_list.append(scheduler.get_last_lr()[0]) 230 | scheduler.step() 231 | 232 | import numpy as np 233 | import matplotlib.pyplot as plt 234 | x = np.arange(max(last_step + 1, 0), min(max_iter, 100), 1) 235 | plt.title("function") 236 | plt.plot(x, lr_list) 237 | plt.show() 238 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | #include "neighbors.h" 6 | #include 7 | 8 | 9 | void brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose) 10 | { 11 | 12 | // Initialize variables 13 | // ****************** 14 | 15 | // square radius 16 | float r2 = radius * radius; 17 | 18 | // indices 19 | int i0 = 0; 20 | 21 | // Counting vector 22 | int max_count = 0; 23 | vector> tmp(queries.size()); 24 | 25 | // Search neigbors indices 26 | // *********************** 27 | 28 | for (auto& p0 : queries) 29 | { 30 | int i = 0; 31 | for (auto& p : supports) 32 | { 33 | if ((p0 - p).sq_norm() < r2) 34 | { 35 | tmp[i0].push_back(i); 36 | if (tmp[i0].size() > max_count) 37 | max_count = tmp[i0].size(); 38 | } 39 | i++; 40 | } 41 | i0++; 42 | } 43 | 44 | // Reserve the memory 45 | neighbors_indices.resize(queries.size() * max_count); 46 | i0 = 0; 47 | for (auto& inds : tmp) 48 | { 49 | for (int j = 0; j < max_count; j++) 50 | { 51 | if (j < inds.size()) 52 | neighbors_indices[i0 * max_count + j] = inds[j]; 53 | else 54 | neighbors_indices[i0 * max_count + j] = -1; 55 | } 56 | i0++; 57 | } 58 | 59 | return; 60 | } 61 | 62 | void ordered_neighbors(vector& queries, 63 | vector& supports, 64 | vector& neighbors_indices, 65 | float radius) 66 | { 67 | 68 | // Initialize variables 69 | // ****************** 70 | 71 | // square radius 72 | float r2 = radius * radius; 73 | 74 | // indices 75 | int i0 = 0; 76 | 77 | // Counting vector 78 | int max_count = 0; 79 | float d2; 80 | vector> tmp(queries.size()); 81 | vector> dists(queries.size()); 82 | 83 | // Search neigbors indices 84 | // *********************** 85 | 86 | for (auto& p0 : queries) 87 | { 88 | int i = 0; 89 | for (auto& p : supports) 90 | { 91 | d2 = (p0 - p).sq_norm(); 92 | if (d2 < r2) 93 | { 94 | // Find order of the new point 95 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 96 | int index = std::distance(dists[i0].begin(), it); 97 | 98 | // Insert element 99 | dists[i0].insert(it, d2); 100 | tmp[i0].insert(tmp[i0].begin() + index, i); 101 | 102 | // Update max count 103 | if (tmp[i0].size() > max_count) 104 | max_count = tmp[i0].size(); 105 | } 106 | i++; 107 | } 108 | i0++; 109 | } 110 | 111 | // Reserve the memory 112 | neighbors_indices.resize(queries.size() * max_count); 113 | i0 = 0; 114 | for (auto& inds : tmp) 115 | { 116 | for (int j = 0; j < max_count; j++) 117 | { 118 | if (j < inds.size()) 119 | neighbors_indices[i0 * max_count + j] = inds[j]; 120 | else 121 | neighbors_indices[i0 * max_count + j] = -1; 122 | } 123 | i0++; 124 | } 125 | 126 | return; 127 | } 128 | 129 | void batch_ordered_neighbors(vector& queries, 130 | vector& supports, 131 | vector& q_batches, 132 | vector& s_batches, 133 | vector& neighbors_indices, 134 | float radius) 135 | { 136 | 137 | // Initialize variables 138 | // ****************** 139 | 140 | // square radius 141 | float r2 = radius * radius; 142 | 143 | // indices 144 | int i0 = 0; 145 | 146 | // Counting vector 147 | int max_count = 0; 148 | float d2; 149 | vector> tmp(queries.size()); 150 | vector> dists(queries.size()); 151 | 152 | // batch index 153 | int b = 0; 154 | int sum_qb = 0; 155 | int sum_sb = 0; 156 | 157 | 158 | // Search neigbors indices 159 | // *********************** 160 | 161 | for (auto& p0 : queries) 162 | { 163 | // Check if we changed batch 164 | if (i0 == sum_qb + q_batches[b]) 165 | { 166 | sum_qb += q_batches[b]; 167 | sum_sb += s_batches[b]; 168 | b++; 169 | } 170 | 171 | // Loop only over the supports of current batch 172 | vector::iterator p_it; 173 | int i = 0; 174 | for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ ) 175 | { 176 | d2 = (p0 - *p_it).sq_norm(); 177 | if (d2 < r2) 178 | { 179 | // Find order of the new point 180 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 181 | int index = std::distance(dists[i0].begin(), it); 182 | 183 | // Insert element 184 | dists[i0].insert(it, d2); 185 | tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i); 186 | 187 | // Update max count 188 | if (tmp[i0].size() > max_count) 189 | max_count = tmp[i0].size(); 190 | } 191 | i++; 192 | } 193 | i0++; 194 | } 195 | 196 | // Reserve the memory 197 | neighbors_indices.resize(queries.size() * max_count); 198 | i0 = 0; 199 | for (auto& inds : tmp) 200 | { 201 | for (int j = 0; j < max_count; j++) 202 | { 203 | if (j < inds.size()) 204 | neighbors_indices[i0 * max_count + j] = inds[j]; 205 | else 206 | neighbors_indices[i0 * max_count + j] = supports.size(); 207 | } 208 | i0++; 209 | } 210 | 211 | return; 212 | } 213 | 214 | 215 | void batch_nanoflann_neighbors(vector& queries, 216 | vector& supports, 217 | vector& q_batches, 218 | vector& s_batches, 219 | vector& neighbors_indices, 220 | float radius) 221 | { 222 | 223 | // Initialize variables 224 | // ****************** 225 | 226 | // indices 227 | int i0 = 0; 228 | 229 | // Square radius 230 | float r2 = radius * radius; 231 | 232 | // Counting vector 233 | int max_count = 0; 234 | float d2; 235 | vector>> all_inds_dists(queries.size()); 236 | 237 | // batch index 238 | int b = 0; 239 | int sum_qb = 0; 240 | int sum_sb = 0; 241 | 242 | // Nanoflann related variables 243 | // *************************** 244 | 245 | // CLoud variable 246 | PointCloud current_cloud; 247 | 248 | // Tree parameters 249 | nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); 250 | 251 | // KDTree type definition 252 | typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , 253 | PointCloud, 254 | 3 > my_kd_tree_t; 255 | 256 | // Pointer to trees 257 | my_kd_tree_t* index; 258 | 259 | // Build KDTree for the first batch element 260 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 261 | index = new my_kd_tree_t(3, current_cloud, tree_params); 262 | index->buildIndex(); 263 | 264 | 265 | // Search neigbors indices 266 | // *********************** 267 | 268 | // Search params 269 | nanoflann::SearchParameters search_params; 270 | search_params.sorted = true; 271 | 272 | for (auto& p0 : queries) 273 | { 274 | 275 | // Check if we changed batch 276 | if (i0 == sum_qb + q_batches[b]) 277 | { 278 | sum_qb += q_batches[b]; 279 | sum_sb += s_batches[b]; 280 | b++; 281 | 282 | // Change the points 283 | current_cloud.pts.clear(); 284 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 285 | 286 | // Build KDTree of the current element of the batch 287 | delete index; 288 | index = new my_kd_tree_t(3, current_cloud, tree_params); 289 | index->buildIndex(); 290 | } 291 | 292 | // Initial guess of neighbors size 293 | all_inds_dists[i0].reserve(max_count); 294 | 295 | // Find neighbors 296 | float query_pt[3] = { p0.x, p0.y, p0.z}; 297 | 298 | size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params); 299 | 300 | // Update max count 301 | if (nMatches > max_count) 302 | max_count = nMatches; 303 | 304 | // Increment query idx 305 | i0++; 306 | } 307 | 308 | // Reserve the memory 309 | neighbors_indices.resize(queries.size() * max_count); 310 | i0 = 0; 311 | sum_sb = 0; 312 | sum_qb = 0; 313 | b = 0; 314 | for (auto& inds_dists : all_inds_dists) 315 | { 316 | // Check if we changed batch 317 | if (i0 == sum_qb + q_batches[b]) 318 | { 319 | sum_qb += q_batches[b]; 320 | sum_sb += s_batches[b]; 321 | b++; 322 | } 323 | 324 | for (int j = 0; j < max_count; j++) 325 | { 326 | if (j < inds_dists.size()) 327 | neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb; 328 | else 329 | neighbors_indices[i0 * max_count + j] = supports.size(); 330 | } 331 | i0++; 332 | } 333 | 334 | delete index; 335 | 336 | return; 337 | } 338 | 339 | void batch_nanoflann_knn(vector& queries, 340 | vector& supports, 341 | vector& q_batches, 342 | vector& s_batches, 343 | vector& neighbors_indices, 344 | const size_t K) 345 | { 346 | 347 | // Initialize variables 348 | // ****************** 349 | 350 | // indices 351 | int i0 = 0; 352 | /*# vector> all_inds_dists(queries.size()); */ 353 | 354 | // batch index 355 | int b = 0; 356 | int sum_qb = 0; 357 | int sum_sb = 0; 358 | 359 | // Nanoflann related variables 360 | // *************************** 361 | 362 | // CLoud variable 363 | PointCloud current_cloud; 364 | 365 | // Tree parameters 366 | nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); 367 | 368 | // KDTree type definition 369 | typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , 370 | PointCloud, 371 | 3 > my_kd_tree_t; 372 | 373 | // Pointer to trees 374 | my_kd_tree_t* index; 375 | 376 | // Build KDTree for the first batch element 377 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 378 | index = new my_kd_tree_t(3, current_cloud, tree_params); 379 | index->buildIndex(); 380 | 381 | 382 | // Search neigbors indices 383 | // *********************** 384 | 385 | // Search params 386 | nanoflann::SearchParameters search_params; 387 | search_params.sorted = true; 388 | neighbors_indices.resize(queries.size() * K); 389 | 390 | for (auto& p0 : queries) 391 | { 392 | 393 | // Check if we changed batch 394 | if (i0 == sum_qb + q_batches[b]) 395 | { 396 | sum_qb += q_batches[b]; 397 | sum_sb += s_batches[b]; 398 | b++; 399 | 400 | // Change the points 401 | current_cloud.pts.clear(); 402 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 403 | 404 | // Build KDTree of the current element of the batch 405 | delete index; 406 | index = new my_kd_tree_t(3, current_cloud, tree_params); 407 | index->buildIndex(); 408 | } 409 | // all_inds_dists[i0].reserve(K); 410 | 411 | // Find neighbors 412 | float query_pt[3] = { p0.x, p0.y, p0.z}; 413 | vector out_dist_sqr; 414 | out_dist_sqr.reserve(K); 415 | index->knnSearch(query_pt, K, &neighbors_indices[i0*K], &out_dist_sqr[0]); 416 | 417 | 418 | // Increment query idx 419 | i0++; 420 | } 421 | // Reserve the memory 422 | i0 = 0; 423 | sum_sb = 0; 424 | sum_qb = 0; 425 | b = 0; 426 | for (i0 = 0; i0 0) 436 | for (int j = 0; j < K; j++) 437 | neighbors_indices[i0 * K + j] += sum_sb; 438 | } 439 | 440 | delete index; 441 | 442 | return; 443 | } 444 | -------------------------------------------------------------------------------- /transforms.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | import random 7 | 8 | import numpy as np 9 | import scipy 10 | import scipy.ndimage 11 | import scipy.interpolate 12 | from numpy import cross 13 | from scipy.linalg import expm, norm 14 | import torch 15 | 16 | 17 | class Compose(object): 18 | """ 19 | Composes several transforms together. 20 | Parameters: 21 | transforms: The transforms that will be combined together 22 | Call: 23 | *args: Usually coords, features, label, norms, the coordinates/features/labels/normals of the points, respectively 24 | Return: 25 | args: Usually coords, features, label, norms, the coordinates/features/labels/normals of the points, respectively 26 | """ 27 | 28 | def __init__(self, transforms): 29 | self.transforms = transforms 30 | 31 | def __call__(self, *args): 32 | for t in self.transforms: 33 | args = t(*args) 34 | return args 35 | 36 | 37 | class RandomDropColor(object): 38 | """ 39 | Random drop color augmentation. 40 | Parameters: 41 | p: probability the augmentation is applied. Default to be 0.8 (apply the augmentation 80% of the time) 42 | color_augment: the amount of color drop, default to 0.0 (completely remove color) 43 | Call: 44 | coords: input coordinates for each point 45 | color: input color for each point 46 | label: input label for each point 47 | norms: input normals for each point 48 | Return: 49 | coords: input coordinates for each point 50 | color: color after random dropping 51 | label: input label for each point 52 | norms: input normals for each point 53 | for this augmentation, coords, label and norms will not be changed 54 | """ 55 | def __init__(self, p=0.8, color_augment=0.0): 56 | self.p = p 57 | self.color_augment = color_augment 58 | 59 | def __call__(self, coords, color, labels, norms): 60 | t = torch.rand(1).numpy()[0] 61 | # print(t) 62 | if color is not None and t > self.p: 63 | color *= self.color_augment 64 | return coords, color, labels, norms 65 | 66 | def __repr__(self): 67 | return 'RandomDropColor(color_augment: {}, p: {})'.format(self.color_augment, self.p) 68 | 69 | 70 | class RandomDropout(object): 71 | """ 72 | Random dropout points augmentation. This will randomly drop some points from the point cloud. 73 | Parameters: 74 | dropout_ratio: probability a point is dropped, default to 0.2 (20% chance) 75 | dropout_application_ratio: probability the dropout is applied to the point cloud, default to 0.5 (50% chance) 76 | Call: 77 | coords: input coordinates for each point 78 | feats: input features for each point 79 | label: input label for each point 80 | norms: input normals for each point 81 | Return: 82 | coords: coordinates after dropout 83 | feats: features after dropout 84 | label: labels after dropout 85 | norms: normals after dropout 86 | """ 87 | def __init__(self, dropout_ratio=0.2, dropout_application_ratio=0.5): 88 | self.dropout_ratio = dropout_ratio 89 | self.dropout_application_ratio = dropout_application_ratio 90 | 91 | def __call__(self, coords, feats, labels, norms): 92 | if random.random() < self.dropout_application_ratio: 93 | N = len(coords) 94 | inds = np.random.choice(N, int(N * (1 - self.dropout_ratio)), replace=False) 95 | return coords[inds], feats[inds], labels[inds], norms[inds] 96 | return coords, feats, labels, norms 97 | 98 | 99 | class RandomHorizontalFlip(object): 100 | """ 101 | Random horizontal flip augmentation. This will flip the object in the xy plane. 102 | Parameters: 103 | apply_likelihood: determines how likely the transformation will be applied, default 0.95 (95% chance) 104 | axis_flip_likelihood: determines how likely each axis will be flipped if the point cloud is to be flipped, default 0.5 (50% chance) 105 | upright_axis: determines which axis is the z dimension, usually, 'z' for using the regular 'z' (axis index = 2 assuming xyz indexing) 106 | Call: 107 | coords: input coordinates for each point 108 | feats: input features for each point 109 | label: input label for each point 110 | norms: input normals for each point 111 | Return: 112 | coords: coordinates after flipping 113 | feats: input features 114 | label: input labels 115 | norms: normals after flipping 116 | for this augmentation, feats and label will not be changed 117 | """ 118 | def __init__(self, upright_axis, apply_likelihood=0.95, axis_flip_likelihood=0.5): 119 | """ 120 | upright_axis: axis index among x,y,z, i.e. 2 for z 121 | """ 122 | self.D = 3 123 | self.apply_likelihood = apply_likelihood 124 | self.axis_flip_likelihood = axis_flip_likelihood 125 | self.upright_axis = {'x': 0, 'y': 1, 'z': 2}[upright_axis.lower()] 126 | # Use the rest of axes for flipping. 127 | self.horz_axes = set(range(self.D)) - set([self.upright_axis]) 128 | 129 | def __call__(self, coords, feats, labels, norms): 130 | if random.random() < self.apply_likelihood: 131 | for curr_ax in self.horz_axes: 132 | if random.random() < self.axis_flip_likelihood: 133 | coords[:, curr_ax] = - coords[:, curr_ax] 134 | norms[:, curr_ax] = - norms[:, curr_ax] 135 | return coords, feats, labels, norms 136 | 137 | 138 | class ChromaticTranslation(object): 139 | """ 140 | Add random color to the point cloud, input must be an array in [0,255] or a PIL image 141 | By default, the first 3 dimensions of the features will be assumed as color 142 | Parameters: 143 | apply_likelihood: determines how likely the transformation will be applied, default 0.95 (95% chance) 144 | trans_range_ratio: ratio of translation i.e. 255 * 2 * ratio * rand(-0.5, 0.5), default 0.1 145 | Call: 146 | coords: input coordinates for each point 147 | feats: input features for each point 148 | label: input label for each point 149 | norms: input normals for each point 150 | Return: 151 | coords: input coordinates for each point 152 | feats: features after application of the transformation 153 | label: input label for each point 154 | norms: input normals for each point 155 | for this augmentation, coords, label and norms will not be changed 156 | """ 157 | 158 | def __init__(self, trans_range_ratio=1e-1, apply_likelihood=0.95): 159 | self.apply_likelihood = apply_likelihood 160 | self.trans_range_ratio = trans_range_ratio 161 | 162 | def __call__(self, coords, feats, labels, norms): 163 | if torch.rand(1).numpy()[0] < self.apply_likelihood: 164 | tr = (torch.rand(1, 3).numpy() - 0.5) * 255 * 2 * self.trans_range_ratio 165 | feats[:, :3] = np.clip(tr + feats[:, :3], 0, 255) 166 | return coords, feats, labels, norms 167 | 168 | 169 | class ChromaticAutoContrast(object): 170 | """ 171 | Blend features with another version of the features by changing the color contrast, input must be an array in [0,255] or a PIL image 172 | By default, the first 3 dimensions of the features will be assumed as color 173 | Parameters: 174 | randomize_blend_factor: use random blending factors for each point cloud or not (default True) 175 | blend_factor: blending ratio between the original color (1-blend_factor) and the color from the new contrast (blend_factor) 176 | Call: 177 | coords: input coordinates for each point 178 | feats: input features for each point 179 | label: input label for each point 180 | norms: input normals for each point 181 | Return: 182 | coords: input coordinates for each point 183 | feats: features after application of the transformation 184 | label: input label for each point 185 | norms: input normals for each point 186 | for this augmentation, coords, label and norms will not be changed 187 | """ 188 | def __init__(self, randomize_blend_factor=True, blend_factor=0.5): 189 | self.randomize_blend_factor = randomize_blend_factor 190 | self.blend_factor = blend_factor 191 | 192 | def __call__(self, coords, feats, labels, norms): 193 | if torch.rand(1).numpy()[0] < 0.2: 194 | # mean = np.mean(feats, 0, keepdims=True) 195 | # std = np.std(feats, 0, keepdims=True) 196 | # lo = mean - std 197 | # hi = mean + std 198 | lo = np.min(feats[:, :3], 0, keepdims=True) 199 | hi = np.max(feats[:, :3], 0, keepdims=True) 200 | 201 | scale = 255 / (hi - lo) 202 | 203 | contrast_feats = (feats[:, :3] - lo) * scale 204 | 205 | blend_factor = torch.rand(1).numpy()[0] if self.randomize_blend_factor else self.blend_factor 206 | feats[:, :3] = (1 - blend_factor) * feats[:, :3] + blend_factor * contrast_feats 207 | return coords, feats, labels, norms 208 | 209 | 210 | class ChromaticJitter(object): 211 | """ 212 | Jitter the color of points (add random noise on the color) 213 | By default, the first 3 dimensions of the features will be assumed as color 214 | Parameters: 215 | std: standard deviation of the color jitter assuming color is distributed as N(0,1) (will be multiplied by 255 for color ranging [0,255]) 216 | apply_likelihood: determines how likely the transformation will be applied, default 0.95 (95% chance) 217 | Call: 218 | coords: input coordinates for each point 219 | feats: input features for each point 220 | label: input label for each point 221 | norms: input normals for each point 222 | Return: 223 | coords: input coordinates for each point 224 | feats: features after application of the transformation 225 | label: input label for each point 226 | norms: input normals for each point 227 | for this augmentation, coords, label and norms will not be changed 228 | """ 229 | def __init__(self, std=0.01, apply_likelihood=0.95): 230 | self.apply_likelihood = apply_likelihood 231 | self.std = std 232 | 233 | def __call__(self, coords, feats, labels, norms): 234 | if torch.rand(1).numpy()[0] < self.apply_likelihood: 235 | # noise = np.random.randn(feats.shape[0], 3) 236 | noise = torch.randn(feats.shape[0], 3).numpy() 237 | noise *= self.std * 255 238 | feats[:, :3] = np.clip(noise + feats[:, :3], 0, 255) 239 | return coords, feats, labels, norms 240 | 241 | 242 | def elastic_distortion(pointcloud, granularity, magnitude): 243 | """ 244 | Apply elastic distortion on sparse coordinate space. 245 | Call: 246 | pointcloud: numpy array of (number of points, at least 3 spatial dims) 247 | granularity: size of the noise grid (in same scale[m/cm] as the voxel grid) 248 | magnitude: noise multiplier 249 | Return: 250 | pointcloud: point cloud after elastic distortions 251 | """ 252 | blurx = np.ones((3, 1, 1, 1)).astype('float32') / 3 253 | blury = np.ones((1, 3, 1, 1)).astype('float32') / 3 254 | blurz = np.ones((1, 1, 3, 1)).astype('float32') / 3 255 | coords = pointcloud[:, :3] 256 | coords_min = coords.min(0) 257 | 258 | # Create Gaussian noise tensor of the size given by granularity. 259 | noise_dim = ((coords - coords_min).max(0) // granularity).astype(int) + 3 260 | noise = np.random.randn(*noise_dim, 3).astype(np.float32) 261 | 262 | # Smoothing. 263 | for _ in range(2): 264 | noise = scipy.ndimage.filters.convolve(noise, blurx, mode='constant', cval=0) 265 | noise = scipy.ndimage.filters.convolve(noise, blury, mode='constant', cval=0) 266 | noise = scipy.ndimage.filters.convolve(noise, blurz, mode='constant', cval=0) 267 | 268 | # Trilinear interpolate noise filters for each spatial dimensions. 269 | ax = [ 270 | np.linspace(d_min, d_max, d) 271 | for d_min, d_max, d in zip(coords_min - granularity, coords_min + granularity * 272 | (noise_dim - 2), noise_dim) 273 | ] 274 | interp = scipy.interpolate.RegularGridInterpolator(ax, noise, bounds_error=0, fill_value=0) 275 | pointcloud[:, :3] = coords + interp(coords) * magnitude 276 | return pointcloud 277 | 278 | 279 | # Rotation matrix along axis with angle theta 280 | def M(axis, theta): 281 | return expm(cross(np.eye(3), axis / norm(axis) * theta)) 282 | 283 | LOCFEAT_IDX = 2 284 | 285 | 286 | def get_transformation_matrix(rotation_augmentation_bound, scale_augmentation_bound, rotation_angle=None): 287 | """ 288 | Obtain a random transformation matrix. 289 | Call: 290 | rotation_augmentation_bound: maximal degrees of rotation 291 | scale_augmentation_bound: maximal scale 292 | Return: 293 | scale_matrix: scale matrix 294 | rotation_matrix: rotation matrix 295 | """ 296 | scale_matrix = np.eye(4) 297 | rotation_matrix = np.eye(4) 298 | 299 | # Random rotation 300 | rot_mat = np.eye(3) 301 | 302 | rot_mats = [] 303 | for axis_ind, rot_bound in enumerate(rotation_augmentation_bound): 304 | theta = 0 305 | axis = np.zeros(3) 306 | axis[axis_ind] = 1 307 | if rot_bound is not None: 308 | theta = np.random.uniform(*rot_bound) 309 | rot_mats.append(M(axis, theta)) 310 | 311 | # Use random order 312 | np.random.shuffle(rot_mats) 313 | rot_mat = rot_mats[0] @ rot_mats[1] @ rot_mats[2] 314 | 315 | if rotation_angle is not None: 316 | axis = np.zeros(3) 317 | axis[LOCFEAT_IDX] = 1 318 | rot_mat = M(axis, rotation_angle) 319 | 320 | rotation_matrix[:3, :3] = rot_mat 321 | 322 | # Scale 323 | scale = np.random.uniform(scale_augmentation_bound) 324 | np.fill_diagonal(scale_matrix[:3, :3], scale) 325 | return scale_matrix, rotation_matrix 326 | -------------------------------------------------------------------------------- /test_ScanNet_voting.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | 7 | # Testing code with voting, can be used to benchmark on the testing set 8 | # Usage: python test_ScanNet_voting.py --config config_file --pretrain_path model --split split 9 | # where config_file is the training config file (e.g. configFLPCF_10cm.yaml, model is the .pth file 10 | # stored from training, and split is the split you want to test on (e.g. 'validation')) 11 | 12 | import os 13 | import time 14 | import argparse 15 | import datetime 16 | import yaml 17 | import numpy as np 18 | from easydict import EasyDict as edict 19 | 20 | import torch 21 | import torch.nn.functional as F 22 | import torch.utils.data 23 | 24 | import open3d as o3d 25 | 26 | from util.common_util import AverageMeter, intersectionAndUnion, replace_batchnorm, to_device, init_seeds 27 | from util.logger import get_logger 28 | 29 | from model_architecture import PointConvFormer_Segmentation as Main_Model 30 | from train_ScanNet_DDP_WarmUP import get_default_configs 31 | from datasetCommon import collect_fn 32 | from scannet_data_loader_color_DDP import ScanNetDataset 33 | 34 | 35 | def collect_fn_test(data_list): 36 | ''' 37 | Crop the points from a data list into several different ones if there are too many points. 38 | Besides the normal collect function, additionally return idx_data for the different point crops. 39 | (check the 'multiple' mode of the voxelization) 40 | ''' 41 | all_features = [] 42 | all_pointclouds = [] 43 | all_edges_self = [] 44 | all_edges_forward = [] 45 | all_edges_propagate = [] 46 | all_target = [] 47 | all_norms = [] 48 | idx_data = [] 49 | for data in data_list: 50 | the_start = 0 51 | while the_start < len(data): 52 | count = 0 53 | the_end = len(data) 54 | crop_idx = np.zeros(0) 55 | for i, crop in enumerate(data[the_start:]): 56 | count += len(crop['crop_idx']) 57 | if count > args.MAX_POINTS_NUM: 58 | the_end = the_start+i 59 | break 60 | crop_idx = np.concatenate((crop_idx, crop['crop_idx'])) 61 | features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms = collect_fn(data[the_start:the_end]) 62 | all_features.append(features) 63 | all_pointclouds.append(pointclouds) 64 | all_edges_self.append(edges_self) 65 | all_edges_forward.append(edges_forward) 66 | all_edges_propagate.append(edges_propagate) 67 | all_target.append(target) 68 | all_norms.append(norms) 69 | idx_data.append(crop_idx) 70 | the_start = the_end 71 | 72 | return all_features, all_pointclouds, all_edges_self, all_edges_forward, \ 73 | all_edges_propagate, all_target, all_norms, idx_data 74 | 75 | 76 | def get_parser(): 77 | ''' 78 | Get the arguments of the call. 79 | ''' 80 | parser = argparse.ArgumentParser( 81 | description='PointConvFormer for Semantic Segmentation') 82 | parser.add_argument( 83 | '--config', 84 | default='./configFLPCF_10cm.yaml', 85 | type=str, 86 | help='config file') 87 | parser.add_argument( 88 | '--pretrain_path', 89 | default='./', 90 | type=str, 91 | help='the path of pretrain weights') 92 | parser.add_argument( 93 | '--vote_num', 94 | default=1, 95 | type=int, 96 | help='number of vote') 97 | parser.add_argument( 98 | '--split', 99 | default='validation', 100 | type=str, 101 | help='the path of pretrain weights') 102 | parser.add_argument( 103 | '--init_deg', 104 | default=0., 105 | type=float, 106 | help='inital degree [0 ~ 1]') 107 | 108 | args = parser.parse_args() 109 | assert args.config is not None 110 | cfg = edict(yaml.safe_load(open(args.config, 'r'))) 111 | cfg = get_default_configs(cfg) 112 | cfg.pretrain_path = args.pretrain_path 113 | cfg.config = args.config 114 | cfg.vote_num = args.vote_num 115 | cfg.split = args.split 116 | cfg.init_deg = args.init_deg 117 | return cfg 118 | 119 | 120 | def main_vote(): 121 | global args, logger 122 | args = get_parser() 123 | 124 | file_dir = os.path.join(args.eval_path, 'TIME_%s_SemSeg-' % 125 | (args.model_name) + 126 | str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) 127 | args.file_dir = file_dir 128 | ply_dir = os.path.join(args.file_dir, 'ply') 129 | txt_dir = os.path.join(args.file_dir, 'txt') 130 | prob_dir = os.path.join(args.file_dir, 'prob') 131 | 132 | if not os.path.exists(args.eval_path): 133 | os.makedirs(args.eval_path) 134 | 135 | if not os.path.exists(file_dir): 136 | os.makedirs(file_dir) 137 | 138 | if not os.path.exists(ply_dir): 139 | os.makedirs(ply_dir) 140 | 141 | if not os.path.exists(txt_dir): 142 | os.makedirs(txt_dir) 143 | 144 | if not os.path.exists(prob_dir): 145 | os.makedirs(prob_dir) 146 | 147 | logger = get_logger(file_dir) 148 | logger.info(args) 149 | assert args.num_classes > 1 150 | logger.info("=> creating model ...") 151 | logger.info("Classes: {}".format(args.classes)) 152 | 153 | # get model 154 | model = Main_Model(args) 155 | 156 | logger.info(model) 157 | 158 | init_seeds(args.manual_seed) 159 | 160 | if os.path.isfile(args.pretrain_path): 161 | logger.info("=> loading checkpoint '{}'".format(args.pretrain_path)) 162 | checkpoint = torch.load(args.pretrain_path) 163 | state_dict = checkpoint['state_dict'] 164 | model.load_state_dict(state_dict, strict=True) 165 | logger.info( 166 | "=> loaded checkpoint '{}' (epoch {})".format( 167 | args.pretrain_path, 168 | checkpoint['epoch'])) 169 | args.epoch = checkpoint['epoch'] 170 | else: 171 | raise RuntimeError( 172 | "=> no checkpoint found at '{}'".format( 173 | args.pretrain_path)) 174 | 175 | logger.info( 176 | '>>>>>>>>>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>>>>>>>>>>>') 177 | intersection_meter = AverageMeter() 178 | union_meter = AverageMeter() 179 | target_meter = AverageMeter() 180 | 181 | colormap = np.array(create_color_palette40()) / 255.0 182 | mapper = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 183 | 12, 14, 16, 24, 28, 33, 34, 36, 39]) 184 | 185 | model.eval() 186 | replace_batchnorm(model) 187 | model = model.to('cuda') 188 | 189 | num_vote = args.vote_num 190 | rotate_deg_list = np.arange(num_vote) / num_vote + args.init_deg 191 | logger.info("rotate_degs: {}".format(rotate_deg_list)) 192 | 193 | torch.cuda.empty_cache() 194 | 195 | output_dict = {} 196 | 197 | args.BATCH_SIZE = 1 198 | 199 | pcd = o3d.geometry.PointCloud() 200 | time_list = [] 201 | 202 | for vote_id in range(num_vote): 203 | 204 | print( 205 | "vote_id : ", 206 | vote_id, 207 | "rotate_deg : ", 208 | rotate_deg_list[vote_id] * 209 | 360) 210 | rotate_deg = rotate_deg_list[vote_id] 211 | # Rotate differently for each vote 212 | val_dataset = ScanNetDataset(args, dataset=args.split, rotate_deg=rotate_deg, voxelize_mode='multiple') 213 | dataset_loader = torch.utils.data.DataLoader(val_dataset, 214 | batch_size=args.BATCH_SIZE, 215 | collate_fn=collect_fn_test, 216 | num_workers=args.NUM_WORKERS, 217 | sampler=torch.utils.data.SequentialSampler(val_dataset), 218 | pin_memory=True) 219 | 220 | rotate_deg = int(rotate_deg_list[vote_id] * 360) 221 | cur_prob_dir = os.path.join(prob_dir, str(rotate_deg)) 222 | if not os.path.exists(cur_prob_dir): 223 | os.makedirs(cur_prob_dir) 224 | 225 | for idx, all_data in enumerate(dataset_loader): 226 | scene_name = val_dataset.get_scene_name(idx) 227 | features, pointclouds, edges_self, edges_forward, edges_propagate, label, norms, idx_data = all_data 228 | features, pointclouds, edges_self, edges_forward, edges_propagate, label, norms = \ 229 | to_device(features), to_device(pointclouds), \ 230 | to_device(edges_self), to_device(edges_forward), \ 231 | to_device(edges_propagate), to_device(label), to_device(norms) 232 | all_size = int(max([item.max() for item in idx_data]) + 1) 233 | pred = torch.zeros((all_size, args.num_classes)).cuda() 234 | all_labels = torch.zeros((all_size), dtype=torch.long).cuda() 235 | for i in range(len(features)): 236 | with torch.no_grad(): 237 | torch.cuda.synchronize() 238 | st = time.time() 239 | pred_part = model(features[i], pointclouds[i], edges_self[i], edges_forward[i], edges_propagate[i], norms[i]) 240 | 241 | torch.cuda.synchronize() 242 | et = time.time() 243 | time_list.append(et - st) 244 | pred_part = pred_part.contiguous().view(-1, args.num_classes) 245 | pred_part = F.softmax(pred_part, dim=-1) 246 | 247 | torch.cuda.empty_cache() 248 | pred[idx_data[i], :] += pred_part 249 | all_labels[idx_data[i]] = label[i] 250 | logger.info('Test: {}/{}, {}, {}, {}'.format(idx + 1, 251 | len(dataset_loader), 252 | i + 1, idx_data[i].shape[0], et - st)) 253 | 254 | pred = pred / (pred.sum(-1)[:, None] + 1e-8) 255 | # pred = pred.max(1)[1].data.cpu().numpy() 256 | raw_coord = val_dataset.get_raw_coord(idx) 257 | 258 | np.save(os.path.join(cur_prob_dir, scene_name + '.npy'), 259 | {'pred': pred, 'target': label, 'xyz': raw_coord}) 260 | 261 | if scene_name in output_dict: 262 | output_dict[scene_name]['pred'] += pred 263 | else: 264 | cur_dict = { 265 | 'pred': pred, 266 | 'target': all_labels, 267 | 'xyz': raw_coord 268 | } 269 | output_dict[scene_name] = cur_dict 270 | torch.cuda.empty_cache() 271 | 272 | pcd = o3d.geometry.PointCloud() 273 | for scene_name, output_data in output_dict.items(): 274 | print('scene_name : ', scene_name) 275 | target = output_data['target'].cpu().numpy().astype(np.int32) 276 | pred = output_data['pred'].cpu().numpy() 277 | xyz = output_data['xyz'] 278 | 279 | output = pred.argmax(axis=1) 280 | 281 | intersection, union, target = intersectionAndUnion( 282 | output, target, args.num_classes, args.ignore_label) 283 | intersection_meter.update(intersection) 284 | union_meter.update(union) 285 | target_meter.update(target) 286 | 287 | output = mapper[output.astype(np.int32)] 288 | pcd.points = o3d.utility.Vector3dVector(xyz) 289 | pcd.colors = o3d.utility.Vector3dVector(colormap[output]) 290 | o3d.io.write_point_cloud(str(ply_dir) + '/' + scene_name, pcd) 291 | 292 | fp = open(str(txt_dir) + '/' + scene_name.replace("_vh_clean_2.ply", ".txt"), "w") 293 | for l in output: 294 | fp.write(str(int(l)) + '\n') 295 | fp.close() 296 | 297 | iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) 298 | accuracy_class = intersection_meter.sum / (target_meter.sum + 1e-10) 299 | mIoU = np.mean(iou_class) 300 | mAcc = np.mean(accuracy_class) 301 | allAcc = sum(intersection_meter.sum) / (sum(target_meter.sum) + 1e-10) 302 | logger.info( 303 | 'Val result: mIoU/mAcc/allAcc {:.4f}/{:.4f}/{:.4f}.'.format(mIoU, mAcc, allAcc)) 304 | for i in range(args.num_classes): 305 | logger.info('Class_{} Result: iou/accuracy {:.4f}/{:.4f}.'.format(i, 306 | iou_class[i], accuracy_class[i])) 307 | logger.info('<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<') 308 | 309 | logger.info('Average running time per frame: ', np.mean(time_list)) 310 | 311 | 312 | def create_color_palette40(): 313 | return [ 314 | (0, 0, 0), 315 | (174, 199, 232), # wall 316 | (152, 223, 138), # floor 317 | (31, 119, 180), # cabinet 318 | (255, 187, 120), # bed 319 | (188, 189, 34), # chair 320 | (140, 86, 75), # sofa 321 | (255, 152, 150), # table 322 | (214, 39, 40), # door 323 | (197, 176, 213), # window 324 | (148, 103, 189), # bookshelf 325 | (196, 156, 148), # picture 326 | (23, 190, 207), # counter 327 | (178, 76, 76), 328 | (247, 182, 210), # desk 329 | (66, 188, 102), 330 | (219, 219, 141), # curtain 331 | (140, 57, 197), 332 | (202, 185, 52), 333 | (51, 176, 203), 334 | (200, 54, 131), 335 | (92, 193, 61), 336 | (78, 71, 183), 337 | (172, 114, 82), 338 | (255, 127, 14), # refrigerator 339 | (91, 163, 138), 340 | (153, 98, 156), 341 | (140, 153, 101), 342 | (158, 218, 229), # shower curtain 343 | (100, 125, 154), 344 | (178, 127, 135), 345 | (120, 185, 128), 346 | (146, 111, 194), 347 | (44, 160, 44), # toilet 348 | (112, 128, 144), # sink 349 | (96, 207, 209), 350 | (227, 119, 194), # bathtub 351 | (213, 92, 176), 352 | (94, 106, 211), 353 | (82, 84, 163), # otherfurn 354 | (100, 85, 144) 355 | ] 356 | 357 | 358 | if __name__ == '__main__': 359 | main_vote() 360 | -------------------------------------------------------------------------------- /scannet_data_loader_color_DDP.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | # Data Loader for ScanNet 6 | 7 | import random 8 | import torch 9 | import numpy as np 10 | import glob 11 | from torch.utils.data import Dataset 12 | import transforms as t 13 | from util.voxelize import voxelize 14 | 15 | from datasetCommon import subsample_and_knn, compute_weight, collect_fn 16 | 17 | 18 | class ScanNetDataset(Dataset): 19 | def __init__(self, cfg, dataset="training", rotate_deg=0.0, rotate_aug=False, flip_aug=False, scale_aug=False, transform_aug=False, color_aug=False, crop=False, shuffle_index=False, mix3D=False, voxelize_mode='random'): 20 | 21 | self.data = [] 22 | self.cfg = cfg 23 | self.set = dataset 24 | 25 | self.rotate_aug = rotate_aug 26 | self.flip_aug = flip_aug 27 | self.scale_aug = scale_aug 28 | self.transform = transform_aug 29 | self.trans_std = [0.02, 0.02, 0.02] 30 | self.color_aug = color_aug 31 | self.crop = crop 32 | self.shuffle_index = shuffle_index 33 | self.mix3D = mix3D 34 | self.rotate_deg = rotate_deg 35 | self.voxelize_mode = voxelize_mode 36 | 37 | if 'noisy_points' in self.cfg and self.cfg.noisy_points: 38 | self.noisy_points = t.NoisyPoints(self.cfg.noise_level, self.cfg.noise_pct, self.cfg.ignore_label) 39 | self.add_noise = True 40 | else: 41 | self.add_noise = False 42 | 43 | # if self.color_aug: 44 | ''' 45 | color_transform = [t.ChromaticAutoContrast(), 46 | t.ChromaticTranslation(0.05), 47 | t.ChromaticJitter(0.05)] 48 | ''' 49 | # Random drop color is so important just let it be there now 50 | if self.color_aug: 51 | color_transform = [t.RandomDropColor()] 52 | self.color_transform = t.Compose(color_transform) 53 | 54 | if self.set == "training": 55 | data_files = glob.glob(self.cfg.train_data_path) 56 | elif self.set == "validation": 57 | data_files = glob.glob(self.cfg.val_data_path) 58 | elif self.set == "trainval": 59 | data_files_train = glob.glob(self.cfg.train_data_path) 60 | data_files_val = glob.glob(self.cfg.val_data_path) 61 | data_files = data_files_train + data_files_val 62 | else: # 'test' 63 | data_files = glob.glob(self.cfg.test_data_path) 64 | 65 | for x in torch.utils.data.DataLoader( 66 | data_files, 67 | collate_fn=lambda x: torch.load(x[0]), num_workers=cfg.NUM_WORKERS): 68 | self.data.append(x) 69 | 70 | print('%s examples: %d' % (self.set, len(self.data))) 71 | 72 | if self.cfg.USE_WEIGHT: 73 | weights = compute_weight(self.data) 74 | else: 75 | weights = [1.0] * 20 76 | print("label weights", weights) 77 | self.cfg.weights = weights 78 | 79 | def __len__(self): 80 | """ 81 | Return the length of data here 82 | """ 83 | return len(self.data) 84 | 85 | def _augment_data(self, coord, color, norm, label): 86 | ######################### 87 | # random data augmentation 88 | ######################### 89 | 90 | # random augmentation by rotation 91 | if self.rotate_aug: 92 | # rotate_rad = np.deg2rad(np.random.random() * 360) - np.pi 93 | rotate_rad = np.deg2rad(torch.rand(1).numpy()[0] * 360) - np.pi 94 | c, s = np.cos(rotate_rad), np.sin(rotate_rad) 95 | j = np.matrix([[c, s], [-s, c]]) 96 | coord[:, :2] = np.dot(coord[:, :2], j) 97 | norm[:, :2] = np.dot(norm[:, :2], j) 98 | # print(rotate_rad) 99 | 100 | # random augmentation by flip x, y or x+y 101 | if self.flip_aug: 102 | flip_type = torch.randint(4, (1,)).numpy()[0] # np.random.choice(4, 1) 103 | # print(flip_type) 104 | if flip_type == 1: 105 | coord[:, 0] = -coord[:, 0] 106 | norm[:, 0] = -norm[:, 0] 107 | elif flip_type == 2: 108 | coord[:, 1] = -coord[:, 1] 109 | norm[:, 1] = -norm[:, 1] 110 | elif flip_type == 3: 111 | coord[:, :2] = -coord[:, :2] 112 | norm[:, :2] = -norm[:, :2] 113 | 114 | if self.scale_aug: 115 | noise_scale = torch.rand(1).numpy()[0] * 0.4 + 0.8 # np.random.uniform(0.8, 1.2) 116 | # print(noise_scale) 117 | coord[:, 0] = noise_scale * coord[:, 0] 118 | coord[:, 1] = noise_scale * coord[:, 1] 119 | 120 | if self.transform: 121 | # noise_translate = np.array([np.random.normal(0, self.trans_std[0], 1), 122 | # np.random.normal(0, self.trans_std[1], 1), 123 | # np.random.normal(0, self.trans_std[2], 1)]).T 124 | # noise_translate = np.array([torch.randn(1).numpy()[0] * self.trans_std[0], 125 | # torch.randn(1).numpy()[0] * self.trans_std[1], 126 | # torch.randn(1).numpy()[0] * self.trans_std[2]]).T 127 | num_points = coord.shape[0] 128 | noise_translate = torch.randn(num_points, 3).numpy() 129 | noise_translate[:, 0] *= self.trans_std[0] 130 | noise_translate[:, 1] *= self.trans_std[1] 131 | noise_translate[:, 2] *= self.trans_std[2] 132 | # print("before range: ", coord.min(0), coord.max(0)) 133 | coord[:, 0:3] += noise_translate 134 | # print("after range: ", coord.min(0), coord.max(0)) 135 | # print(noise_translate) 136 | 137 | if self.color_aug: 138 | # color = (color + 1) * 127.5 139 | # color *= 255. 140 | coord, color, label, norm = self.color_transform(coord, color, label, norm) 141 | # color = color / 127.5 - 1 142 | # color /= 255. 143 | 144 | # crop half of the scene 145 | if self.crop: 146 | points = coord - coord.mean(0) 147 | if torch.rand(1).numpy()[0] < 0.5: 148 | inds = np.all([points[:, 0] >= 0.0], axis=0) 149 | else: 150 | inds = np.all([points[:, 0] < 0.0], axis=0) 151 | 152 | coord, color, norm, label = ( 153 | coord[~inds], 154 | color[~inds], 155 | norm[~inds], 156 | label[~inds] 157 | ) 158 | 159 | return coord, color, norm, label 160 | 161 | def get_raw_coord(self, indx): 162 | return self.data[indx][0] 163 | 164 | def get_scene_name(self, indx): 165 | return self.data[indx][3] 166 | 167 | def __getitem__(self, indx): 168 | coord, features, label, _ = self.data[indx] 169 | 170 | color, norm = features[:, :3], features[:, 3:] 171 | 172 | # move z to 0+ 173 | z_min = coord[:, 2].min() 174 | coord[:, 2] -= z_min 175 | 176 | # Specific rotation 177 | if self.rotate_deg != 0.: 178 | rotate_rad = np.deg2rad(self.rotate_deg * 360) - np.pi 179 | c, s = np.cos(rotate_rad), np.sin(rotate_rad) 180 | j = np.matrix([[c, s], [-s, c]]) 181 | coord[:, :2] = np.dot(coord[:, :2], j) 182 | norm[:, :2] = np.dot(norm[:, :2], j) 183 | 184 | coord, color, norm, label = self._augment_data(coord, color, norm, label) 185 | 186 | # Add noise to examine robustness 187 | if self.add_noise: 188 | coords, color, norm, label = self.noisy_points(coord, color, norm, label) 189 | 190 | 191 | # Mix3D augmentation from 3DV 2021 paper 192 | if self.mix3D and random.random() < 0.8: 193 | mix_idx = np.random.randint(len(self.data)) 194 | coords2, features2, labels2, _ = self.data[mix_idx] 195 | 196 | color2, norm2 = features2[:, :3], features2[:, 3:] 197 | z_min = coords2[:, 2].min() 198 | coords2[:, 2] -= z_min 199 | coords2, color2, norm2, labels2 = self._augment_data(coords2, color2, norm2, labels2) 200 | coord = np.concatenate((coord, coords2), axis=0) 201 | color = np.concatenate((color, color2), axis=0) 202 | norm = np.concatenate((norm, norm2), axis=0) 203 | label = np.concatenate((label, labels2), axis=0) 204 | 205 | # input normalize 206 | coord_min = np.min(coord, 0) 207 | coord -= coord_min 208 | 209 | # voxelize coords 210 | uniq_idx = voxelize(coord, self.cfg.grid_size[0], mode=self.voxelize_mode) 211 | if self.voxelize_mode != 'multiple': 212 | all_data = {} 213 | uniq_idx = voxelize(coord, self.cfg.grid_size[0], mode=self.voxelize_mode) 214 | coord, color, norm, label = coord[uniq_idx], color[uniq_idx], norm[uniq_idx], label[uniq_idx] 215 | else: 216 | all_data = [] 217 | # For testing set, output multiple sets of data for the purpose of testing on all points before voxelization 218 | for crop_idx in uniq_idx: 219 | data = {} 220 | coord_part, norm_part = coord[crop_idx], norm[crop_idx] 221 | point_list, nei_forward_list, nei_propagate_list, nei_self_list, norm_list \ 222 | = subsample_and_knn(coord_part, norm_part, grid_size=self.cfg.grid_size, K_self=self.cfg.K_self, 223 | K_forward=self.cfg.K_forward, K_propagate=self.cfg.K_propagate) 224 | data['point_list'] = point_list 225 | data['nei_forward_list'] = nei_forward_list 226 | data['nei_propagate_list'] = nei_propagate_list 227 | data['nei_self_list'] = nei_self_list 228 | data['surface_normal_list'] = norm_list 229 | data['feature_list'] = [color[crop_idx].astype(np.float32)] 230 | data['label_list'] = [label[crop_idx].astype(np.int32)] 231 | data['crop_idx'] = crop_idx 232 | all_data.append(data) 233 | return all_data 234 | 235 | # subsample points during training time 236 | if (self.set == "training" or self.set == "trainval") and label.shape[0] > self.cfg.MAX_POINTS_NUM: 237 | init_idx = torch.randint(label.shape[0], (1,)).numpy()[0] 238 | crop_idx = np.argsort(np.sum(np.square(coord - coord[init_idx]), 1))[:self.cfg.MAX_POINTS_NUM] 239 | coord, color, norm, label = coord[crop_idx], color[crop_idx], norm[crop_idx], label[crop_idx] 240 | 241 | # shuffle_index 242 | if self.shuffle_index: 243 | shuf_idx = torch.randperm(coord.shape[0]).numpy() 244 | coord, color, norm, label = coord[shuf_idx], color[shuf_idx], norm[shuf_idx], label[shuf_idx] 245 | 246 | 247 | # print("number of points: ", coord.shape[0]) 248 | 249 | point_list, nei_forward_list, nei_propagate_list, nei_self_list, norm_list = \ 250 | subsample_and_knn(coord, norm, grid_size=self.cfg.grid_size, K_self=self.cfg.K_self, 251 | K_forward=self.cfg.K_forward, K_propagate=self.cfg.K_propagate) 252 | all_data['point_list'] = point_list 253 | all_data['nei_forward_list'] = nei_forward_list 254 | all_data['nei_propagate_list'] = nei_propagate_list 255 | all_data['nei_self_list'] = nei_self_list 256 | all_data['surface_normal_list'] = norm_list 257 | all_data['feature_list'] = [color.astype(np.float32)] 258 | all_data['label_list'] = [label.astype(np.int32)] 259 | 260 | return all_data 261 | 262 | 263 | def getdataLoadersDDP(cfg): 264 | ''' 265 | Get ScanNet data loaders for the training and validation set. If cfg.DDP is set then a DistributedSampler is used, otherwise, a RandomSampler is used 266 | Input: 267 | cfg: config dictionary 268 | Output: 269 | train_loader: data loader for the training set 270 | val_loader: data loader for the validation set 271 | ''' 272 | # Initialize datasets 273 | if cfg.DDP: 274 | train_loader, val_loader, _, _ = getdataLoaders(cfg, torch.utils.data.distributed.DistributedSampler) 275 | else: 276 | train_loader, val_loader, _, _ = getdataLoaders(cfg, torch.utils.data.RandomSampler) 277 | return train_loader, val_loader 278 | 279 | 280 | def getdataLoaders(cfg, sampler): 281 | ''' 282 | Get ScanNet data loaders with the correct augmentations 283 | Output: 284 | train_loader: data loader for the training set 285 | val_loader: data loader for the validation set 286 | train_dataset: a ScanNetDataset object for the training set 287 | validation_dataset: a ScanNetDataset object for the validation set 288 | ''' 289 | # Initialize datasets 290 | training_dataset = ScanNetDataset(cfg, dataset="training", 291 | rotate_aug=cfg.rotate_aug, 292 | flip_aug=cfg.flip_aug, 293 | scale_aug=cfg.scale_aug, 294 | transform_aug=cfg.transform_aug, 295 | color_aug=cfg.color_aug, 296 | crop=cfg.crop, 297 | shuffle_index=cfg.shuffle_index, 298 | mix3D=cfg.mix3D) 299 | # No data augmentation for validation 300 | validation_dataset = ScanNetDataset(cfg, dataset="validation") 301 | 302 | training_sampler = sampler(training_dataset) 303 | 304 | validation_sampler = sampler(validation_dataset) 305 | 306 | train_data_loader = torch.utils.data.DataLoader(training_dataset, 307 | batch_size=cfg.BATCH_SIZE, 308 | collate_fn=collect_fn, 309 | num_workers=cfg.NUM_WORKERS, 310 | pin_memory=True, 311 | sampler=training_sampler, 312 | drop_last=True) 313 | 314 | val_data_loader = torch.utils.data.DataLoader(validation_dataset, 315 | batch_size=cfg.BATCH_SIZE, 316 | collate_fn=collect_fn, 317 | num_workers=cfg.NUM_WORKERS, 318 | sampler=validation_sampler, 319 | pin_memory=True) 320 | 321 | return train_data_loader, val_data_loader, training_dataset, validation_dataset 322 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_neighbors/wrapper.cpp: -------------------------------------------------------------------------------- 1 | // MIT License 2 | 3 | // Copyright (c) 2019 HuguesTHOMAS 4 | 5 | 6 | #include 7 | #include 8 | #include "neighbors/neighbors.h" 9 | #include 10 | 11 | 12 | 13 | // docstrings for our module 14 | // ************************* 15 | 16 | static char module_docstring[] = "This module provides two methods to compute radius neighbors from pointclouds or batch of pointclouds"; 17 | 18 | static char batch_query_docstring[] = "Method to get radius neighbors in a batch of stacked pointclouds"; 19 | static char batch_kquery_docstring[] = "Method to get k nearest neighbors in a batch of stacked pointclouds"; 20 | 21 | 22 | // Declare the functions 23 | // ********************* 24 | 25 | static PyObject *batch_neighbors(PyObject *self, PyObject *args, PyObject *keywds); 26 | 27 | static PyObject *batch_kneighbors(PyObject *self, PyObject *args, PyObject *keywds); 28 | 29 | // Specify the members of the module 30 | // ********************************* 31 | 32 | static PyMethodDef module_methods[] = 33 | { 34 | { "batch_query", (PyCFunction)batch_neighbors, METH_VARARGS | METH_KEYWORDS, batch_query_docstring }, 35 | { "batch_kquery", (PyCFunction)batch_kneighbors, METH_VARARGS | METH_KEYWORDS, batch_kquery_docstring }, 36 | {NULL, NULL, 0, NULL} 37 | }; 38 | 39 | 40 | // Initialize the module 41 | // ********************* 42 | 43 | static struct PyModuleDef moduledef = 44 | { 45 | PyModuleDef_HEAD_INIT, 46 | "radius_neighbors", // m_name 47 | module_docstring, // m_doc 48 | -1, // m_size 49 | module_methods, // m_methods 50 | NULL, // m_reload 51 | NULL, // m_traverse 52 | NULL, // m_clear 53 | NULL, // m_free 54 | }; 55 | 56 | PyMODINIT_FUNC PyInit_radius_neighbors(void) 57 | { 58 | import_array(); 59 | return PyModule_Create(&moduledef); 60 | } 61 | 62 | 63 | // Definition of the batch_subsample method 64 | // ********************************** 65 | 66 | static PyObject* batch_neighbors(PyObject* self, PyObject* args, PyObject* keywds) 67 | { 68 | 69 | // Manage inputs 70 | // ************* 71 | 72 | // Args containers 73 | PyObject* queries_obj = NULL; 74 | PyObject* supports_obj = NULL; 75 | PyObject* q_batches_obj = NULL; 76 | PyObject* s_batches_obj = NULL; 77 | 78 | // Keywords containers 79 | static char* kwlist[] = { "queries", "supports", "q_batches", "s_batches", "radius", NULL }; 80 | float radius = 0.1; 81 | 82 | // Parse the input 83 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "OOOO|$f", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &radius)) 84 | { 85 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 86 | return NULL; 87 | } 88 | 89 | 90 | // Interpret the input objects as numpy arrays. 91 | PyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY); 92 | PyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY); 93 | PyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY); 94 | PyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY); 95 | 96 | // Verify data was load correctly. 97 | if (queries_array == NULL) 98 | { 99 | Py_XDECREF(queries_array); 100 | Py_XDECREF(supports_array); 101 | Py_XDECREF(q_batches_array); 102 | Py_XDECREF(s_batches_array); 103 | PyErr_SetString(PyExc_RuntimeError, "Error converting query points to numpy arrays of type float32"); 104 | return NULL; 105 | } 106 | if (supports_array == NULL) 107 | { 108 | Py_XDECREF(queries_array); 109 | Py_XDECREF(supports_array); 110 | Py_XDECREF(q_batches_array); 111 | Py_XDECREF(s_batches_array); 112 | PyErr_SetString(PyExc_RuntimeError, "Error converting support points to numpy arrays of type float32"); 113 | return NULL; 114 | } 115 | if (q_batches_array == NULL) 116 | { 117 | Py_XDECREF(queries_array); 118 | Py_XDECREF(supports_array); 119 | Py_XDECREF(q_batches_array); 120 | Py_XDECREF(s_batches_array); 121 | PyErr_SetString(PyExc_RuntimeError, "Error converting query batches to numpy arrays of type int32"); 122 | return NULL; 123 | } 124 | if (s_batches_array == NULL) 125 | { 126 | Py_XDECREF(queries_array); 127 | Py_XDECREF(supports_array); 128 | Py_XDECREF(q_batches_array); 129 | Py_XDECREF(s_batches_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting support batches to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3) 136 | { 137 | Py_XDECREF(queries_array); 138 | Py_XDECREF(supports_array); 139 | Py_XDECREF(q_batches_array); 140 | Py_XDECREF(s_batches_array); 141 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : query.shape is not (N, 3)"); 142 | return NULL; 143 | } 144 | if ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3) 145 | { 146 | Py_XDECREF(queries_array); 147 | Py_XDECREF(supports_array); 148 | Py_XDECREF(q_batches_array); 149 | Py_XDECREF(s_batches_array); 150 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : support.shape is not (N, 3)"); 151 | return NULL; 152 | } 153 | if ((int)PyArray_NDIM(q_batches_array) > 1) 154 | { 155 | Py_XDECREF(queries_array); 156 | Py_XDECREF(supports_array); 157 | Py_XDECREF(q_batches_array); 158 | Py_XDECREF(s_batches_array); 159 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : queries_batches.shape is not (B,) "); 160 | return NULL; 161 | } 162 | if ((int)PyArray_NDIM(s_batches_array) > 1) 163 | { 164 | Py_XDECREF(queries_array); 165 | Py_XDECREF(supports_array); 166 | Py_XDECREF(q_batches_array); 167 | Py_XDECREF(s_batches_array); 168 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : supports_batches.shape is not (B,) "); 169 | return NULL; 170 | } 171 | if ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0)) 172 | { 173 | Py_XDECREF(queries_array); 174 | Py_XDECREF(supports_array); 175 | Py_XDECREF(q_batches_array); 176 | Py_XDECREF(s_batches_array); 177 | PyErr_SetString(PyExc_RuntimeError, "Wrong number of batch elements: different for queries and supports "); 178 | return NULL; 179 | } 180 | 181 | // Number of points 182 | int Nq = (int)PyArray_DIM(queries_array, 0); 183 | int Ns= (int)PyArray_DIM(supports_array, 0); 184 | 185 | // Number of batches 186 | int Nb = (int)PyArray_DIM(q_batches_array, 0); 187 | 188 | // Call the C++ function 189 | // ********************* 190 | 191 | // Convert PyArray to Cloud C++ class 192 | vector queries; 193 | vector supports; 194 | vector q_batches; 195 | vector s_batches; 196 | queries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq); 197 | supports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns); 198 | q_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb); 199 | s_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb); 200 | 201 | // Create result containers 202 | vector neighbors_indices; 203 | 204 | // Compute results 205 | //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 206 | batch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 207 | 208 | // Check result 209 | if (neighbors_indices.size() < 1) 210 | { 211 | PyErr_SetString(PyExc_RuntimeError, "Error"); 212 | return NULL; 213 | } 214 | 215 | // Manage outputs 216 | // ************** 217 | 218 | // Maximal number of neighbors 219 | int max_neighbors = neighbors_indices.size() / Nq; 220 | 221 | // Dimension of output containers 222 | npy_intp* neighbors_dims = new npy_intp[2]; 223 | neighbors_dims[0] = Nq; 224 | neighbors_dims[1] = max_neighbors; 225 | 226 | // Create output array 227 | PyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_INT); 228 | PyObject* ret = NULL; 229 | 230 | // Fill output array with values 231 | size_t size_in_bytes = Nq * max_neighbors * sizeof(int); 232 | memcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes); 233 | 234 | // Merge results 235 | ret = Py_BuildValue("N", res_obj); 236 | 237 | // Clean up 238 | // ******** 239 | 240 | Py_XDECREF(queries_array); 241 | Py_XDECREF(supports_array); 242 | Py_XDECREF(q_batches_array); 243 | Py_XDECREF(s_batches_array); 244 | 245 | return ret; 246 | } 247 | 248 | 249 | // Definition of the batch_subsample method 250 | // ********************************** 251 | 252 | static PyObject* batch_kneighbors(PyObject* self, PyObject* args, PyObject* keywds) 253 | { 254 | 255 | // Manage inputs 256 | // ************* 257 | 258 | // Args containers 259 | PyObject* queries_obj = NULL; 260 | PyObject* supports_obj = NULL; 261 | PyObject* q_batches_obj = NULL; 262 | PyObject* s_batches_obj = NULL; 263 | 264 | // Keywords containers 265 | static char* kwlist[] = { "queries", "supports", "q_batches", "s_batches", "K", NULL }; 266 | int K = 16; 267 | 268 | // Parse the input 269 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "OOOO|i", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &K)) 270 | { 271 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 272 | return NULL; 273 | } 274 | 275 | 276 | // Interpret the input objects as numpy arrays. 277 | PyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY); 278 | PyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY); 279 | PyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY); 280 | PyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY); 281 | 282 | // Verify data was load correctly. 283 | if (queries_array == NULL) 284 | { 285 | Py_XDECREF(queries_array); 286 | Py_XDECREF(supports_array); 287 | Py_XDECREF(q_batches_array); 288 | Py_XDECREF(s_batches_array); 289 | PyErr_SetString(PyExc_RuntimeError, "Error converting query points to numpy arrays of type float32"); 290 | return NULL; 291 | } 292 | if (supports_array == NULL) 293 | { 294 | Py_XDECREF(queries_array); 295 | Py_XDECREF(supports_array); 296 | Py_XDECREF(q_batches_array); 297 | Py_XDECREF(s_batches_array); 298 | PyErr_SetString(PyExc_RuntimeError, "Error converting support points to numpy arrays of type float32"); 299 | return NULL; 300 | } 301 | if (q_batches_array == NULL) 302 | { 303 | Py_XDECREF(queries_array); 304 | Py_XDECREF(supports_array); 305 | Py_XDECREF(q_batches_array); 306 | Py_XDECREF(s_batches_array); 307 | PyErr_SetString(PyExc_RuntimeError, "Error converting query batches to numpy arrays of type int32"); 308 | return NULL; 309 | } 310 | if (s_batches_array == NULL) 311 | { 312 | Py_XDECREF(queries_array); 313 | Py_XDECREF(supports_array); 314 | Py_XDECREF(q_batches_array); 315 | Py_XDECREF(s_batches_array); 316 | PyErr_SetString(PyExc_RuntimeError, "Error converting support batches to numpy arrays of type int32"); 317 | return NULL; 318 | } 319 | 320 | // Check that the input array respect the dims 321 | if ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3) 322 | { 323 | Py_XDECREF(queries_array); 324 | Py_XDECREF(supports_array); 325 | Py_XDECREF(q_batches_array); 326 | Py_XDECREF(s_batches_array); 327 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : query.shape is not (N, 3)"); 328 | return NULL; 329 | } 330 | if ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3) 331 | { 332 | Py_XDECREF(queries_array); 333 | Py_XDECREF(supports_array); 334 | Py_XDECREF(q_batches_array); 335 | Py_XDECREF(s_batches_array); 336 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : support.shape is not (N, 3)"); 337 | return NULL; 338 | } 339 | if ((int)PyArray_NDIM(q_batches_array) > 1) 340 | { 341 | Py_XDECREF(queries_array); 342 | Py_XDECREF(supports_array); 343 | Py_XDECREF(q_batches_array); 344 | Py_XDECREF(s_batches_array); 345 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : queries_batches.shape is not (B,) "); 346 | return NULL; 347 | } 348 | if ((int)PyArray_NDIM(s_batches_array) > 1) 349 | { 350 | Py_XDECREF(queries_array); 351 | Py_XDECREF(supports_array); 352 | Py_XDECREF(q_batches_array); 353 | Py_XDECREF(s_batches_array); 354 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : supports_batches.shape is not (B,) "); 355 | return NULL; 356 | } 357 | if ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0)) 358 | { 359 | Py_XDECREF(queries_array); 360 | Py_XDECREF(supports_array); 361 | Py_XDECREF(q_batches_array); 362 | Py_XDECREF(s_batches_array); 363 | PyErr_SetString(PyExc_RuntimeError, "Wrong number of batch elements: different for queries and supports "); 364 | return NULL; 365 | } 366 | 367 | // Number of points 368 | int Nq = (int)PyArray_DIM(queries_array, 0); 369 | int Ns= (int)PyArray_DIM(supports_array, 0); 370 | 371 | // Number of batches 372 | int Nb = (int)PyArray_DIM(q_batches_array, 0); 373 | 374 | // Call the C++ function 375 | // ********************* 376 | 377 | // Convert PyArray to Cloud C++ class 378 | vector queries; 379 | vector supports; 380 | vector q_batches; 381 | vector s_batches; 382 | queries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq); 383 | supports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns); 384 | q_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb); 385 | s_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb); 386 | 387 | // Create result containers 388 | vector neighbors_indices; 389 | 390 | // Compute results 391 | //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 392 | batch_nanoflann_knn(queries, supports, q_batches, s_batches, neighbors_indices, K); 393 | 394 | // Check result 395 | if (neighbors_indices.size() < 1) 396 | { 397 | PyErr_SetString(PyExc_RuntimeError, "Error"); 398 | return NULL; 399 | } 400 | 401 | // Manage outputs 402 | // ************** 403 | 404 | // Dimension of output containers 405 | npy_intp* neighbors_dims = new npy_intp[2]; 406 | neighbors_dims[0] = Nq; 407 | neighbors_dims[1] = K; 408 | 409 | // Create output array 410 | PyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_ULONG); 411 | PyObject* ret = NULL; 412 | 413 | // Fill output array with values 414 | size_t size_in_bytes = Nq * K * sizeof(size_t); 415 | memcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes); 416 | 417 | // Merge results 418 | ret = Py_BuildValue("N", res_obj); 419 | 420 | // Clean up 421 | // ******** 422 | 423 | Py_XDECREF(queries_array); 424 | Py_XDECREF(supports_array); 425 | Py_XDECREF(q_batches_array); 426 | Py_XDECREF(s_batches_array); 427 | 428 | return ret; 429 | } 430 | -------------------------------------------------------------------------------- /datasetCommon.py: -------------------------------------------------------------------------------- 1 | # 2 | # For licensing see accompanying LICENSE file. 3 | # Copyright (C) 2022-2023 Apple Inc. All Rights Reserved. 4 | # 5 | 6 | import torch 7 | import numpy as np 8 | from sklearn.neighbors import KDTree 9 | import cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 10 | try: 11 | import cpp_wrappers.cpp_neighbors.radius_neighbors as cpp_neighbors 12 | except BaseException: 13 | print('Failed to import cpp_neighbors, nanoflann kNN is not loaded. Only sklearn kNN is available') 14 | 15 | 16 | def grid_subsampling( 17 | points, 18 | features=None, 19 | labels=None, 20 | sampleDl=0.1, 21 | verbose=0): 22 | """ 23 | This function comes from https://github.com/HuguesTHOMAS/KPConv 24 | Copyright Hugues Thomas 2018 25 | CPP wrapper for a grid subsampling (method = barycenter for points and features 26 | Input: 27 | points: (N, 3) matrix of input points 28 | features: optional (N, d) matrix of features (floating number) 29 | labels: optional (N,) matrix of integer labels 30 | sampleDl: parameter defining the size of grid voxels 31 | verbose: 1 to display 32 | subsampled points, with features and/or labels depending of the input 33 | Output: subsampled points 34 | """ 35 | 36 | # method = "voxelcenters" # "barycenters" "voxelcenters" 37 | method = "barycenters" 38 | 39 | if (features is None) and (labels is None): 40 | return cpp_subsampling.compute( 41 | points, 42 | sampleDl=sampleDl, 43 | verbose=verbose, 44 | method=method) 45 | elif (labels is None): 46 | return cpp_subsampling.compute( 47 | points, 48 | features=features, 49 | sampleDl=sampleDl, 50 | verbose=verbose, 51 | method=method) 52 | elif (features is None): 53 | return cpp_subsampling.compute( 54 | points, 55 | classes=labels, 56 | sampleDl=sampleDl, 57 | verbose=verbose, 58 | method=method) 59 | else: 60 | return cpp_subsampling.compute( 61 | points, 62 | features=features, 63 | classes=labels, 64 | sampleDl=sampleDl, 65 | verbose=verbose, 66 | method=method) 67 | 68 | 69 | def compute_weight(train_data, num_class=20): 70 | """ 71 | Compute the class weights for ScanNet classes based on square root of inverse proportional weighting 72 | Input: 73 | training data: (points, features, labels, norms), only labels are used 74 | Output: 75 | a list of weights for each class 76 | """ 77 | weights = np.array([0.0 for i in range(num_class)]) 78 | 79 | num_rooms = len(train_data) 80 | for i in range(num_rooms): 81 | _, _, labels, _ = train_data[i] 82 | # rm invalid labels 83 | labels = labels[labels >= 0] 84 | for j in range(num_class): 85 | weights[j] += np.sum(labels == j) 86 | 87 | ratio = weights / float(sum(weights)) 88 | ce_label_weight = 1 / (np.power(ratio, 1 / 2)) 89 | return list(ce_label_weight) 90 | 91 | 92 | def compute_knn(ref_points, query_points, K, dilated_rate=1, method='sklearn'): 93 | """ 94 | Compute KNN 95 | Input: 96 | ref_points: reference points (MxD) 97 | query_points: query points (NxD) 98 | K: the amount of neighbors for each point 99 | dilated_rate: If set to larger than 1, then select more neighbors and then choose from them 100 | (Engelmann et al. Dilated Point Convolutions: On the Receptive Field Size of Point Convolutions on 3D Point Clouds. ICRA 2020) 101 | method: Choose between two approaches: Scikit-Learn ('sklearn') or nanoflann ('nanoflann'). In general nanoflann should be faster, but sklearn is more stable 102 | Output: 103 | neighbors_idx: for each query point, its K nearest neighbors among the reference points (N x K) 104 | """ 105 | num_ref_points = ref_points.shape[0] 106 | 107 | if num_ref_points < K or num_ref_points < dilated_rate * K: 108 | num_query_points = query_points.shape[0] 109 | inds = np.random.choice( 110 | num_ref_points, (num_query_points, K)).astype( 111 | np.int32) 112 | 113 | return inds 114 | if method == 'sklearn': 115 | kdt = KDTree(ref_points) 116 | neighbors_idx = kdt.query( 117 | query_points, 118 | k=K * dilated_rate, 119 | return_distance=False) 120 | elif method == 'nanoflann': 121 | neighbors_idx = batch_neighbors( 122 | query_points, ref_points, [ 123 | query_points.shape[0]], [num_ref_points], K * dilated_rate) 124 | else: 125 | raise Exception('compute_knn: unsupported knn algorithm') 126 | if dilated_rate > 1: 127 | neighbors_idx = np.array( 128 | neighbors_idx[:, ::dilated_rate], dtype=np.int32) 129 | 130 | return neighbors_idx 131 | 132 | 133 | def crop(points, x_min, y_min, z_min, x_max, y_max, z_max): 134 | """ 135 | Crop all points within a 3D bounding box defined by (x_min, y_min, z_min) and (x_max, y_max, z_max) 136 | Input: 137 | points: an array of points (nx3) 138 | x_min, y_min, z_min, x_max, y_max, z_max: bounding box extremal coordinates 139 | Output: 140 | inds: indices of the points that are within the bounding box defined by (x_min, y_min, z_min) and (x_max, y_max, z_max) 141 | """ 142 | if x_max <= x_min or y_max <= y_min or z_max <= z_min: 143 | raise ValueError( 144 | "We should have x_min < x_max and y_min < y_max and z_min < z_max. But we got" 145 | " (x_min = {x_min}, y_min = {y_min}, z_min = {z_min}," 146 | " x_max = {x_max}, y_max = {y_max}, z_max = {z_max})".format( 147 | x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max, z_min=z_min, z_max=z_max, )) 148 | inds = np.all( 149 | [ 150 | (points[:, 0] >= x_min), 151 | (points[:, 0] < x_max), 152 | (points[:, 1] >= y_min), 153 | (points[:, 1] < y_max), 154 | (points[:, 2] >= z_min), 155 | (points[:, 2] < z_max), 156 | ], 157 | axis=0, 158 | ) 159 | return inds 160 | 161 | 162 | def tensorizeList(nplist, is_index=False): 163 | """ 164 | Make all numpy arrays inside a list into torch tensors 165 | """ 166 | ret_list = [] 167 | for npitem in nplist: 168 | if is_index: 169 | if npitem is None: 170 | ret_list.append(None) 171 | else: 172 | ret_list.append( 173 | torch.from_numpy( 174 | npitem).long().unsqueeze(0)) 175 | else: 176 | ret_list.append(torch.from_numpy(npitem).float().unsqueeze(0)) 177 | 178 | return ret_list 179 | 180 | 181 | def tensorize( 182 | features, 183 | pointclouds, 184 | edges_self, 185 | edges_forward, 186 | edges_propagate, 187 | target, 188 | norms): 189 | """ 190 | Convert numpy arrays from inside lists into torch tensors for all input data 191 | """ 192 | pointclouds = tensorizeList(pointclouds) 193 | norms = tensorizeList(norms) 194 | edges_self = tensorizeList(edges_self, True) 195 | edges_forward = tensorizeList(edges_forward, True) 196 | edges_propagate = tensorizeList(edges_propagate, True) 197 | 198 | target = torch.from_numpy(target).long().unsqueeze(0) 199 | features = torch.from_numpy(features).float().unsqueeze(0) 200 | 201 | return features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms 202 | 203 | 204 | def listToBatch( 205 | features, 206 | pointclouds, 207 | edges_self, 208 | edges_forward, 209 | edges_propagate, 210 | target, 211 | norms): 212 | """ 213 | ListToBatch transforms a batch of multiple clouds into one point cloud so that we do not have to pad them to the same length 214 | The way this works is that all point clouds are concatenated one after another e.g., if you have point cloud 1 which is [5154,3], point cloud 2 which is [4749, 3] 215 | then it creates a point cloud as if it has batch size 1, which is a tensor of shape [1, 5154+4749, 3] 216 | It also modifies the edges (indices of k-nearest-neighbors) so that they point to the correct points 217 | For example, for point cloud 2, we add 5154 to all its neighbor indices so that they 218 | link to the points in point cloud 2 in this combined tensor 219 | Input: List versions of all the input 220 | Output: Batched versions of all the input 221 | """ 222 | # import ipdb; ipdb.set_trace() 223 | num_sample = len(pointclouds) 224 | 225 | # process sample 0 226 | featureBatch = features[0][0] 227 | pointcloudsBatch = pointclouds[0] 228 | pointcloudsNormsBatch = norms[0] 229 | if target: 230 | targetBatch = target[0][0] 231 | else: 232 | targetBatch = np.array(0) 233 | 234 | edgesSelfBatch = edges_self[0] 235 | edgesForwardBatch = edges_forward[0] 236 | edgesPropagateBatch = edges_propagate[0] 237 | 238 | points_stored = [val.shape[0] for val in pointcloudsBatch] 239 | 240 | for i in range(1, num_sample): 241 | if target: 242 | targetBatch = np.concatenate([targetBatch, target[i][0]], 0) 243 | featureBatch = np.concatenate([featureBatch, features[i][0]], 0) 244 | 245 | for j in range(len(edges_forward[i])): 246 | tempMask = edges_forward[i][j] == -1 247 | edges_forwardAdd = edges_forward[i][j] + points_stored[j] 248 | edges_forwardAdd[tempMask] = -1 249 | edgesForwardBatch[j] = np.concatenate([edgesForwardBatch[j], 250 | edges_forwardAdd], 0) 251 | 252 | tempMask2 = edges_propagate[i][j] == -1 253 | edges_propagateAdd = edges_propagate[i][j] + points_stored[j + 1] 254 | edges_propagateAdd[tempMask2] = -1 255 | edgesPropagateBatch[j] = np.concatenate([edgesPropagateBatch[j], 256 | edges_propagateAdd], 0) 257 | 258 | for j in range(len(pointclouds[i])): 259 | tempMask3 = edges_self[i][j] == -1 260 | edges_selfAdd = edges_self[i][j] + points_stored[j] 261 | edges_selfAdd[tempMask3] = -1 262 | edgesSelfBatch[j] = np.concatenate([edgesSelfBatch[j], 263 | edges_selfAdd], 0) 264 | 265 | pointcloudsBatch[j] = np.concatenate( 266 | [pointcloudsBatch[j], pointclouds[i][j]], 0) 267 | pointcloudsNormsBatch[j] = np.concatenate( 268 | [pointcloudsNormsBatch[j], norms[i][j]], 0) 269 | 270 | points_stored[j] += pointclouds[i][j].shape[0] 271 | 272 | return featureBatch, pointcloudsBatch, edgesSelfBatch, edgesForwardBatch, edgesPropagateBatch, \ 273 | targetBatch, pointcloudsNormsBatch 274 | 275 | 276 | def prepare( 277 | features, 278 | pointclouds, 279 | edges_self, 280 | edges_forward, 281 | edges_propagate, 282 | target, 283 | norms): 284 | """ 285 | Prepare data coming from data loader (lists of numpy arrays) into torch tensors ready to send to training 286 | """ 287 | 288 | features_out, pointclouds_out, edges_self_out, edges_forward_out, edges_propagate_out, target_out, norms_out = [], [], [], [], [], [], [] 289 | 290 | features_out, pointclouds_out, edges_self_out, edges_forward_out, edges_propagate_out, target_out, norms_out = \ 291 | listToBatch(features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms) 292 | 293 | features_out, pointclouds_out, edges_self_out, edges_forward_out, edges_propagate_out, target_out, norms_out = \ 294 | tensorize(features_out, pointclouds_out, edges_self_out, edges_forward_out, edges_propagate_out, target_out, norms_out) 295 | 296 | return features_out, pointclouds_out, edges_self_out, edges_forward_out, edges_propagate_out, target_out, norms_out 297 | 298 | 299 | def collect_fn(data_list): 300 | """ 301 | collect data from the data dictionary and outputs pytorch tensors 302 | """ 303 | features = [] 304 | pointclouds = [] 305 | target = [] 306 | norms = [] 307 | edges_forward = [] 308 | edges_propagate = [] 309 | edges_self = [] 310 | for i, data in enumerate(data_list): 311 | features.append(data['feature_list']) 312 | pointclouds.append(data['point_list']) 313 | if 'label_list' in data.keys(): 314 | target.append(data['label_list']) 315 | norms.append(data['surface_normal_list']) 316 | 317 | edges_forward.append(data['nei_forward_list']) 318 | edges_propagate.append(data['nei_propagate_list']) 319 | edges_self.append(data['nei_self_list']) 320 | 321 | features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms = \ 322 | prepare(features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms) 323 | 324 | return features, pointclouds, edges_self, edges_forward, edges_propagate, target, norms 325 | 326 | 327 | def subsample_and_knn( 328 | coord, 329 | norm, 330 | grid_size=[0.1], 331 | K_self=16, 332 | K_forward=16, 333 | K_propagate=16): 334 | """ 335 | Perform grid subsampling and compute kNN at each subsampling level 336 | Input: 337 | coord: N x 3 coordinates 338 | norm: N x 3 surface normals 339 | grid_size: all the downsampling levels (in cm) you want to use, e.g. [0.05, 0.1, 0.2, 0.4, 0.8] 340 | K_self: number of neighbors within each downsampling level 341 | K_forward: number of neighbors from one downsampling level to the next one (with less points), used in the downsampling PointConvs in the encoder 342 | K_propagate: number of neighbors from one downsampling level to the previous one (with more points), used in the upsampling PointConvs in the decoder 343 | Outputs: 344 | point_list: list of length len(grid_size) 345 | nei_forward_list: downsampling kNN neighbors (K_forward neighbors for each point) 346 | nei_propagate_list: upsampling kNN neighbors (K_propagate neighbors for each point) 347 | nei_self_list: kNN neighbors between the same layer 348 | norm_list: list of surface normals averaged within each voxel at each grid_size 349 | """ 350 | point_list, norm_list = [], [] 351 | nei_forward_list, nei_propagate_list, nei_self_list = [], [], [] 352 | 353 | for j, grid_s in enumerate(grid_size): 354 | if j == 0: 355 | sub_point, sub_norm = coord.astype( 356 | np.float32), norm.astype( 357 | np.float32) 358 | 359 | point_list.append(sub_point) 360 | norm_list.append(sub_norm) 361 | # compute edges 362 | nself = compute_knn(sub_point, sub_point, K_self[j]) 363 | nei_self_list.append(nself) 364 | 365 | else: 366 | sub_point, sub_norm = grid_subsampling( 367 | points=point_list[-1], features=norm_list[-1], sampleDl=grid_s) 368 | 369 | if sub_point.shape[0] <= K_self[j]: 370 | sub_point, sub_norm = point_list[-1], norm_list[-1] 371 | 372 | # compute edges, nforward is for downsampling, npropagate is for upsampling, 373 | # nself is for normal PointConv layers (between the same set of 374 | # points) 375 | nforward = compute_knn(point_list[-1], sub_point, K_forward[j]) 376 | npropagate = compute_knn(sub_point, point_list[-1], K_propagate[j]) 377 | nself = compute_knn(sub_point, sub_point, K_self[j]) 378 | # point_list is a list with len(grid_size) length, each item is a numpy array 379 | # of num_points x dimensionality 380 | 381 | point_list.append(sub_point) 382 | norm_list.append(sub_norm) 383 | nei_forward_list.append(nforward) 384 | nei_propagate_list.append(npropagate) 385 | nei_self_list.append(nself) 386 | 387 | return point_list, nei_forward_list, nei_propagate_list, nei_self_list, norm_list 388 | 389 | 390 | def getdataLoader(cfg, dataset, loader_set, sampler): 391 | """ 392 | Dataset should generate items as a dictionary with: 393 | color_list: colors 394 | point_list: point xyz coordinates 395 | surface_normal_list: surface normals 396 | nei_forward_list: neighbors from upper level to lower level 397 | nei_propagate_list: neighbors from lower level to upper level 398 | nei_self_list: neighbors at the same level 399 | label_list: optional labels, if no labels is given 400 | These can be prepared by the subsample_and_knn function which takes the input of coordinates, color, 401 | surface normals and optional labels. Surface normals can be computed with Open3D, but could also be 402 | obtained e.g. from meshes etc. 403 | """ 404 | this_dataset = dataset(cfg, dataset=loader_set) 405 | this_sampler = sampler(this_dataset) 406 | data_loader = torch.utils.data.DataLoader( 407 | this_dataset, 408 | batch_size=cfg.BATCH_SIZE, 409 | collate_fn=collect_fn, 410 | num_workers=cfg.NUM_WORKERS, 411 | sampler=this_sampler, 412 | pin_memory=True) 413 | return data_loader, this_dataset 414 | 415 | 416 | def batch_neighbors(queries, supports, q_batches, s_batches, K): 417 | """ 418 | Computes neighbors for a batch of queries and supports 419 | :param queries: (N1, 3) the query points 420 | :param supports: (N2, 3) the support points 421 | :param q_batches: (B) the list of lengths of batch elements in queries 422 | :param s_batches: (B)the list of lengths of batch elements in supports 423 | :param K: long 424 | :return: neighbors indices 425 | """ 426 | 427 | return cpp_neighbors.batch_kquery( 428 | queries, supports, q_batches, s_batches, K=int(K)) 429 | --------------------------------------------------------------------------------