├── LICENSE ├── README.md ├── assets └── teaser.png ├── checkpoints ├── model_cls.pth └── model_seg.pth ├── compute_cer.py ├── data └── PLACEHOLDER ├── data_prep.py ├── data_utils.py ├── docs └── 4972.pdf ├── environment.yaml ├── eval_modelnet40C.sh ├── eval_modelnet40noise.sh ├── logs └── PLACEHOLDER ├── main_cls.py ├── main_partseg.py ├── main_seg.py ├── misc.py └── models ├── diffConv_cls.py ├── diffConv_partseg.py ├── diffConv_seg.py ├── modules.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Manxi Lin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # diffConv: Analyzing Irregular Point Clouds with an Irregular View 2 | Standard spatial convolutions assume input data with a regular neighborhood structure. Existing methods typically generalize convolution to the irregular point cloud domain by fixing a regular "view" through e.g. a fixed neighborhood size, where the convolution kernel size remains the same for each point. However, since point clouds are not as structured as images, the fixed neighbor number gives an unfortunate inductive bias. We present a novel graph convolution named Difference Graph Convolution (diffConv), which does not rely on a regular view. diffConv operates on spatially-varying and density-dilated neighborhoods, which are further adapted by a learned masked attention mechanism. Experiments show that our model is very robust to the noise, obtaining state-of-the-art performance in 3D shape classification and scene understanding tasks, along with a faster inference speed. 3 | 4 | [[Arxiv]](https://arxiv.org/abs/2111.14658) [[ECCV]](https://www.ecva.net/papers/eccv_2022/papers_ECCV/papers/136630375.pdf) 5 | 6 | ![Alt text](assets/teaser.png) 7 | 8 | 9 | ## Dependencies 10 | - Python (tested on 3.7.11) 11 | - PyTorch (tested on 1.9.0) 12 | - CUDA (tested on 11.6) 13 | - other packages: sklearn, h5py, open3d 14 | - Install [CUDA accelerated PointNet++ library](https://github.com/daveredrum/Pointnet2.ScanNet/tree/master/pointnet2) under `models/pointnet2`. 15 | 16 | ## 3D Object Shape Classification 17 | ### ModelNet40 18 | **Prepare dataset** 19 | 20 | python3 data_prep.py --dataset=modelnet40 21 | 22 | **Train the model with default hyperparameters** 23 | 24 | python3 main_cls.py --exp_name=md40_cls --dataset=modelnet40 25 | 26 | There are many hyperparameters to customize, call 27 | 28 | python3 main_cls.py --help 29 | 30 | for details. 31 | 32 | **Evaluate with our pretrained model** 33 | 34 | python3 main_cls.py --exp_name=md40_cls_eval --dataset=modelnet40 --eval=True --model_path=checkpoints/model_cls.pth 35 | 36 | `--model_path` can be any trained parameters. 37 | 38 | **Evaluate model performance under noise** 39 | 40 | . eval_modelnet40noise.sh 41 | 42 | **Train model on resplited ModelNet40** 43 | 44 | python3 main_cls.py --exp_name=md40_resplit --dataset=modelnet40resplit 45 | 46 | Note that everytime the dataset is randomly resplitted. 47 | 48 | ### ModelNet40-C 49 | **Prepare dataset** 50 | 51 | Follow the [official instruction](https://github.com/jiachens/ModelNet40-C), then move `ModelNet40-C/data/modelnet40_c` to `data/modelnet40_c` folder. 52 | 53 | **Evaluate with our pretrained model** 54 | 55 | . eval_modelnet40C.sh 56 | 57 | ### ScanObjectNN 58 | **Prepare dataset** 59 | 60 | Download the [dataset](https://hkust-vgd.github.io/scanobjectnn/) and unzip it at `data/h5_files`. 61 | 62 | **Train the model with default hyperparameters** 63 | 64 | python3 main_cls.py --exp_name=sonn_cls --dataset=scanobjectnn --bg=False 65 | 66 | set `--bg` to `True` to train the model on the pointcloud with backgrounds. 67 | 68 | **Evaluation** 69 | 70 | Same as ModelNet40. 71 | 72 | ## 3D Scene Segmentation 73 | **NB: Please be aware that there could be an error on the Toronto3D segmentation, as reported in issues, causing the model to show constant (and anormal) IoU during training. I failed to reproduce the error when I ran the code from scratch, and thus not able to debug it. I am really sorry for this. My best guess is that this bug could be fixed by importing the anaconda environment from environment.yaml.** 74 | 75 | ### Toronto3D 76 | **Prepare dataset (may require torch 1.8.x)** 77 | 78 | Download the [dataset](https://github.com/WeikaiTan/Toronto-3D) and unzip it to `data/Toronto_3D`, then run 79 | 80 | python3 data_prep.py --dataset=toronto3d 81 | 82 | **Train the model with default hyperparameters** 83 | 84 | python3 main_seg.py --exp_name=trt_seg 85 | 86 | **Evaluate with our pretrained model** 87 | 88 | python3 main_seg.py --exp_name=trt_seg --eval=True --model_path=checkpoints/model_seg.pth 89 | 90 | ## 3D Object Shape Segmentation 91 | ### ShapeNetPart 92 | **Prepare dataset** 93 | 94 | python3 data_prep.py --dataset=shapenetpart 95 | 96 | **Train the model with default hyperparameters** 97 | 98 | python3 main_partseg.py --exp_name=spnetpt_seg 99 | 100 | **Evaluation** 101 | 102 | Same as other tasks. 103 | 104 | ## Citation 105 | Please cite this paper if you find this work helpful to your research, 106 | 107 | @inproceedings{lin2021diffconv, 108 | title={diffconv: Analyzing Irregular Point Clouds with an Irregular View}, 109 | author={Lin, Manxi and Feragen, Aasa}, 110 | booktitle={Proceedings of the European Conference on Computer Vision (ECCV)}, 111 | year={2022} 112 | } 113 | 114 | ## License 115 | MIT License 116 | 117 | ## Acknowledgements 118 | Part of this codebase is borrowed from [PointNet](https://github.com/charlesq34/pointnet), [DGCNN](https://github.com/WangYueFt/dgcnn), [dgcnn.pytorch](https://github.com/AnTao97/dgcnn.pytorch), [CurveNet](https://github.com/tiangexiang/CurveNet), [Pointnet2.ScanNet](https://github.com/daveredrum/Pointnet2.ScanNet). Sincere appreciation to their works! 119 | 120 | -------------------------------------------------------------------------------- /assets/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/assets/teaser.png -------------------------------------------------------------------------------- /checkpoints/model_cls.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/checkpoints/model_cls.pth -------------------------------------------------------------------------------- /checkpoints/model_seg.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/checkpoints/model_seg.pth -------------------------------------------------------------------------------- /compute_cer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | with open('logs/eval_modelnet40C.log', 'r') as f: 3 | logs = f.read() 4 | content = logs.split('test acc: ')[1:] 5 | cer = [1 - eval(c[:5]) for c in content] 6 | 7 | print('CER: %.4f'%np.mean(cer)) -------------------------------------------------------------------------------- /data/PLACEHOLDER: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/data/PLACEHOLDER -------------------------------------------------------------------------------- /data_prep.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import numpy as np 4 | from collections import defaultdict 5 | 6 | def split_Toronto3D(idx, n=8192, block_size=5, data_dir='./data'): 7 | import open3d.ml.torch as ml3d 8 | assert idx in [1,2,3,4] 9 | if idx in [1,3,4]: 10 | head = "train" 11 | elif idx == 2: 12 | head = "test" 13 | data_dir = os.path.join(data_dir, "Toronto_3D") 14 | # construct a dataset by specifying dataset_path 15 | dataset = ml3d.datasets.Toronto3D(dataset_path=data_dir) 16 | 17 | # get the 'all' split that combines training, validation and test set 18 | all_split = dataset.get_split('all') 19 | 20 | # print the shape of the first point cloud 21 | pointcloud = all_split.get_data(idx) 22 | 23 | pc = pointcloud['point'] 24 | label = pointcloud['label'] 25 | 26 | if not os.path.exists(os.path.join(data_dir, head)): 27 | os.mkdir(os.path.join(data_dir, head)) 28 | 29 | mask = np.any(np.isnan(pc), axis=1) 30 | 31 | pc = pc[~mask] 32 | 33 | pc_dict = defaultdict(list) 34 | label_dict = defaultdict(list) 35 | 36 | point_num = pc.shape[0] 37 | assert label.shape[0] == point_num 38 | 39 | print('There are %d points in the cloud'%point_num) 40 | 41 | for i in range(point_num): 42 | p = pc[i,:] 43 | x, y = p[0]//block_size, p[1]//block_size 44 | pc_dict['%d%d'%(x,y)].append(p) 45 | label_dict['%d%d'%(x,y)].append(label[i]) 46 | print("Processing %.2f%% points in Region %d"%(i/point_num*100, idx)) 47 | 48 | for k in pc_dict.keys(): 49 | pc = np.array(pc_dict[k]) 50 | label = np.array(label_dict[k]) 51 | if pc.shape[0] >= n: 52 | ind = np.arange(0, pc.shape[0], 1, np.int32) 53 | np.random.shuffle(ind) 54 | pc = pc[ind[:n], :] 55 | label = label[ind[:n]] 56 | np.save(os.path.join(data_dir, head, 'L%d_%s_point.npy'%(idx, k)), pc) 57 | np.save(os.path.join(data_dir, head, 'L%d_%s_label.npy'%(idx, k)), label) 58 | 59 | def make_folder(data_folder): 60 | base_dir = os.path.dirname(os.path.abspath(__file__)) 61 | data_dir = os.path.join(base_dir, data_folder) 62 | if not os.path.exists(data_dir): 63 | os.mkdir(data_dir) 64 | return data_dir 65 | 66 | def download_modelnet40(data_dir): 67 | if not os.path.exists(os.path.join(data_dir, 'modelnet40_ply_hdf5_2048')): 68 | site = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip' 69 | zipfile = os.path.basename(site) 70 | os.system('wget %s --no-check-certificate; unzip %s' % (site, zipfile)) 71 | os.system('mv %s %s' % ('modelnet40_ply_hdf5_2048', data_dir)) 72 | os.system('rm %s' % (zipfile)) 73 | 74 | print('ModelNet40 dataset has been prepared.') 75 | 76 | def download_toronto3d(data_dir): 77 | for i in range(1, 5): 78 | split_Toronto3D(i, data_dir=data_dir) 79 | 80 | print('Toronto3D dataset has been prepared.') 81 | 82 | def download_shapenetpart(data_dir): 83 | if not os.path.exists(os.path.join(data_dir, 'shapenet_part_seg_hdf5_data')): 84 | www = 'https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip' 85 | zipfile = os.path.basename(www) 86 | os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile)) 87 | os.system('mv %s %s' % ('hdf5_data', os.path.join(data_dir, 'shapenet_part_seg_hdf5_data'))) 88 | os.system('rm %s' % (zipfile)) 89 | 90 | print('ShapeNetPart dataset has been prepared.') 91 | 92 | if __name__=="__main__": 93 | parser = argparse.ArgumentParser(description='Data Preparation') 94 | parser.add_argument('--dataset', type=str, default='modelnet40', metavar='N', 95 | choices=['modelnet40', 'toronto3d', 'shapenetpart']) 96 | parser.add_argument('--data_folder', type=str, default='data') 97 | args = parser.parse_args() 98 | 99 | data_folder = args.data_folder 100 | data_dir = make_folder(data_folder) 101 | 102 | dataset = args.dataset 103 | 104 | eval("download_%s(data_dir)"%dataset) 105 | 106 | 107 | 108 | 109 | 110 | 111 | -------------------------------------------------------------------------------- /data_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | @Author: Yue Wang 3 | @Contact: yuewangx@mit.edu 4 | @File: data.py 5 | @Time: 2018/10/13 6:21 PM 6 | 7 | Modified by 8 | @Author: Manxi Lin 9 | @Contact: manli@dtu.dk 10 | @Time: 2022/7/7 3:00 PM 11 | """ 12 | import os 13 | import glob 14 | import h5py 15 | import numpy as np 16 | import torch 17 | from torch.utils.data import Dataset 18 | from models.utils import get_dist 19 | 20 | def load_data_mdn(partition): 21 | base_dir = os.path.dirname(os.path.abspath(__file__)) 22 | data_dir = os.path.join(base_dir, 'data') 23 | all_data = [] 24 | all_label = [] 25 | for h5_name in glob.glob(os.path.join(data_dir, 'modelnet40_ply_hdf5_2048', '*%s*.h5'%partition)): 26 | f = h5py.File(h5_name, 'r+') 27 | data = f['data'][:].astype('float32') 28 | label = f['label'][:].astype('int64') 29 | f.close() 30 | all_data.append(data) 31 | all_label.append(label) 32 | all_data = np.concatenate(all_data, axis=0) 33 | all_label = np.concatenate(all_label, axis=0) 34 | return all_data, all_label 35 | 36 | def load_data_sonn(partition, bg): 37 | base_dir = os.path.dirname(os.path.abspath(__file__)) 38 | data_dir = os.path.join(base_dir, 'data') 39 | if bg: 40 | head = 'main_split' 41 | else: 42 | head = 'main_split_nobg' 43 | if partition == 'train': 44 | partition = 'training' 45 | h5_name = os.path.join(data_dir, 'h5_files', head, '%s_objectdataset.h5'%(partition)) 46 | f = h5py.File(h5_name, 'r+') 47 | data = f['data'][:].astype('float32') 48 | label = f['label'][:].astype('int64') 49 | return data, label 50 | 51 | def load_data_seg(partition): 52 | base_dir = os.path.dirname(os.path.abspath(__file__)) 53 | data_dir = os.path.join(base_dir, 'data') 54 | point_dir = os.path.join(data_dir, 'Toronto_3D', partition, '*_point.npy') 55 | label_dir = os.path.join(data_dir, 'Toronto_3D', partition, '*_label.npy') 56 | all_data = glob.glob(point_dir) 57 | all_label = glob.glob(label_dir) 58 | 59 | return all_data, all_label 60 | 61 | 62 | def load_data_partseg(partition): 63 | base_dir = os.path.dirname(os.path.abspath(__file__)) 64 | data_dir = os.path.join(base_dir, 'data') 65 | all_data = [] 66 | all_label = [] 67 | all_seg = [] 68 | if partition == 'trainval': 69 | file = glob.glob(os.path.join(data_dir, 'shapenet_part_seg_hdf5_data', '*train*.h5')) \ 70 | + glob.glob(os.path.join(data_dir, 'shapenet_part_seg_hdf5_data', '*val*.h5')) 71 | else: 72 | file = glob.glob(os.path.join(data_dir, 'shapenet_part_seg_hdf5_data', '*%s*.h5'%partition)) 73 | for h5_name in file: 74 | f = h5py.File(h5_name, 'r+') 75 | data = f['data'][:].astype('float32') 76 | label = f['label'][:].astype('int64') 77 | seg = f['pid'][:].astype('int64') 78 | f.close() 79 | all_data.append(data) 80 | all_label.append(label) 81 | all_seg.append(seg) 82 | all_data = np.concatenate(all_data, axis=0) 83 | all_label = np.concatenate(all_label, axis=0) 84 | all_seg = np.concatenate(all_seg, axis=0) 85 | return all_data, all_label, all_seg 86 | 87 | 88 | def translate_pointcloud(pointcloud): 89 | xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) 90 | xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) 91 | 92 | translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') 93 | return translated_pointcloud 94 | 95 | def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02): 96 | N, C = pointcloud.shape 97 | pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip) 98 | return pointcloud 99 | 100 | def rotate_pointcloud(pointcloud): 101 | theta = np.pi*2 * np.random.uniform() 102 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]]) 103 | pointcloud[:,[0,2]] = pointcloud[:,[0,2]].dot(rotation_matrix) # random rotation (x,z) 104 | return pointcloud 105 | 106 | class ModelNet40(Dataset): 107 | def __init__(self, num_points, partition='train'): 108 | self.data, self.label = load_data_mdn(partition) 109 | self.num_points = num_points 110 | self.partition = partition 111 | 112 | def __getitem__(self, item): 113 | pointcloud = self.data[item][:self.num_points] 114 | label = self.label[item] 115 | 116 | if self.partition == 'train': 117 | pointcloud = translate_pointcloud(pointcloud) 118 | np.random.shuffle(pointcloud) 119 | return pointcloud, label 120 | 121 | def __len__(self): 122 | return self.data.shape[0] 123 | 124 | class ModelNet40Noise(Dataset): 125 | def __init__(self, num_points, num_noise, partition='test'): 126 | assert partition == "test",'Noise study can only be applied during evaluation' 127 | self.data, self.label = load_data_mdn(partition) 128 | self.num_points = num_points 129 | self.partition = partition 130 | self.num_noise = num_noise 131 | assert self.num_noise <= self.num_points,'number of noise points should be less than the point number' 132 | 133 | def __getitem__(self, item): 134 | pointcloud = self.data[item][:self.num_points] 135 | label = self.label[item] 136 | 137 | pointcloud = pointcloud[:-self.num_noise, :] 138 | noise = np.random.rand(self.num_noise, 3)*2-1 139 | pointcloud = np.concatenate((pointcloud, noise), axis=0).astype('float32') 140 | np.random.shuffle(pointcloud) 141 | 142 | return pointcloud, label 143 | 144 | def __len__(self): 145 | return self.data.shape[0] 146 | 147 | class ModelNet40Resplit(Dataset): 148 | def __init__(self, num_points, partition='train'): 149 | train_data, train_label = load_data_mdn('train') 150 | test_data, test_label = load_data_mdn('test') 151 | self.num_points = num_points 152 | self.partition = partition 153 | all_data = np.concatenate((train_data, test_data), axis=0) 154 | all_label = np.concatenate((train_label, test_label), axis=0) 155 | indices = list(range(all_data.shape[0])) 156 | np.random.shuffle(indices) 157 | if partition == 'train': 158 | self.data = all_data[indices[:8617], ...] 159 | self.label = all_label[indices[:8617], ...] 160 | elif partition == 'test': 161 | self.data = all_data[indices[8617:8617+1847], ...] 162 | self.label = all_label[indices[8617:8617+1847], ...] 163 | elif partition == 'vali': 164 | self.data = all_data[indices[8617+1847:], ...] 165 | self.label = all_label[indices[8617+1847:], ...] 166 | else: 167 | raise NameError 168 | 169 | def __getitem__(self, item): 170 | pointcloud = self.data[item][:self.num_points] 171 | label = self.label[item] 172 | 173 | if self.partition == 'train': 174 | pointcloud = translate_pointcloud(pointcloud) 175 | np.random.shuffle(pointcloud) 176 | return pointcloud, label 177 | 178 | def __len__(self): 179 | return self.data.shape[0] 180 | 181 | class ScanObjectNN(Dataset): 182 | def __init__(self, num_points, partition='train', bg=False): 183 | self.data, self.label = load_data_sonn(partition, bg) 184 | self.num_points = num_points 185 | self.partition = partition 186 | 187 | def __getitem__(self, item): 188 | pointcloud = self.data[item][:self.num_points] 189 | label = self.label[item] 190 | 191 | if self.partition == 'train': 192 | pointcloud = translate_pointcloud(pointcloud) 193 | pointcloud = rotate_pointcloud(pointcloud) 194 | pointcloud = jitter_pointcloud(pointcloud) 195 | np.random.shuffle(pointcloud) 196 | return pointcloud, label 197 | 198 | def __len__(self): 199 | return self.data.shape[0] 200 | 201 | class ModelNet40C(Dataset): 202 | def __init__(self, corruption, severity): 203 | base_dir = os.path.dirname(os.path.abspath(__file__)) 204 | data_dir = os.path.join(base_dir, 'data', 'modelnet40_c') 205 | DATA_DIR = os.path.join(data_dir, 'data_' + corruption + '_' +str(severity) + '.npy') 206 | LABEL_DIR = os.path.join(data_dir, 'label.npy') 207 | self.data = np.load(DATA_DIR) 208 | self.label = np.load(LABEL_DIR) 209 | 210 | def __getitem__(self, item): 211 | pointcloud = self.data[item] 212 | label = self.label[item] 213 | label = label 214 | return torch.from_numpy(pointcloud), torch.from_numpy(label).long() 215 | 216 | def __len__(self): 217 | return self.data.shape[0] 218 | 219 | class Toronto3D(Dataset): 220 | def __init__(self, num_points=2048, partition='train'): 221 | self.data, self.seg = load_data_seg(partition) 222 | self.num_points = num_points 223 | self.partition = partition 224 | 225 | def __getitem__(self, item): 226 | pointcloud = np.load(self.data[item]) 227 | seg = np.load(self.seg[item]) 228 | pointcloud = pointcloud[:self.num_points] 229 | seg = seg[:self.num_points] 230 | # normalize point cloud 231 | pointcloud = pointcloud - np.min(pointcloud, axis=0) 232 | pointcloud /= 5 # normalize point cloud in 5x5 blocks 233 | 234 | if self.partition == 'train': 235 | indices = list(range(pointcloud.shape[0])) 236 | pointcloud = translate_pointcloud(pointcloud) 237 | pointcloud = jitter_pointcloud(pointcloud) 238 | pointcloud[:,:3] = rotate_pointcloud(pointcloud[:,:3]) 239 | np.random.shuffle(indices) 240 | pointcloud = pointcloud[indices] 241 | seg = seg[indices] 242 | seg = torch.LongTensor(seg) 243 | pointcloud = torch.from_numpy(pointcloud) 244 | return pointcloud, seg 245 | 246 | def __len__(self): 247 | return len(self.data) 248 | 249 | class ShapeNetPart(Dataset): 250 | def __init__(self, num_points, partition='train'): 251 | self.data, self.label, self.seg = load_data_partseg(partition) 252 | self.seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 253 | self.index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 254 | self.num_points = num_points 255 | self.partition = partition 256 | self.seg_num_all = 50 257 | self.seg_start_index = 0 258 | 259 | def __getitem__(self, item): 260 | pointcloud = self.data[item][:self.num_points] 261 | label = self.label[item] 262 | seg = self.seg[item][:self.num_points] 263 | 264 | if self.partition == 'trainval': 265 | pointcloud = translate_pointcloud(pointcloud) 266 | indices = list(range(pointcloud.shape[0])) 267 | np.random.shuffle(indices) 268 | pointcloud = pointcloud[indices] 269 | seg = seg[indices] 270 | 271 | return pointcloud, label, seg 272 | 273 | def __len__(self): 274 | return self.data.shape[0] 275 | 276 | 277 | class ShapeNetPartNoise(Dataset): 278 | def __init__(self, num_points, num_noise, partition='test'): 279 | assert partition == "test",'Noise study can only be applied during evaluation' 280 | self.data, self.label, self.seg = load_data_partseg(partition) 281 | self.seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 282 | self.index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 283 | self.seg_num_all = 50 284 | self.seg_start_index = 0 285 | self.num_points = num_points 286 | self.num_noise = num_noise 287 | assert self.num_noise <= self.num_points,'number of noise points should be less than the point number' 288 | 289 | def __getitem__(self, item): 290 | pointcloud = self.data[item][:self.num_points] 291 | label = self.label[item] 292 | seg = self.seg[item][:self.num_points] 293 | 294 | noise = np.random.rand(self.num_noise, 3)*2-1 295 | 296 | dist = get_dist(torch.from_numpy(noise.astype('float32')).cuda().unsqueeze(0), 297 | torch.from_numpy(pointcloud.astype('float32')).cuda().unsqueeze(0)) 298 | 299 | noise_idx = torch.min(dist, dim=-1)[1].squeeze(0).cpu().numpy() 300 | noise_seg = seg[noise_idx] 301 | 302 | pointcloud = pointcloud[:-self.num_noise, :] 303 | seg = seg[:-self.num_noise] 304 | 305 | pointcloud = np.concatenate((pointcloud, noise), axis=0).astype('float32') 306 | seg = np.concatenate((seg, noise_seg), axis=0) 307 | 308 | indices = list(range(pointcloud.shape[0])) 309 | np.random.shuffle(indices) 310 | 311 | pointcloud = pointcloud[indices] 312 | seg = seg[indices] 313 | 314 | return pointcloud, label, seg 315 | 316 | def __len__(self): 317 | return self.data.shape[0] 318 | -------------------------------------------------------------------------------- /docs/4972.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/docs/4972.pdf -------------------------------------------------------------------------------- /environment.yaml: -------------------------------------------------------------------------------- 1 | name: dgcnn 2 | channels: 3 | - pytorch 4 | - conda-forge 5 | - defaults 6 | dependencies: 7 | - _libgcc_mutex=0.1 8 | - _openmp_mutex=4.5 9 | - ca-certificates=2021.7.5 10 | - certifi=2021.5.30 11 | - ld_impl_linux-64=2.35.1 12 | - libffi=3.3 13 | - libgcc-ng=9.3.0 14 | - libgomp=9.3.0 15 | - libstdcxx-ng=9.3.0 16 | - ncurses=6.2 17 | - openssl=1.1.1l 18 | - python=3.7.11 19 | - readline=8.1 20 | - sqlite=3.36.0 21 | - tk=8.6.10 22 | - wheel=0.37.0 23 | - xz=5.2.5 24 | - zlib=1.2.11 25 | - pip: 26 | - absl-py==1.0.0 27 | - addict==2.4.0 28 | - albumentations==1.2.1 29 | - anyio==3.6.1 30 | - argon2-cffi==21.3.0 31 | - argon2-cffi-bindings==21.2.0 32 | - attrs==21.4.0 33 | - babel==2.10.1 34 | - backcall==0.2.0 35 | - bboxtools==1.0.4 36 | - beautifulsoup4==4.11.1 37 | - bleach==5.0.0 38 | - cachecontrol==0.12.11 39 | - cached-property==1.5.2 40 | - cachetools==5.0.0 41 | - cachy==0.3.0 42 | - cffi==1.15.0 43 | - charset-normalizer==2.0.12 44 | - cleo==0.8.1 45 | - click==8.1.3 46 | - clikit==0.6.2 47 | - crashtest==0.3.1 48 | - cryptography==37.0.2 49 | - cycler==0.11.0 50 | - debugpy==1.6.0 51 | - decorator==5.1.1 52 | - defusedxml==0.7.1 53 | - deprecation==2.1.0 54 | - distlib==0.3.4 55 | - easydict==1.9 56 | - efficientnet-pytorch==0.7.1 57 | - entrypoints==0.4 58 | - exif==1.3.5 59 | - fastjsonschema==2.15.3 60 | - filelock==3.7.0 61 | - fonttools==4.29.1 62 | - frechetdist==0.6 63 | - google-auth==2.6.4 64 | - google-auth-oauthlib==0.4.6 65 | - grpcio==1.44.0 66 | - gudhi==3.5.0 67 | - h5py==3.3.0 68 | - html5lib==1.1 69 | - idna==3.3 70 | - imagecorruptions==1.1.2 71 | - imageio==2.19.1 72 | - imgaug==0.4.0 73 | - importlib-metadata==4.11.4 74 | - importlib-resources==5.7.1 75 | - ipykernel==6.13.0 76 | - ipython==7.33.0 77 | - ipython-genutils==0.2.0 78 | - ipywidgets==7.7.0 79 | - jedi==0.18.1 80 | - jeepney==0.8.0 81 | - jinja2==3.1.2 82 | - joblib==1.0.1 83 | - json5==0.9.8 84 | - jsonschema==4.5.1 85 | - jupyter-client==7.3.1 86 | - jupyter-core==4.10.0 87 | - jupyter-packaging==0.12.0 88 | - jupyter-server==1.17.0 89 | - jupyterlab==3.4.2 90 | - jupyterlab-pygments==0.2.2 91 | - jupyterlab-server==2.14.0 92 | - jupyterlab-widgets==1.1.0 93 | - keyring==22.3.0 94 | - kiwisolver==1.3.2 95 | - libtiff==0.4.2 96 | - lockfile==0.12.2 97 | - markdown==3.3.6 98 | - markupsafe==2.1.1 99 | - matplotlib==3.5.1 100 | - matplotlib-inline==0.1.3 101 | - mistune==0.8.4 102 | - mogutda==0.3.3 103 | - monai==1.0.1 104 | - msgpack==1.0.3 105 | - munch==2.5.0 106 | - nbclassic==0.3.7 107 | - nbclient==0.6.3 108 | - nbconvert==6.5.0 109 | - nbformat==5.4.0 110 | - nest-asyncio==1.5.5 111 | - networkx==2.6.3 112 | - ninja==1.10.2.2 113 | - notebook==6.4.11 114 | - notebook-shim==0.1.0 115 | - numpy==1.21.1 116 | - nvidia-cublas-cu11==11.10.3.66 117 | - nvidia-cuda-nvrtc-cu11==11.7.99 118 | - nvidia-cuda-runtime-cu11==11.7.99 119 | - nvidia-cudnn-cu11==8.5.0.96 120 | - nvidia-htop==1.0.5 121 | - oauthlib==3.2.0 122 | - open3d==0.15.2 123 | - opencv-contrib-python==4.6.0.66 124 | - opencv-python==4.6.0.66 125 | - opencv-python-headless==4.6.0.66 126 | - packaging==20.9 127 | - pandas==1.3.5 128 | - pandocfilters==1.5.0 129 | - parso==0.8.3 130 | - pastel==0.2.1 131 | - pexpect==4.8.0 132 | - pickleshare==0.7.5 133 | - pillow==9.1.0 134 | - pip==22.2.2 135 | - pkginfo==1.8.2 136 | - platformdirs==2.5.2 137 | - plum-py==0.8.2 138 | - plyfile==0.7.4 139 | - poetry==1.1.13 140 | - poetry-core==1.0.8 141 | - pointnet2==3.0.0 142 | - pointnet2-ops==3.0.0 143 | - pretrainedmodels==0.7.4 144 | - prometheus-client==0.14.1 145 | - prompt-toolkit==3.0.29 146 | - protobuf==3.18.0 147 | - psutil==5.9.1 148 | - ptflops==0.6.8 149 | - ptyprocess==0.7.0 150 | - pyasn1==0.4.8 151 | - pyasn1-modules==0.2.8 152 | - pycparser==2.21 153 | - pyexiv2==2.8.0 154 | - pygments==2.12.0 155 | - pylev==1.4.0 156 | - pyparsing==3.0.7 157 | - pyquaternion==0.9.9 158 | - pyrsistent==0.18.1 159 | - python-dateutil==2.8.2 160 | - pytz==2021.3 161 | - pywavelets==1.3.0 162 | - pyyaml==6.0 163 | - pyzmq==23.0.0 164 | - qudida==0.0.4 165 | - requests==2.27.1 166 | - requests-oauthlib==1.3.1 167 | - requests-toolbelt==0.9.1 168 | - rsa==4.8 169 | - scikit-image==0.19.2 170 | - scikit-learn==0.24.2 171 | - scipy==1.7.1 172 | - secretstorage==3.3.2 173 | - segmentation-models-pytorch==0.3.1 174 | - send2trash==1.8.0 175 | - setuptools==62.3.2 176 | - shapely==1.8.4 177 | - shellingham==1.4.0 178 | - similaritymeasures==0.4.4 179 | - six==1.16.0 180 | - sklearn==0.0 181 | - sniffio==1.2.0 182 | - soupsieve==2.3.2.post1 183 | - tensorboard==2.8.0 184 | - tensorboard-data-server==0.6.1 185 | - tensorboard-plugin-wit==1.8.1 186 | - tensorboardx==2.4 187 | - termcolor==1.1.0 188 | - terminado==0.15.0 189 | - thop==0.0.31-2005241907 190 | - threadpoolctl==2.2.0 191 | - tifffile==2021.11.2 192 | - timm==0.4.12 193 | - tinycss2==1.1.1 194 | - tomlkit==0.11.0 195 | - topologylayer==0.0.0 196 | - torch==1.8.1 197 | - torch-tb-profiler==0.4.0 198 | - torchaudio==0.9.0 199 | - torchsummary==1.5.1 200 | - torchvision==0.10.0 201 | - tornado==6.1 202 | - tqdm==4.64.0 203 | - traitlets==5.2.1.post0 204 | - typing-extensions==3.10.0.0 205 | - urllib3==1.26.9 206 | - virtualenv==20.14.1 207 | - wcwidth==0.2.5 208 | - webencodings==0.5.1 209 | - websocket-client==1.3.2 210 | - werkzeug==2.1.1 211 | - widgetsnbextension==3.6.0 212 | - zipp==3.8.0 213 | prefix: /home/manli/.conda/envs/dgcnn 214 | -------------------------------------------------------------------------------- /eval_modelnet40C.sh: -------------------------------------------------------------------------------- 1 | for cor in 'uniform' 'gaussian' 'background' 'impulse' 'upsampling' 'distortion_rbf' 'distortion_rbf_inv' 'density' 'density_inc' 'shear' 'rotation' 'cutout' 'distortion' 'occlusion' 'lidar'; do 2 | 3 | for sev in 1 2 3 4 5; do 4 | 5 | CUDA_VISIBLE_DEVICES=0 python3 main_cls.py --eval=True --model_path=checkpoints/model_cls.pth --dataset=modelnet40C --exp_name=eval_modelnet40C --corruption=${cor} --severity=${sev} 6 | 7 | done 8 | done 9 | 10 | python3 compute_cer.py -------------------------------------------------------------------------------- /eval_modelnet40noise.sh: -------------------------------------------------------------------------------- 1 | for nn in 1 10 50 100; do 2 | 3 | CUDA_VISIBLE_DEVICES=0 python3 main_cls.py --eval=True --model_path=checkpoints/model_cls.pth --dataset=modelnet40noise --exp_name=eval_modelnet40noise --num_noise=${nn} 4 | 5 | done 6 | -------------------------------------------------------------------------------- /logs/PLACEHOLDER: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mmmmimic/diffConvNet/8e07b04a67ab039601d427692602d7d953c0b4a9/logs/PLACEHOLDER -------------------------------------------------------------------------------- /main_cls.py: -------------------------------------------------------------------------------- 1 | """ 2 | @Author: Yue Wang 3 | @Contact: yuewangx@mit.edu 4 | @File: main_cls.py 5 | @Time: 2018/10/13 10:39 PM 6 | 7 | Modified by 8 | @Author: Manxi Lin 9 | @Contact: manli@dtu.dk 10 | @Time: 2022/07/07 17:10 PM 11 | """ 12 | 13 | from __future__ import print_function 14 | import os 15 | import argparse 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | import torch.optim as optim 20 | from torch.optim.lr_scheduler import CosineAnnealingLR 21 | from data_utils import ModelNet40, ModelNet40C, ModelNet40Noise, ModelNet40Resplit, ScanObjectNN 22 | from models.diffConv_cls import Model 23 | import numpy as np 24 | from torch.utils.data import DataLoader 25 | from misc import cal_loss, IOStream 26 | import sklearn.metrics as metrics 27 | 28 | torch.cuda.synchronize() 29 | 30 | def train(args, io): 31 | if args.dataset == 'modelnet40': 32 | train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=32, 33 | batch_size=args.batch_size, shuffle=True, drop_last=True) 34 | test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=32, 35 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 36 | output_channels = 40 37 | elif args.dataset == 'modelnet40resplit': 38 | train_loader = DataLoader(ModelNet40Resplit(partition='train', num_points=args.num_points), num_workers=32, 39 | batch_size=args.batch_size, shuffle=True, drop_last=True) 40 | test_loader = DataLoader(ModelNet40Resplit(partition='vali', num_points=args.num_points), num_workers=32, 41 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 42 | output_channels = 40 43 | elif args.dataset == 'scanobjectnn': 44 | train_loader = DataLoader(ScanObjectNN(partition='train', num_points=args.num_points, bg=args.bg), num_workers=32, 45 | batch_size=args.batch_size, shuffle=True, drop_last=True) 46 | test_loader = DataLoader(ScanObjectNN(partition='test', num_points=args.num_points, bg=args.bg), num_workers=32, 47 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 48 | output_channels = 15 49 | 50 | device = "cuda" if torch.cuda.is_available() else "cpu" 51 | 52 | model = Model(args, output_channels) 53 | 54 | print('begin experiment: %s'%args.exp_name) 55 | 56 | model = model.to(device) 57 | print(str(model)) 58 | 59 | model = nn.DataParallel(model) 60 | print("Let's use", torch.cuda.device_count(), "GPUs!") 61 | 62 | opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) 63 | 64 | scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3) 65 | 66 | criterion = cal_loss 67 | 68 | best_test_acc = 0 69 | for epoch in range(args.epochs): 70 | #################### 71 | # Train 72 | #################### 73 | train_loss = 0.0 74 | count = 0.0 75 | model.train() 76 | train_pred = [] 77 | train_true = [] 78 | for data, label in train_loader: 79 | data, label = data.to(device), label.to(device).squeeze() 80 | batch_size = data.size()[0] 81 | opt.zero_grad() 82 | logits = model(data) 83 | loss = criterion(logits, label) 84 | loss.backward() 85 | opt.step() 86 | preds = logits.max(dim=1)[1] 87 | count += batch_size 88 | train_loss += loss.item() * batch_size 89 | train_true.append(label.cpu().numpy()) 90 | train_pred.append(preds.detach().cpu().numpy()) 91 | scheduler.step() 92 | 93 | train_true = np.concatenate(train_true) 94 | train_pred = np.concatenate(train_pred) 95 | outstr = 'Train %d, loss: %.4f, train acc: %.4f, train avg acc: %.4f' % (epoch, 96 | train_loss*1.0/count, 97 | metrics.accuracy_score( 98 | train_true, train_pred), 99 | metrics.balanced_accuracy_score( 100 | train_true, train_pred)) 101 | io.cprint(outstr) 102 | 103 | #################### 104 | # Validation 105 | #################### 106 | test_loss = 0.0 107 | count = 0.0 108 | model.eval() 109 | test_pred = [] 110 | test_true = [] 111 | for data, label in test_loader: 112 | data, label = data.to(device), label.to(device).squeeze() 113 | batch_size = data.size()[0] 114 | with torch.no_grad(): 115 | logits = model(data) 116 | loss = criterion(logits, label) 117 | preds = logits.max(dim=1)[1] 118 | count += batch_size 119 | test_loss += loss.item() * batch_size 120 | test_true.append(label.cpu().numpy()) 121 | test_pred.append(preds.detach().cpu().numpy()) 122 | test_true = np.concatenate(test_true) 123 | test_pred = np.concatenate(test_pred) 124 | test_acc = metrics.accuracy_score(test_true, test_pred) 125 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) 126 | outstr = 'Test %d, loss: %.4f, test acc: %.4f, test avg acc: %.4f' % (epoch, 127 | test_loss*1.0/count, 128 | test_acc, 129 | avg_per_class_acc) 130 | io.cprint(outstr) 131 | if test_acc >= best_test_acc: 132 | best_test_acc = test_acc 133 | torch.save(model.state_dict(), './checkpoints/%s.pth' % args.exp_name) 134 | 135 | def test(args, io): 136 | if args.dataset == 'modelnet40': 137 | test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=32, 138 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 139 | output_channels = 40 140 | elif args.dataset == 'modelnet40C': 141 | test_loader = DataLoader(ModelNet40C(args.corruption, args.severity), num_workers=32, 142 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 143 | output_channels = 40 144 | elif args.dataset == 'modelnet40noise': 145 | test_loader = DataLoader(ModelNet40Noise(args.num_points, args.num_noise), num_workers=32, 146 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 147 | output_channels = 40 148 | elif args.dataset == 'modelnet40resplit': 149 | test_loader = DataLoader(ModelNet40Resplit(partition='test', num_points=args.num_points), num_workers=32, 150 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 151 | output_channels = 40 152 | elif args.dataset == 'scanobjectnn': 153 | test_loader = DataLoader(ScanObjectNN(partition='test', num_points=args.num_points, bg=args.bg), num_workers=32, 154 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 155 | output_channels = 15 156 | 157 | device = "cuda" if torch.cuda.is_available() else "cpu" 158 | 159 | #Try to load models 160 | model = Model(args, output_channels) 161 | 162 | model = model.to(device) 163 | 164 | model = nn.DataParallel(model) 165 | model.load_state_dict(torch.load(args.model_path)) 166 | model = model.eval() 167 | test_acc = 0.0 168 | test_true = [] 169 | test_pred = [] 170 | for data, label in test_loader: 171 | data, label = data.to(device), label.to(device).squeeze() 172 | with torch.no_grad(): 173 | logits = model(data) 174 | preds = logits.max(dim=1)[1] 175 | test_true.append(label.cpu().numpy()) 176 | test_pred.append(preds.detach().cpu().numpy()) 177 | test_true = np.concatenate(test_true) 178 | test_pred = np.concatenate(test_pred) 179 | test_acc = metrics.accuracy_score(test_true, test_pred) 180 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) 181 | outstr = 'Test :: test acc: %.3f, test avg acc: %.3f'%(test_acc, avg_per_class_acc) 182 | io.cprint(outstr) 183 | 184 | if __name__ == "__main__": 185 | # Training settings 186 | parser = argparse.ArgumentParser(description='Point Cloud Recognition') 187 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 188 | help='Name of the experiment') 189 | parser.add_argument('--dataset', type=str, default='modelnet40', metavar='N', 190 | choices=['modelnet40', 'modelnet40C', 'modelnet40noise', 'modelnet40resplit', 'scanobjectnn']) 191 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 192 | help='Size of batch)') 193 | parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size', 194 | help='Size of batch)') 195 | parser.add_argument('--epochs', type=int, default=2000, metavar='N', 196 | help='number of episode to train') 197 | parser.add_argument('--lr', type=float, default=0.001, metavar='LR', 198 | help='learning rate (default: 0.001, 0.1 if using sgd)') 199 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 200 | help='SGD momentum (default: 0.9)') 201 | parser.add_argument('--seed', type=int, default=42, metavar='S', 202 | help='random seed (default: 42)') 203 | parser.add_argument('--eval', type=bool, default=False, 204 | help='evaluate the model') 205 | parser.add_argument('--num_points', type=int, default=1024, 206 | help='num of points to use') 207 | parser.add_argument('--dropout', type=float, default=0.5, 208 | help='initial dropout rate') 209 | parser.add_argument('--emb_dims', type=int, default=512, metavar='N', 210 | help='Dimension of embeddings') 211 | parser.add_argument('--model_path', type=str, default='', metavar='N', 212 | help='Pretrained model path') 213 | parser.add_argument('--radius', type=float, default=0.005, 214 | help='search radius') 215 | parser.add_argument('--corruption', type=str, default='uniform', metavar='N', 216 | help='corruption of ModelNetC') 217 | parser.add_argument('--severity', type=int, default=1, metavar='S', 218 | help='severity of ModelNetC') 219 | parser.add_argument('--num_noise', type=int, default=100, 220 | help='number of noise points in noise study') 221 | parser.add_argument('--bg', type=bool, default=False, 222 | help='whether to add background in scanobjectnn') 223 | 224 | args = parser.parse_args() 225 | 226 | io = IOStream(os.path.join('./logs', args.exp_name)) 227 | io.cprint(str(args)) 228 | 229 | torch.manual_seed(args.seed) 230 | if torch.cuda.is_available(): 231 | io.cprint( 232 | 'Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') 233 | torch.cuda.manual_seed(args.seed) 234 | else: 235 | io.cprint('Using CPU') 236 | 237 | if not args.eval: 238 | train(args, io) 239 | else: 240 | test(args, io) 241 | -------------------------------------------------------------------------------- /main_partseg.py: -------------------------------------------------------------------------------- 1 | """ 2 | @Author: An Tao, Pengliang Ji 3 | @Contact: ta19@mails.tsinghua.edu.cn, jpl1723@buaa.edu.cn 4 | @File: main_partseg.py 5 | @Time: 2021/7/20 7:49 PM 6 | 7 | Modified by 8 | @Author: Manxi Lin 9 | @Contact: manli@dtu.dk 10 | @Time: 2022/07/11 15:29 PM 11 | """ 12 | 13 | from __future__ import print_function 14 | import os 15 | import argparse 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | import torch.optim as optim 20 | from torch.optim.lr_scheduler import CosineAnnealingLR 21 | from data_utils import ShapeNetPart, ShapeNetPartNoise 22 | from models.diffConv_partseg import Model 23 | import numpy as np 24 | from torch.utils.data import DataLoader 25 | from misc import cal_loss, IOStream 26 | import sklearn.metrics as metrics 27 | 28 | seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 29 | index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 30 | 31 | def calculate_shape_IoU(pred_np, seg_np, label): 32 | shape_ious = [] 33 | for shape_idx in range(seg_np.shape[0]): 34 | parts = range(seg_num[label[0]]) 35 | part_ious = [] 36 | for part in parts: 37 | I = np.sum(np.logical_and(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 38 | U = np.sum(np.logical_or(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 39 | if U == 0: 40 | iou = 1 41 | else: 42 | iou = I / float(U) 43 | part_ious.append(iou) 44 | shape_ious.append(np.mean(part_ious)) 45 | return shape_ious 46 | 47 | def train(args, io): 48 | if args.dataset == "shapenetpart": 49 | train_dataset = ShapeNetPart(partition='trainval', num_points=args.num_points) 50 | test_dataset = ShapeNetPart(partition='test', num_points=args.num_points) 51 | train_loader = DataLoader(train_dataset, num_workers=32, batch_size=args.batch_size, shuffle=True, drop_last=False) 52 | test_loader = DataLoader(test_dataset, 53 | num_workers=32, batch_size=args.test_batch_size, shuffle=True, drop_last=False) 54 | 55 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 56 | 57 | #Try to load models 58 | seg_num_all = train_loader.dataset.seg_num_all 59 | seg_start_index = train_loader.dataset.seg_start_index 60 | 61 | model = Model(args).to(device) 62 | print(str(model)) 63 | 64 | model = nn.DataParallel(model) 65 | print("Let's use ", torch.cuda.device_count(), " GPUs!") 66 | 67 | opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) 68 | 69 | scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3) 70 | criterion = cal_loss 71 | 72 | best_test_iou = 0 73 | for epoch in range(args.epochs): 74 | #################### 75 | # Train 76 | #################### 77 | train_loss = 0.0 78 | count = 0.0 79 | model.train() 80 | train_true_cls = [] 81 | train_pred_cls = [] 82 | train_true_seg = [] 83 | train_pred_seg = [] 84 | train_label_seg = [] 85 | 86 | for data, label, seg in train_loader: 87 | seg = seg - seg_start_index 88 | label_one_hot = np.zeros((label.shape[0], 16)) 89 | for idx in range(label.shape[0]): 90 | label_one_hot[idx, label[idx]] = 1 91 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 92 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 93 | batch_size = data.size()[0] 94 | opt.zero_grad() 95 | seg_pred = model(data, label_one_hot) 96 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 97 | loss = criterion(seg_pred.view(-1, seg_num_all), seg.view(-1,1).squeeze()) 98 | 99 | loss.backward() 100 | torch.nn.utils.clip_grad_norm_(model.parameters(), 2) 101 | opt.step() 102 | pred = seg_pred.max(dim=2)[1] 103 | count += batch_size 104 | train_loss += loss.item() * batch_size 105 | seg_np = seg.cpu().numpy() 106 | pred_np = pred.detach().cpu().numpy() 107 | train_true_cls.append(seg_np.reshape(-1)) 108 | train_pred_cls.append(pred_np.reshape(-1)) 109 | train_true_seg.append(seg_np) 110 | train_pred_seg.append(pred_np) 111 | train_label_seg.append(label.reshape(-1)) 112 | scheduler.step() 113 | train_true_cls = np.concatenate(train_true_cls) 114 | train_pred_cls = np.concatenate(train_pred_cls) 115 | train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls) 116 | avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls, train_pred_cls) 117 | train_true_seg = np.concatenate(train_true_seg, axis=0) 118 | train_pred_seg = np.concatenate(train_pred_seg, axis=0) 119 | train_label_seg = np.concatenate(train_label_seg) 120 | train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg) 121 | outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % (epoch, 122 | train_loss*1.0/count, 123 | train_acc, 124 | avg_per_class_acc, 125 | np.mean(train_ious)) 126 | io.cprint(outstr) 127 | 128 | #################### 129 | # Test 130 | #################### 131 | test_loss = 0.0 132 | count = 0.0 133 | model.eval() 134 | test_true_cls = [] 135 | test_pred_cls = [] 136 | test_true_seg = [] 137 | test_pred_seg = [] 138 | test_label_seg = [] 139 | for data, label, seg in test_loader: 140 | seg = seg - seg_start_index 141 | label_one_hot = np.zeros((label.shape[0], 16)) 142 | for idx in range(label.shape[0]): 143 | label_one_hot[idx, label[idx]] = 1 144 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 145 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 146 | batch_size = data.size()[0] 147 | seg_pred = model(data, label_one_hot) 148 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 149 | loss = criterion(seg_pred.view(-1, seg_num_all), seg.view(-1,1).squeeze()) 150 | 151 | pred = seg_pred.max(dim=2)[1] 152 | count += batch_size 153 | test_loss += loss.item() * batch_size 154 | seg_np = seg.cpu().numpy() 155 | pred_np = pred.detach().cpu().numpy() 156 | test_true_cls.append(seg_np.reshape(-1)) 157 | test_pred_cls.append(pred_np.reshape(-1)) 158 | test_true_seg.append(seg_np) 159 | test_pred_seg.append(pred_np) 160 | test_label_seg.append(label.reshape(-1)) 161 | test_true_cls = np.concatenate(test_true_cls) 162 | test_pred_cls = np.concatenate(test_pred_cls) 163 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 164 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 165 | test_true_seg = np.concatenate(test_true_seg, axis=0) 166 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 167 | test_label_seg = np.concatenate(test_label_seg) 168 | test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg) 169 | outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (epoch, 170 | test_loss*1.0/count, 171 | test_acc, 172 | avg_per_class_acc, 173 | np.mean(test_ious)) 174 | io.cprint(outstr) 175 | 176 | if np.mean(test_ious) >= best_test_iou: 177 | best_test_iou = np.mean(test_ious) 178 | torch.save(model.state_dict(), 'checkpoints/%s.pth' % args.exp_name) 179 | 180 | def test(args, io): 181 | if args.dataset == "shapenetpart": 182 | test_dataset = ShapeNetPart(partition='test', num_points=args.num_points) 183 | elif args.dataset == "shapenetpartnoise": 184 | test_dataset = ShapeNetPartNoise(partition='test', num_points=args.num_points, num_noise=args.num_noise) 185 | 186 | test_loader = DataLoader(test_dataset, 187 | batch_size=args.test_batch_size, shuffle=False, drop_last=False) 188 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 189 | 190 | #Try to load models 191 | seg_num_all = test_loader.dataset.seg_num_all 192 | seg_start_index = test_loader.dataset.seg_start_index 193 | 194 | model = Model(args, seg_num_all).to(device) 195 | 196 | model = nn.DataParallel(model) 197 | model.load_state_dict(torch.load(args.model_path)) 198 | model = model.eval() 199 | test_acc = 0.0 200 | test_true_cls = [] 201 | test_pred_cls = [] 202 | test_true_seg = [] 203 | test_pred_seg = [] 204 | test_label_seg = [] 205 | for data, label, seg in test_loader: 206 | seg = seg - seg_start_index 207 | label_one_hot = np.zeros((label.shape[0], 16)) 208 | for idx in range(label.shape[0]): 209 | label_one_hot[idx, label[idx]] = 1 210 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 211 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 212 | seg_pred = model(data, label_one_hot) 213 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 214 | pred = seg_pred.max(dim=2)[1] 215 | seg_np = seg.cpu().numpy() 216 | pred_np = pred.detach().cpu().numpy() 217 | test_true_cls.append(seg_np.reshape(-1)) 218 | test_pred_cls.append(pred_np.reshape(-1)) 219 | test_true_seg.append(seg_np) 220 | test_pred_seg.append(pred_np) 221 | test_label_seg.append(label.reshape(-1)) 222 | test_true_cls = np.concatenate(test_true_cls) 223 | test_pred_cls = np.concatenate(test_pred_cls) 224 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 225 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 226 | test_true_seg = np.concatenate(test_true_seg, axis=0) 227 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 228 | test_label_seg = np.concatenate(test_label_seg) 229 | test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg) 230 | outstr = 'Test :: test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (test_acc, 231 | avg_per_class_acc, 232 | np.mean(test_ious)) 233 | io.cprint(outstr) 234 | 235 | if __name__ == "__main__": 236 | # Training settings 237 | parser = argparse.ArgumentParser(description='Point Cloud Part Segmentation') 238 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 239 | help='Name of the experiment') 240 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 241 | help='Size of batch)') 242 | parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size', 243 | help='Size of batch)') 244 | parser.add_argument('--epochs', type=int, default=2000, metavar='N', 245 | help='number of episode to train ') 246 | parser.add_argument('--lr', type=float, default=0.001, metavar='LR', 247 | help='learning rate (default: 0.001, 0.1 if using sgd)') 248 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 249 | help='SGD momentum (default: 0.9)') 250 | parser.add_argument('--seed', type=int, default=42, metavar='S', 251 | help='random seed (default: 42)') 252 | parser.add_argument('--eval', type=bool, default=False, 253 | help='evaluate the model') 254 | parser.add_argument('--num_points', type=int, default=2048, 255 | help='num of points to use') 256 | parser.add_argument('--dropout', type=float, default=0.5, 257 | help='dropout rate') 258 | parser.add_argument('--radius', type=float, default=0.005, 259 | help='searching radius') 260 | parser.add_argument('--model_path', type=str, default='', metavar='N', 261 | help='Pretrained model path') 262 | parser.add_argument('--dataset', type=str, default='shapenetpart', metavar='N', 263 | choices=['shapenetpart', 'shapenetpartnoise']) 264 | parser.add_argument('--num_noise', type=int, default=100, 265 | help='number of noise points in noise study') 266 | 267 | args = parser.parse_args() 268 | 269 | io = IOStream(os.path.join('./logs', args.exp_name)) 270 | io.cprint(str(args)) 271 | 272 | torch.manual_seed(args.seed) 273 | if torch.cuda.is_available(): 274 | io.cprint( 275 | 'Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') 276 | torch.cuda.manual_seed(args.seed) 277 | else: 278 | io.cprint('Using CPU') 279 | 280 | if not args.eval: 281 | train(args, io) 282 | else: 283 | test(args, io) -------------------------------------------------------------------------------- /main_seg.py: -------------------------------------------------------------------------------- 1 | """ 2 | @Author: Yue Wang 3 | @Contact: yuewangx@mit.edu 4 | @File: main_seg.py 5 | @Time: 2018/10/13 10:39 PM 6 | 7 | Modified by 8 | @Author: Manxi Lin 9 | @Contact: manli@dtu.dk 10 | @Time: 2022/07/10 16:45 PM 11 | """ 12 | from __future__ import print_function 13 | import os 14 | import argparse 15 | import torch 16 | import torch.nn as nn 17 | import torch.nn.functional as F 18 | import torch.optim as optim 19 | from torch.optim.lr_scheduler import CosineAnnealingLR 20 | from data_utils import Toronto3D 21 | from models.diffConv_seg import Model 22 | import numpy as np 23 | from torch.utils.data import DataLoader 24 | from misc import cal_loss, IOStream 25 | import sklearn.metrics as metrics 26 | 27 | def calculate_sem_IoU(pred_np, seg_np): 28 | I_all = np.zeros(9) 29 | U_all = np.zeros(9) 30 | for sem_idx in range(seg_np.shape[0]): 31 | for sem in range(9): 32 | I = np.sum(np.logical_and(pred_np[sem_idx] == sem, seg_np[sem_idx] == sem)) 33 | U = np.sum(np.logical_or(pred_np[sem_idx] == sem, seg_np[sem_idx] == sem)) 34 | I_all[sem] += I 35 | U_all[sem] += U 36 | # 0 is unlabelled, thus is not considered 37 | I_all = I_all[1:] 38 | U_all = U_all[1:] 39 | return I_all / U_all 40 | 41 | def calculate_class_IoU(pred_np, seg_np): 42 | ''' 43 | return iou for each category 44 | ''' 45 | IOU = np.zeros((seg_np.shape[0], 9)) 46 | for sem_idx in range(seg_np.shape[0]): 47 | n = 9 48 | for sem in range(9): 49 | I = np.sum(np.logical_and(pred_np[sem_idx] == sem, seg_np[sem_idx] == sem)) 50 | U = np.sum(np.logical_or(pred_np[sem_idx] == sem, seg_np[sem_idx] == sem)) 51 | if U == 0: 52 | I = 1 53 | U = 1 54 | n -= 1 55 | IOU[sem_idx, sem] = I / U 56 | return IOU[:,1:] # remove the unlabelled class (0) 57 | 58 | def train(args, io): 59 | train_loader = DataLoader(Toronto3D(partition='train'), 60 | num_workers=16, batch_size=args.batch_size, shuffle=True, drop_last=True) 61 | test_loader = DataLoader(Toronto3D(partition='test'), 62 | num_workers=16, batch_size=args.test_batch_size, shuffle=False, drop_last=False) 63 | 64 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 65 | 66 | #Try to load models 67 | model = Model(args, num_classes=9) 68 | model.to(device) 69 | 70 | model = nn.DataParallel(model) 71 | print("Let's use", torch.cuda.device_count(), "GPUs!") 72 | 73 | opt = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) 74 | 75 | scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-4) 76 | criterion = cal_loss 77 | 78 | best_mean_iou = 0 79 | for epoch in range(args.epochs): 80 | #################### 81 | # Train 82 | #################### 83 | train_loss = 0.0 84 | count = 0.0 85 | model.train() 86 | train_true_cls = [] 87 | train_pred_cls = [] 88 | train_true_seg = [] 89 | train_pred_seg = [] 90 | for data, seg in train_loader: 91 | data, seg = data.to(device), seg.to(device) 92 | batch_size = data.size()[0] 93 | opt.zero_grad() 94 | seg_pred = model(data) 95 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 96 | loss = criterion(seg_pred.view(-1, 9), seg.view(-1,1).squeeze()) 97 | loss.backward() 98 | opt.step() 99 | 100 | pred = seg_pred.argmax(dim=2) 101 | count += batch_size 102 | train_loss += loss.item() * batch_size 103 | seg_np = seg.cpu().numpy() 104 | pred_np = pred.detach().cpu().numpy() 105 | train_true_cls.append(seg_np.reshape(-1)) 106 | train_pred_cls.append(pred_np.reshape(-1)) 107 | train_true_seg.append(seg_np) 108 | train_pred_seg.append(pred_np) 109 | scheduler.step() 110 | train_true_cls = np.concatenate(train_true_cls) 111 | train_pred_cls = np.concatenate(train_pred_cls) 112 | train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls) 113 | avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls, train_pred_cls) 114 | train_true_seg = np.concatenate(train_true_seg, axis=0) 115 | train_pred_seg = np.concatenate(train_pred_seg, axis=0) 116 | train_ious = calculate_sem_IoU(train_pred_seg, train_true_seg) 117 | ious = calculate_class_IoU(train_pred_seg, train_true_seg) 118 | outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, train iou: %.6f, mean iou: %.6f' % (epoch, 119 | train_loss*1.0/count, 120 | train_acc, 121 | avg_per_class_acc, 122 | np.mean(train_ious), 123 | np.mean(ious)) 124 | io.cprint(outstr) 125 | 126 | #################### 127 | # Test 128 | #################### 129 | test_loss = 0.0 130 | count = 0.0 131 | model.eval() 132 | test_true_cls = [] 133 | test_pred_cls = [] 134 | test_true_seg = [] 135 | test_pred_seg = [] 136 | for data, seg in test_loader: 137 | data, seg = data.to(device), seg.to(device) 138 | batch_size = data.size()[0] 139 | with torch.no_grad(): 140 | seg_pred = model(data) 141 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 142 | loss = criterion(seg_pred.view(-1, 9), seg.view(-1,1).squeeze()) 143 | pred = seg_pred.argmax(dim=2) 144 | count += batch_size 145 | test_loss += loss.item() * batch_size 146 | seg_np = seg.cpu().numpy() 147 | pred_np = pred.detach().cpu().numpy() 148 | test_true_cls.append(seg_np.reshape(-1)) 149 | test_pred_cls.append(pred_np.reshape(-1)) 150 | test_true_seg.append(seg_np) 151 | test_pred_seg.append(pred_np) 152 | test_true_cls = np.concatenate(test_true_cls) 153 | test_pred_cls = np.concatenate(test_pred_cls) 154 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 155 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 156 | test_true_seg = np.concatenate(test_true_seg, axis=0) 157 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 158 | test_ious = calculate_sem_IoU(test_pred_seg, test_true_seg) 159 | ious = calculate_class_IoU(test_pred_seg, test_true_seg) 160 | outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test iou: %.6f, mean iou: %.6f' % (epoch, 161 | test_loss*1.0/count, 162 | test_acc, 163 | avg_per_class_acc, 164 | np.mean(test_ious), 165 | np.mean(ious)) 166 | io.cprint(outstr) 167 | if np.mean(ious) >= best_mean_iou: 168 | best_mean_iou = np.mean(ious) 169 | torch.save(model.state_dict(), 'checkpoints/%s.pth' %args.exp_name) 170 | 171 | def test(args, io): 172 | test_loader = DataLoader(Toronto3D(partition='test'), 173 | num_workers=16, batch_size=args.test_batch_size, shuffle=False, drop_last=False) 174 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 175 | 176 | #Try to load models 177 | model = Model(args, num_classes=9) 178 | model.to(device) 179 | 180 | model = nn.DataParallel(model) 181 | print("Let's use", torch.cuda.device_count(), "GPUs!") 182 | 183 | print('load model') 184 | model.load_state_dict(torch.load(args.model_path)) 185 | 186 | with torch.no_grad(): 187 | #################### 188 | # Test 189 | #################### 190 | count = 0.0 191 | model.eval() 192 | test_true_cls = [] 193 | test_pred_cls = [] 194 | test_true_seg = [] 195 | test_pred_seg = [] 196 | for data, seg in test_loader: 197 | data, seg = data.to(device), seg.to(device) 198 | batch_size = data.size()[0] 199 | with torch.no_grad(): 200 | seg_pred = model(data) 201 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 202 | pred = seg_pred.max(dim=2)[1] 203 | count += batch_size 204 | seg_np = seg.cpu().numpy() 205 | pred_np = pred.detach().cpu().numpy() 206 | test_true_cls.append(seg_np.reshape(-1)) 207 | test_pred_cls.append(pred_np.reshape(-1)) 208 | test_true_seg.append(seg_np) 209 | test_pred_seg.append(pred_np) 210 | test_true_cls = np.concatenate(test_true_cls) 211 | test_pred_cls = np.concatenate(test_pred_cls) 212 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 213 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 214 | test_true_seg = np.concatenate(test_true_seg, axis=0) 215 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 216 | test_ious = calculate_sem_IoU(test_pred_seg, test_true_seg) 217 | ious = calculate_class_IoU(test_pred_seg, test_true_seg) 218 | ious = np.mean(ious, axis=0) 219 | category_names = ['Road', 'Road Mark', 'Natural', 'Building', 'Util. line', 'Pole', 'Car', 'Fence'] 220 | io.cprint('Category-wise iou:') 221 | for c, iou in zip(category_names, ious): 222 | io.cprint('%s: %.4f'%(c, iou)) 223 | outstr = 'Overall Test :: test acc: %.4f, test avg acc: %.4f, test iou: %.4f, mean iou: %.4f' % (test_acc, 224 | avg_per_class_acc, 225 | np.mean(test_ious), 226 | np.mean(ious)) 227 | io.cprint(outstr) 228 | 229 | 230 | if __name__ == "__main__": 231 | # Training settings 232 | parser = argparse.ArgumentParser(description='Point Cloud Scene Segmentation') 233 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 234 | help='Name of the experiment') 235 | parser.add_argument('--batch_size', type=int, default=16, metavar='batch_size', 236 | help='Size of batch)') 237 | parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size', 238 | help='Size of batch)') 239 | parser.add_argument('--epochs', type=int, default=2000, metavar='N', 240 | help='number of episode to train ') 241 | parser.add_argument('--lr', type=float, default=0.001, metavar='LR', 242 | help='learning rate (default: 0.001, 0.1 if using sgd)') 243 | parser.add_argument('--seed', type=int, default=42, metavar='S', 244 | help='random seed (default: 42)') 245 | parser.add_argument('--eval', type=bool, default=False, 246 | help='evaluate the model') 247 | parser.add_argument('--num_points', type=int, default=2048, 248 | help='num of points to use') 249 | parser.add_argument('--dropout', type=float, default=0.5, 250 | help='dropout rate') 251 | parser.add_argument('--radius', type=float, default=0.005, 252 | help='searching radius') 253 | parser.add_argument('--model_path', type=str, default='', metavar='N', 254 | help='Pretrained model path') 255 | args = parser.parse_args() 256 | 257 | io = IOStream(os.path.join('./logs', args.exp_name)) 258 | io.cprint(str(args)) 259 | 260 | torch.manual_seed(args.seed) 261 | if torch.cuda.is_available(): 262 | io.cprint( 263 | 'Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') 264 | torch.cuda.manual_seed(args.seed) 265 | else: 266 | io.cprint('Using CPU') 267 | 268 | if not args.eval: 269 | train(args, io) 270 | else: 271 | test(args, io) 272 | -------------------------------------------------------------------------------- /misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | import os 4 | 5 | def cal_loss(pred, gt, smoothing=True): 6 | ''' Calculate cross entropy loss, apply label smoothing if needed. ''' 7 | 8 | gt = gt.contiguous().view(-1) 9 | 10 | if smoothing: 11 | eps = 0.2 12 | n_class = pred.size(1) 13 | 14 | one_hot = torch.zeros_like(pred).scatter(1, gt.view(-1, 1), 1) 15 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) 16 | log_prb = F.log_softmax(pred, dim=1) 17 | 18 | loss = -(one_hot * log_prb).sum(dim=1).mean() 19 | else: 20 | loss = F.cross_entropy(pred, gt, reduction='mean') 21 | 22 | return loss 23 | 24 | class IOStream(): 25 | def __init__(self, file_name): 26 | self.f = os.path.join(file_name+'.log') 27 | 28 | def cprint(self, text): 29 | print(text) 30 | with open(self.f, 'a') as f: 31 | f.write(text+'\n') 32 | f.flush() -------------------------------------------------------------------------------- /models/diffConv_cls.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from .modules import Conv1x1, diffConv 6 | 7 | class Model(nn.Module): 8 | def __init__(self, args, output_channels=40): 9 | super().__init__() 10 | 11 | self.args = args 12 | 13 | init_feat = 32 14 | 15 | self.le = Conv1x1(3, init_feat) # local encoder 16 | 17 | radius = args.radius 18 | 19 | self.conv1 = diffConv(init_feat, init_feat*2, radius) 20 | 21 | self.conv2 = diffConv(init_feat*2, init_feat*4, radius*2) 22 | 23 | self.conv3 = diffConv(init_feat*4, init_feat*8, radius*4) 24 | 25 | self.conv4 = diffConv(init_feat*8, init_feat*16, radius*8) 26 | 27 | self.last_conv = Conv1x1(init_feat*16, args.emb_dims) 28 | 29 | self.linear = nn.Sequential( 30 | nn.Linear(args.emb_dims*2, 512, bias=False), 31 | nn.BatchNorm1d(512), 32 | nn.GELU(), 33 | nn.Dropout(p=args.dropout), 34 | nn.Linear(512, output_channels) 35 | ) 36 | 37 | def forward(self, x): 38 | xyz = x.clone() 39 | point_num = xyz.size(1) 40 | 41 | feat = self.le(x) 42 | 43 | l1_feat, l1_xyz = self.conv1(feat, xyz, point_num) 44 | 45 | l2_feat, l2_xyz = self.conv2(l1_feat, l1_xyz, point_num//2) 46 | 47 | l3_feat, l3_xyz = self.conv3(l2_feat, l2_xyz, point_num//4) 48 | 49 | l4_feat, _ = self.conv4(l3_feat, l3_xyz, point_num//8) 50 | 51 | x = self.last_conv(l4_feat) 52 | 53 | batch_size = x.size(0) 54 | x = x.transpose(1, 2) 55 | x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) 56 | x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1) 57 | x = torch.cat((x1, x2), dim=1) 58 | 59 | x = self.linear(x) 60 | return x 61 | 62 | if __name__ == "__main__": 63 | parser = argparse.ArgumentParser(description='Point Cloud Classification') 64 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 65 | help='Name of the experiment') 66 | parser.add_argument('--dropout', type=float, default=0.5, 67 | help='dropout rate') 68 | parser.add_argument('--radius', type=float, default=0.005, 69 | help='search radius') 70 | parser.add_argument('--emb_dims', type=int, default=512, metavar='N', 71 | help='Dimension of embeddings') 72 | args = parser.parse_args() 73 | 74 | device = "cuda" if torch.cuda.is_available() else "cpu" 75 | pc = torch.rand(5, 1024, 3).to(device) 76 | model = Model(args).to(device) 77 | feat = model(pc) 78 | print(feat.shape) 79 | 80 | 81 | 82 | -------------------------------------------------------------------------------- /models/diffConv_partseg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from .modules import Conv1x1, diffConv, PointFeaturePropagation 5 | from .pointnet2.pointnet2_modules import PointnetSAModule 6 | import argparse 7 | 8 | class Model(nn.Module): 9 | def __init__(self, args, num_classes=50, category=16): 10 | super().__init__() 11 | 12 | init_channel = 32 13 | self.le0 = Conv1x1(3, init_channel) 14 | 15 | self.le1 = PointnetSAModule( 16 | npoint=args.num_points, 17 | radius=0.05, 18 | nsample=20, 19 | mlp=[init_channel, init_channel, init_channel], 20 | use_xyz=True, 21 | bn=True 22 | ) 23 | 24 | radius = args.radius 25 | 26 | # encoder 27 | self.conv1 = diffConv(init_channel, init_channel*2, radius) 28 | 29 | self.conv2 = diffConv(init_channel*2, init_channel*4, radius*4) 30 | 31 | self.conv3 = diffConv(init_channel*4, init_channel*8, radius*8) 32 | 33 | self.conv4 = diffConv(init_channel*8, init_channel*16, radius*16) 34 | 35 | self.conv5 = diffConv(init_channel*16, init_channel*32, radius*32) 36 | 37 | # decoder 38 | self.fp4 = PointFeaturePropagation(in_channel1=init_channel*32, in_channel2=init_channel*16, out_channel=init_channel*8) 39 | self.up_conv5 = diffConv(init_channel*8, init_channel*16, radius*16) 40 | 41 | self.fp3 = PointFeaturePropagation(in_channel1=init_channel*16, in_channel2=init_channel*8, out_channel=init_channel*4) 42 | self.up_conv4 = diffConv(init_channel*4, init_channel*8, radius*8) 43 | 44 | self.fp2 = PointFeaturePropagation(in_channel1=init_channel*8, in_channel2=init_channel*4, out_channel=init_channel*2) 45 | self.up_conv3 = diffConv(init_channel*2, init_channel*4, radius*4) 46 | 47 | self.fp1 = PointFeaturePropagation(in_channel1=init_channel*4, in_channel2=init_channel*2, out_channel=init_channel*1) 48 | self.up_conv2 = diffConv(init_channel, init_channel*2, radius) 49 | 50 | self.global_conv1 = nn.Sequential(Conv1x1(init_channel*32+3, init_channel*8), 51 | Conv1x1(init_channel*8, init_channel*2) 52 | ) 53 | self.global_conv2 = nn.Sequential(Conv1x1(init_channel*16+3, init_channel*4), 54 | Conv1x1(init_channel*4, init_channel*2) 55 | ) 56 | 57 | self.up_conv1 =Conv1x1(init_channel*(2+2+2)+3+category, 256) 58 | 59 | self.last_conv = nn.Linear(256, num_classes, bias=False) 60 | 61 | self.se = nn.Sequential(nn.AdaptiveAvgPool1d(1), 62 | nn.Conv1d(256, 256//8, 1, bias=False), 63 | nn.BatchNorm1d(256//8), 64 | nn.GELU(), 65 | nn.Conv1d(256//8, 256, 1, bias=False), 66 | nn.Sigmoid()) 67 | 68 | self.drop = nn.Dropout(args.dropout) 69 | 70 | def forward(self, x, l=None): 71 | batch_size = x.size(0) 72 | xyz = x.clone() 73 | point_num = xyz.size(1) 74 | 75 | x = self.le0(x) 76 | 77 | x = x.transpose(1,2).contiguous() 78 | l1_xyz, l1_feat = self.le1(xyz, x) # [B, N, 3] -> [B, N, 16] 79 | l1_feat = l1_feat.transpose(1,2).contiguous() 80 | x = x.transpose(1,2).contiguous() 81 | 82 | # encoder 83 | l1_feat, l1_xyz = self.conv1(l1_feat, l1_xyz, point_num) 84 | l2_feat, l2_xyz = self.conv2(l1_feat, l1_xyz, point_num//4) 85 | l3_feat, l3_xyz = self.conv3(l2_feat, l2_xyz, point_num//8) 86 | l4_feat, l4_xyz = self.conv4(l3_feat, l3_xyz, point_num//16) 87 | l5_feat, l5_xyz = self.conv5(l4_feat, l4_xyz, point_num//32) 88 | 89 | # encode global feature 90 | emb1 = self.global_conv1(torch.cat((l5_xyz, l5_feat), dim=-1)) 91 | emb1 = torch.max(emb1, dim=1, keepdim=True)[0] 92 | 93 | emb2 = self.global_conv2(torch.cat((l4_xyz, l4_feat), dim=-1)) 94 | emb2 = torch.max(emb2, dim=1, keepdim=True)[0] 95 | 96 | if l is not None: 97 | l = l.view(batch_size, 1, -1) # [B, 1, 16] 98 | emb = torch.cat((emb1, emb2, l), dim=-1) 99 | 100 | emb = emb.expand(-1, point_num, -1) 101 | 102 | #decoder 103 | l4_feat = self.fp4(l4_xyz, l5_xyz, l4_feat, l5_feat) 104 | l4_feat, l4_xyz = self.up_conv5(l4_feat, l4_xyz, point_num//16) 105 | 106 | l3_feat = self.fp3(l3_xyz, l4_xyz, l3_feat, l4_feat) 107 | l3_feat, l3_xyz = self.up_conv4(l3_feat, l3_xyz, point_num//8) 108 | 109 | l2_feat = self.fp2(l2_xyz, l3_xyz, l2_feat, l3_feat) 110 | l2_feat, l2_xyz = self.up_conv3(l2_feat, l2_xyz, point_num//4) 111 | 112 | l1_feat = self.fp1(l1_xyz, l2_xyz, l1_feat, l2_feat) 113 | l1_feat, l1_xyz = self.up_conv2(l1_feat, l1_xyz, point_num) 114 | 115 | # feature fusion 116 | feat = torch.cat((l1_xyz, l1_feat, emb), dim=-1) 117 | feat = self.up_conv1(feat) 118 | 119 | feat = feat.transpose(1,2) 120 | score = self.se(feat) 121 | feat = feat.transpose(1,2) 122 | score = score.transpose(1,2) 123 | 124 | feat = feat*score 125 | 126 | feat = self.drop(feat) 127 | feat = self.last_conv(feat) 128 | feat = feat.transpose(1,2) 129 | return feat 130 | 131 | if __name__ == "__main__": 132 | parser = argparse.ArgumentParser(description='Point Cloud Part Segmentation') 133 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 134 | help='Name of the experiment') 135 | parser.add_argument('--dropout', type=float, default=0.5, 136 | help='dropout rate') 137 | parser.add_argument('--radius', type=float, default=0.005, 138 | help='search radius') 139 | args = parser.parse_args() 140 | 141 | device = "cuda" if torch.cuda.is_available() else "cpu" 142 | pc = torch.rand(5, 1024, 3).to(device) 143 | lb = torch.ones(5, 16).to(device) 144 | model = Model(args).to(device) 145 | model = nn.DataParallel(model) 146 | feat = model(pc, lb) 147 | print(feat.shape) -------------------------------------------------------------------------------- /models/diffConv_seg.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from .modules import Conv1x1, diffConv, PointFeaturePropagation 6 | from .pointnet2.pointnet2_modules import PointnetSAModule 7 | 8 | class Model(nn.Module): 9 | def __init__(self, args, num_classes=9): 10 | super().__init__() 11 | 12 | init_channel = 16 13 | 14 | radius = args.radius 15 | 16 | self.le0 = Conv1x1(3, init_channel) 17 | 18 | self.le1 = PointnetSAModule( 19 | npoint=args.num_points, 20 | radius=0.05, 21 | nsample=20, 22 | mlp=[init_channel, init_channel, init_channel], 23 | use_xyz=True, 24 | bn=True 25 | ) 26 | 27 | # encoder 28 | self.conv1 = diffConv(init_channel, init_channel*2, radius) 29 | 30 | self.conv2 = diffConv(init_channel*2, init_channel*4, radius*2) 31 | 32 | self.conv3 = diffConv(init_channel*4, init_channel*8, radius*4) 33 | 34 | self.conv4 = diffConv(init_channel*8, init_channel*16, radius*8) 35 | 36 | self.fp3 = PointFeaturePropagation(in_channel1=init_channel*16, in_channel2=init_channel*8, out_channel=init_channel*8) 37 | self.up_conv4 = diffConv(init_channel*8, init_channel*8, radius*4) 38 | 39 | self.fp2 = PointFeaturePropagation(in_channel1=init_channel*8, in_channel2=init_channel*4, out_channel=init_channel*8) 40 | self.up_conv3 = diffConv(init_channel*8, init_channel*8, radius*2) 41 | 42 | self.fp1 = PointFeaturePropagation(in_channel1=init_channel*8, in_channel2=init_channel*2, out_channel=init_channel*8) 43 | self.up_conv2 = diffConv(init_channel*8, init_channel*8, radius) 44 | 45 | self.fp0 = PointFeaturePropagation(in_channel1=init_channel*8, in_channel2=16, out_channel=init_channel*8) 46 | self.up_conv1 = nn.Sequential(Conv1x1(init_channel*8+3, 256), 47 | nn.Dropout(args.dropout), 48 | Conv1x1(256, 128), 49 | Conv1x1(128, 128)) 50 | 51 | self.conv = nn.Linear(128, num_classes, bias=False) 52 | 53 | def forward(self, x): 54 | xyz = x.clone() 55 | point_num = xyz.size(1) 56 | x = self.le0(x) 57 | 58 | x = x.transpose(1,2).contiguous() 59 | l1_xyz, l1_feat = self.le1(xyz, x) # [B, N, 3] -> [B, N, 16] 60 | l1_feat = l1_feat.transpose(1,2).contiguous() 61 | x = x.transpose(1,2).contiguous() 62 | 63 | # encoder 64 | l1_feat, l1_xyz = self.conv1(l1_feat, l1_xyz, point_num//2) 65 | l2_feat, l2_xyz = self.conv2(l1_feat, l1_xyz, point_num//4) 66 | l3_feat, l3_xyz = self.conv3(l2_feat, l2_xyz, point_num//8) 67 | l4_feat, l4_xyz = self.conv4(l3_feat, l3_xyz, point_num//16) 68 | 69 | l3_feat = self.fp3(l3_xyz, l4_xyz, l3_feat, l4_feat) 70 | l3_feat, l3_xyz = self.up_conv4(l3_feat, l3_xyz, point_num//8) 71 | 72 | l2_feat = self.fp2(l2_xyz, l3_xyz, l2_feat, l3_feat) 73 | l2_feat, l2_xyz = self.up_conv3(l2_feat, l2_xyz, point_num//4) 74 | 75 | l1_feat = self.fp1(l1_xyz, l2_xyz, l1_feat, l2_feat) 76 | l1_feat, l1_xyz = self.up_conv2(l1_feat, l1_xyz, point_num//2) 77 | 78 | # feature fusion 79 | l1_feat = self.fp0(xyz, l1_xyz, x, l1_feat) 80 | feat = torch.cat((xyz, l1_feat), dim=-1) 81 | feat = self.up_conv1(feat) 82 | feat = self.conv(feat) 83 | feat = feat.transpose(1,2) 84 | return feat 85 | 86 | if __name__ == "__main__": 87 | parser = argparse.ArgumentParser(description='Point Cloud Semantic Segmentation') 88 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 89 | help='Name of the experiment') 90 | parser.add_argument('--dropout', type=float, default=0.5, 91 | help='dropout rate') 92 | parser.add_argument('--radius', type=float, default=0.005, 93 | help='search radius') 94 | args = parser.parse_args() 95 | 96 | device = "cuda" if torch.cuda.is_available() else "cpu" 97 | pc = torch.rand(5, 1024, 3).to(device) 98 | model = Model(args).to(device) 99 | model = nn.DataParallel(model) 100 | feat = model(pc) 101 | print(feat.shape) 102 | -------------------------------------------------------------------------------- /models/modules.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from .utils import sample_and_group, index_points, get_dist 6 | from .pointnet2 import pointnet2_utils 7 | 8 | class Conv1x1(nn.Module): 9 | ''' 10 | 1x1 1d convolution 11 | ''' 12 | def __init__(self, in_channels, out_channels, act=nn.GELU(), bias_=False): # nn.LeakyReLU(negative_slope=0.2) 13 | super(Conv1x1, self).__init__() 14 | self.conv = nn.Sequential( 15 | nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=bias_), 16 | nn.BatchNorm1d(out_channels) 17 | ) 18 | self.act = act 19 | nn.init.xavier_normal_(self.conv[0].weight.data) 20 | 21 | def forward(self, x): 22 | x = x.transpose(1, 2).contiguous() 23 | x = self.conv(x) 24 | 25 | x = x.transpose(1, 2).contiguous() 26 | if self.act is not None: 27 | return self.act(x) 28 | else: 29 | return x 30 | 31 | class PositionEncoder(nn.Module): 32 | def __init__(self, out_channel, radius, k=20): 33 | super(PositionEncoder, self).__init__() 34 | self.k = k 35 | 36 | self.xyz2feature = nn.Sequential( 37 | nn.Conv2d(9, out_channel//8, kernel_size=1), 38 | nn.BatchNorm2d(out_channel//8), 39 | nn.GELU() 40 | ) 41 | 42 | self.mlp = nn.Sequential( 43 | Conv1x1(out_channel//8, out_channel//4), 44 | Conv1x1(out_channel//4, out_channel, act=None) 45 | ) 46 | 47 | self.qg = pointnet2_utils.QueryAndGroup(radius, self.k) 48 | 49 | def forward(self, centroid, xyz, radius, dist): 50 | point_feature, _ = sample_and_group(radius, self.k, xyz, xyz, centroid, dist) # [B, N, k, 3] 51 | 52 | points = centroid.unsqueeze(2).repeat(1, 1, self.k, 1) # [B, N, k, 3] 53 | 54 | variance = point_feature - points # [B, N, k, 3] 55 | 56 | point_feature = torch.cat((points, point_feature, variance), dim=-1) # [B, N, k, 9] 57 | 58 | point_feature = point_feature.permute(0, 3, 1, 2).contiguous() # [B, 9, N, k] 59 | 60 | point_feature = self.xyz2feature(point_feature) # [B, 9, N, k] 61 | 62 | point_feature = torch.max(point_feature, dim=-1)[0].transpose(1,2) # [B, N, C] 63 | 64 | point_feature = self.mlp(point_feature) # [B, N, C'] 65 | 66 | return point_feature 67 | 68 | class MaskedAttention(nn.Module): 69 | def __init__(self, in_channels, hid_channels=128): 70 | super().__init__() 71 | if not hid_channels: 72 | hid_channels = 1 73 | self.conv_q = Conv1x1(in_channels+3, hid_channels, act=None)# map query (key points) to another linear space 74 | self.conv_k = Conv1x1(in_channels+3, hid_channels, act=None)# map key (neighbor points) to another linear space 75 | 76 | def forward(self, cent_feat, feat, mask): 77 | ''' 78 | Inputs: 79 | cent_feat: [B, M, C+3] 80 | feat: [B, N, C+3] 81 | mask: [B, M, N] 82 | 83 | Returns: 84 | adj: [B, M, N] 85 | ''' 86 | q = self.conv_q(cent_feat) # [B, M, C+3] -> [B, M, C_int] 87 | 88 | k = self.conv_k(feat) # [B, N, C+3] -> [B, N, C_int] 89 | 90 | adj = torch.bmm(q, k.transpose(1, 2)) # [B, M, C_int] * [B, C_int, N] -> [B, M, N] 91 | 92 | # masked self-attention: masking all non-neighbors (Eq. 9) 93 | adj = adj.masked_fill(mask < 1e-9, -1e9) 94 | adj = torch.softmax(adj, dim=-1) 95 | 96 | # balanced renormalization (Eq. 11) 97 | adj = torch.sqrt(mask + 1e-9) * torch.sqrt(adj + 1e-9) - 1e-9 98 | 99 | adj = F.normalize(adj, p=1, dim=1) # [B, M, N] 100 | adj = F.normalize(adj, p=1, dim=-1) # [B, M, N] 101 | 102 | return adj 103 | 104 | def dilated_ball_query(dist, h, base_radius, max_radius): 105 | ''' 106 | Density-dilated ball query 107 | Inputs: 108 | dist[B, M, N]: distance matrix 109 | h(float): bandwidth 110 | base_radius(float): minimum search radius 111 | max_radius(float): maximum search radius 112 | Returns: 113 | radius[B, M, 1]: search radius of point 114 | ''' 115 | 116 | # kernel density estimation (Eq. 8) 117 | sigma = 1 118 | gauss = torch.exp(-(dist)/(2*(h**2)*(sigma**2))) # K(x-x_i/h), [B, M, N] 119 | kd_dist = torch.sum(gauss, dim=-1).unsqueeze(-1) # kernel distance, [B, M, 1] 120 | 121 | # normalization 122 | kd_score = kd_dist / (torch.max(kd_dist, dim=1)[0].unsqueeze(-1) + 1e-9) # [B, M, 1] 123 | radius = base_radius + (max_radius - base_radius)*kd_score # kd_score -> max, base_radius -> max_radius 124 | 125 | return radius 126 | 127 | class diffConv(nn.Module): 128 | def __init__(self, in_channels, out_channels, base_radius, bottleneck=4): 129 | super().__init__() 130 | self.conv_v = Conv1x1(2*in_channels, out_channels, act=None) 131 | self.mat = MaskedAttention(in_channels, in_channels//bottleneck) 132 | self.pos_conv = PositionEncoder(out_channels, np.sqrt(base_radius)) 133 | self.base_radius = base_radius # squared radius 134 | 135 | def forward(self, x, xyz, cent_num): 136 | ''' 137 | Inputs: 138 | x[B, N, C]: point features 139 | xyz[B, N, 3]: points 140 | cent_num(int): number of key points 141 | 142 | Returns: 143 | x[B, M, C']: updated point features 144 | centroid[B, M, 3]: sampled features 145 | ''' 146 | batch_size, point_num = xyz.size(0), xyz.size(1) 147 | 148 | if cent_num < point_num: 149 | # random sampling 150 | idx = np.arange(point_num) 151 | idx = idx[:cent_num] 152 | idx = torch.from_numpy(idx).unsqueeze(0).repeat(batch_size, 1).int().to(xyz.device) 153 | 154 | # gathering 155 | centroid = index_points(xyz, idx) # [B, M, 3] 156 | cent_feat = index_points(x, idx) # [B, M, C] 157 | else: 158 | centroid = xyz.clone() 159 | cent_feat = x.clone() 160 | 161 | dist = get_dist(centroid, xyz) # disntance matrix, [B, M, N] 162 | 163 | radius = dilated_ball_query(dist, h=0.1, base_radius=self.base_radius, max_radius=self.base_radius*2) 164 | 165 | mask = (dist < radius).float() 166 | 167 | # get attentive mask (adjacency matrix) 168 | emb_cent = torch.cat((cent_feat, centroid), dim=-1) 169 | emb_x = torch.cat((x, xyz), dim=-1) 170 | adj = self.mat(emb_cent, emb_x, mask) # [B, M, N] 171 | 172 | # inner-group attention 173 | smoothed_x = torch.bmm(adj, x) # [B, M, N] * [B, N, C] -> [B, M, C] 174 | variation = smoothed_x - cent_feat # [B, M, C] -> [B, M, C] 175 | 176 | x = torch.cat((variation, cent_feat), dim=-1) # [B, M, C] -> [B, M, 2C] 177 | x = self.conv_v(x) # [B, M, 2C] -> [B, M, C'] 178 | 179 | pos_emb = self.pos_conv(centroid, xyz, radius, dist) 180 | 181 | # feature fusion 182 | x = x + pos_emb 183 | x = F.gelu(x) 184 | 185 | return x, centroid 186 | 187 | class Attention_block(nn.Module): 188 | ''' 189 | attention U-Net is taken from https://github.com/tiangexiang/CurveNet/blob/main/core/models/curvenet_util.py. 190 | ''' 191 | def __init__(self, F_g, F_l, F_int): 192 | super(Attention_block,self).__init__() 193 | self.W_g = nn.Sequential( 194 | nn.Conv1d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True), 195 | nn.BatchNorm1d(F_int) 196 | ) 197 | 198 | self.W_x = nn.Sequential( 199 | nn.Conv1d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True), 200 | nn.BatchNorm1d(F_int) 201 | ) 202 | 203 | self.psi = nn.Sequential( 204 | nn.Conv1d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True), 205 | nn.BatchNorm1d(1), 206 | nn.Sigmoid() 207 | ) 208 | 209 | def forward(self, g, x): 210 | g = g.transpose(1,2) 211 | x = x.transpose(1,2) 212 | g1 = self.W_g(g) 213 | x1 = self.W_x(x) 214 | psi = F.gelu(g1+x1) 215 | psi = self.psi(psi) 216 | psi = psi.transpose(1,2) 217 | 218 | return psi 219 | 220 | class PointFeaturePropagation(nn.Module): 221 | def __init__(self, in_channel1, in_channel2, out_channel): 222 | super(PointFeaturePropagation, self).__init__() 223 | in_channel = in_channel1 + in_channel2 224 | self.conv = nn.Sequential( 225 | Conv1x1(in_channel, in_channel//2), 226 | Conv1x1(in_channel//2, in_channel//2), 227 | Conv1x1(in_channel//2, out_channel) 228 | ) 229 | self.att = Attention_block(in_channel1, in_channel2, in_channel2) 230 | 231 | def forward(self, xyz1, xyz2, feat1, feat2): 232 | """ 233 | Input: 234 | xyz1: input points position data, [B, N, 3] 235 | xyz2: sampled input points position data, [B, M, 3] 236 | feat1: input points data, [B, N, C'] 237 | feat2: input points data, [B, M, C] 238 | Return: 239 | new_points: upsampled points data, [B, N, C+C'] 240 | """ 241 | dists, idx = pointnet2_utils.three_nn(xyz1, xyz2) 242 | dist_recip = 1.0 / (dists + 1e-8) # [B, N, 3] 243 | norm = torch.sum(dist_recip, dim=-1, keepdim=True) # [B, N, 1] 244 | weight = dist_recip / norm # [B, N, 1] 245 | int_feat = pointnet2_utils.three_interpolate(feat2.transpose(1,2).contiguous(), idx, weight).transpose(1,2) 246 | 247 | psix = self.att(int_feat, feat1) 248 | feat1 = feat1 * psix 249 | 250 | if feat1 is not None: 251 | cat_feat = torch.cat((feat1, int_feat), dim=-1) # [B, N, C'], [B, N, C] -> [B, N, C + C'] 252 | else: 253 | cat_feat = int_feat # [B, N, C] 254 | cat_feat = self.conv(cat_feat) # [B, N, C + C'] -> [B, N, C'] 255 | 256 | return cat_feat 257 | 258 | if __name__ == "__main__": 259 | device = "cuda" if torch.cuda.is_available() else "cpu" 260 | pc = torch.rand(5, 1024, 3).to(device) 261 | feat = torch.rand(5, 1024, 16).to(device) 262 | model = diffConv(16, 32, 0.1).to(device) 263 | new_feat, new_pc = model(feat, pc, 512) 264 | print(new_feat.shape, new_pc.shape) 265 | 266 | 267 | -------------------------------------------------------------------------------- /models/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | sys.path.append('./models/pointnet2') 4 | from .pointnet2 import pointnet2_utils 5 | 6 | def get_dist(src, dst): 7 | """ 8 | Calculate the Euclidean distance between each point pair in two point clouds. 9 | Inputs: 10 | src[B, M, 3]: point cloud 1 11 | dst[B, N, 3]: point cloud 2 12 | Return: 13 | dist[B, M, N]: distance matrix 14 | """ 15 | B, N, _ = src.shape 16 | _, M, _ = dst.shape 17 | dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) 18 | dist += torch.sum(src ** 2, -1).view(B, N, 1) 19 | dist += torch.sum(dst ** 2, -1).view(B, 1, M) 20 | return dist 21 | 22 | def index_points(points, idx): 23 | """ 24 | Input: 25 | points[B, N, C]: input point features 26 | idx[B, M]: sample index data 27 | Return: 28 | new_points[B, M, C]: quried point features 29 | """ 30 | new_points = pointnet2_utils.gather_operation(points.transpose(1,2).contiguous(), idx).transpose(1,2).contiguous() 31 | return new_points 32 | 33 | def sample_and_group(radius, k, xyz, feat, centroid, dist): 34 | """ 35 | Input: 36 | radius[B, M, 1]: search radius of each key point 37 | k(int): max number of samples in local region 38 | xyz[B, N, 3]: query points 39 | centroid[B, M, 3]: key points 40 | dist[B, M, N]: distance matrix 41 | feat[B, N, D]: input points features 42 | Return: 43 | cent_feat[B, M, D]: grouped features 44 | idx[B, M, k]: indices of selected neighbors 45 | """ 46 | device = xyz.device 47 | B, N, _ = xyz.shape 48 | _, M, _ = centroid.shape 49 | 50 | idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, M, 1]) 51 | 52 | idx[dist > radius] = N 53 | idx = idx.sort(dim=-1)[0][:, :, :k] 54 | group_first = idx[:, :, 0].view(B, M, 1).repeat([1, 1, k]) 55 | mask = (idx == N) 56 | idx[mask] = group_first[mask] 57 | 58 | torch.cuda.empty_cache() 59 | idx = idx.int().contiguous() 60 | 61 | feat = feat.transpose(1,2).contiguous() 62 | cent_feat = pointnet2_utils.grouping_operation(feat, idx) 63 | cent_feat = cent_feat.transpose(1,2).transpose(-1, -2).contiguous() 64 | torch.cuda.empty_cache() 65 | 66 | return cent_feat, idx 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | --------------------------------------------------------------------------------