├── .gitignore ├── CODE_OF_CONDUCT.md ├── Frechet ├── FPD.py ├── cls_model_39.pth ├── pointnet.py ├── pre_statistics_all.npz └── test.list ├── LICENSE ├── README.md ├── SECURITY.md ├── SUPPORT.md ├── configs ├── atlasnet.yaml ├── base_config.py ├── grnet.yaml ├── model_names.py ├── msn.yaml ├── sparenet.yaml └── sparenet_gan.yaml ├── cuda ├── MDS │ ├── MDS.cpp │ ├── MDS.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── MDS_cuda.cu │ ├── MDS_module.py │ ├── dist │ │ ├── MDS-0.0.0-py3.6-linux-x86_64.egg │ │ └── MDS-0.0.0-py3.7-linux-x86_64.egg │ └── setup.py ├── chamfer_dist │ ├── __init__.py │ ├── chamfer.cu │ ├── chamfer_cuda.cpp │ ├── setup.py │ └── test.py ├── chamfer_distance │ ├── __init__.py │ ├── chamfer_distance.cpp │ ├── chamfer_distance.cu │ └── chamfer_distance.py ├── cubic_feature_sampling │ ├── __init__.py │ ├── cubic_feature_sampling.cu │ ├── cubic_feature_sampling.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── cubic_feature_sampling_cuda.cpp │ ├── dist │ │ └── cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg │ ├── setup.py │ └── test.py ├── emd │ ├── CDEMD.png │ ├── README.md │ ├── dist │ │ ├── emd-0.0.0-py3.6-linux-x86_64.egg │ │ ├── emd-0.0.0-py3.7-linux-x86_64.egg │ │ └── emd-0.0.0-py3.8-linux-x86_64.egg │ ├── emd.cpp │ ├── emd.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── emd_cuda.cu │ ├── emd_module.py │ └── setup.py ├── expansion_penalty │ ├── dist │ │ ├── expansion_penalty-0.0.0-py3.6-linux-x86_64.egg │ │ ├── expansion_penalty-0.0.0-py3.7-linux-x86_64.egg │ │ └── expansion_penalty-0.0.0-py3.8-linux-x86_64.egg │ ├── expansion_penalty.cpp │ ├── expansion_penalty.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── expansion_penalty_cuda.cu │ ├── expansion_penalty_module.py │ └── setup.py ├── gridding │ ├── __init__.py │ ├── dist │ │ └── gridding-2.1.0-py3.6-linux-x86_64.egg │ ├── gridding.cu │ ├── gridding.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── gridding_cuda.cpp │ ├── gridding_reverse.cu │ ├── setup.py │ └── test.py ├── gridding_loss │ ├── __init__.py │ ├── dist │ │ └── gridding_distance-1.0.0-py3.6-linux-x86_64.egg │ ├── gridding_distance.cu │ ├── gridding_distance.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── gridding_distance_cuda.cpp │ └── setup.py └── p2i_op │ ├── .clang-format │ ├── __init__.py │ ├── common.h │ ├── ext.cpp │ ├── p2i_max.cu │ ├── p2i_max.h │ ├── p2i_sum.cu │ ├── p2i_sum.h │ ├── p2i_test.py │ └── utility.h ├── datasets ├── README.md ├── data │ ├── Completion3D.json │ ├── KITTI.json │ └── ShapeNet.json ├── data_loaders.py ├── data_transforms.py └── io.py ├── models ├── __init__.py ├── atlasnet_generator.py ├── grnet_generator.py ├── msn_generator.py ├── sparenet_discriminator.py └── sparenet_generator.py ├── renderer.png ├── requirements.txt ├── runners ├── atlasnet_runner.py ├── base_runner.py ├── grnet_runner.py ├── misc.py ├── msn_runner.py ├── sparenet_gan_runner.py └── sparenet_runner.py ├── scripts ├── atlasnet_test.sh ├── atlasnet_train.sh ├── grnet_test.sh ├── grnet_train.sh ├── msn_test.sh ├── msn_train.sh ├── sparenet_gan_train.sh ├── sparenet_test.sh └── sparenet_train.sh ├── setup_env.sh ├── teaser.png ├── test.py ├── test_fpd.py ├── train.py └── utils ├── __init__.py ├── misc.py ├── model_init.py ├── p2i_utils.py └── visualizer.py /.gitignore: -------------------------------------------------------------------------------- 1 | */.DS_Store 2 | .DS_Store 3 | .vscode/ 4 | output/ 5 | logs/ 6 | __temp__/ 7 | __pycache__/ 8 | backup/ 9 | build/ 10 | lightning_logs/ -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Microsoft Open Source Code of Conduct 2 | 3 | This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). 4 | 5 | Resources: 6 | 7 | - [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/) 8 | - [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) 9 | - Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns 10 | -------------------------------------------------------------------------------- /Frechet/FPD.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pathlib 3 | from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter 4 | 5 | import numpy as np 6 | import torch 7 | from scipy.linalg import sqrtm 8 | from scipy.misc import imread 9 | from torch.nn.functional import adaptive_avg_pool2d 10 | from Frechet.pointnet import PointNetCls 11 | 12 | try: 13 | from tqdm import tqdm 14 | except ImportError: 15 | # If not tqdm is not available, provide a mock version of it 16 | def tqdm(x): return x 17 | 18 | """Calculate Frechet Pointcloud Distance referened by Frechet Inception Distance." 19 | [ref] GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium 20 | github code : (https://github.com/bioinf-jku/TTUR) 21 | paper : (https://arxiv.org/abs/1706.08500) 22 | 23 | """ 24 | 25 | def get_activations(pointclouds, model, batch_size=100, dims=1808, 26 | device=None, verbose=False): 27 | """Calculates the activations of the pool_3 layer for all images. 28 | Params: 29 | -- pointcloud : pytorch Tensor of pointclouds. 30 | -- model : Instance of inception model 31 | -- batch_size : Batch size of images for the model to process at once. 32 | Make sure that the number of samples is a multiple of 33 | the batch size, otherwise some samples are ignored. This 34 | behavior is retained to match the original FID score 35 | implementation. 36 | -- dims : Dimensionality of features returned by Inception 37 | -- device : If set to device, use GPU 38 | -- verbose : If set to True and parameter out_step is given, the number 39 | of calculated batches is reported. 40 | Returns: 41 | -- A numpy array of dimension (num images, dims) that contains the 42 | activations of the given tensor when feeding inception with the 43 | query tensor. 44 | """ 45 | model.eval() 46 | 47 | n_batches = pointclouds.size(0) // batch_size 48 | n_used_imgs = n_batches * batch_size 49 | 50 | pred_arr = np.empty((n_used_imgs, dims)) 51 | 52 | pointclouds = pointclouds.transpose(1,2) 53 | for i in tqdm(range(n_batches)): 54 | if verbose: 55 | print('\rPropagating batch %d/%d' % (i + 1, n_batches), 56 | end='', flush=True) 57 | start = i * batch_size 58 | end = start + batch_size 59 | 60 | pointcloud_batch = pointclouds[start:end] 61 | 62 | if device is not None: 63 | pointcloud_batch = pointcloud_batch.to(device) 64 | 65 | _, _, actv = model(pointcloud_batch) 66 | 67 | # If model output is not scalar, apply global spatial average pooling. 68 | # This happens if you choose a dimensionality not equal 2048. 69 | # if pred.shape[2] != 1 or pred.shape[3] != 1: 70 | # pred = adaptive_avg_pool2d(pred, output_size=(1, 1)) 71 | 72 | pred_arr[start:end] = actv.cpu().data.numpy().reshape(batch_size, -1) 73 | 74 | if verbose: 75 | print(' done') 76 | 77 | return pred_arr 78 | 79 | 80 | def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): 81 | """Numpy implementation of the Frechet Distance. 82 | The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) 83 | and X_2 ~ N(mu_2, C_2) is 84 | d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). 85 | Stable version by Dougal J. Sutherland. 86 | Params: 87 | -- mu1 : Numpy array containing the activations of a layer of the 88 | inception net (like returned by the function 'get_predictions') 89 | for generated samples. 90 | -- mu2 : The sample mean over activations, precalculated on an 91 | representative data set. 92 | -- sigma1: The covariance matrix over activations for generated samples. 93 | -- sigma2: The covariance matrix over activations, precalculated on an 94 | representative data set. 95 | Returns: 96 | -- : The Frechet Distance. 97 | """ 98 | 99 | mu1 = np.atleast_1d(mu1) 100 | mu2 = np.atleast_1d(mu2) 101 | 102 | sigma1 = np.atleast_2d(sigma1) 103 | sigma2 = np.atleast_2d(sigma2) 104 | 105 | assert mu1.shape == mu2.shape, \ 106 | 'Training and test mean vectors have different lengths' 107 | assert sigma1.shape == sigma2.shape, \ 108 | 'Training and test covariances have different dimensions' 109 | 110 | diff = mu1 - mu2 111 | 112 | # Product might be almost singular 113 | covmean, _ = sqrtm(sigma1.dot(sigma2), disp=False) 114 | if not np.isfinite(covmean).all(): 115 | msg = ('fid calculation produces singular product; ' 116 | 'adding %s to diagonal of cov estimates') % eps 117 | print(msg) 118 | offset = np.eye(sigma1.shape[0]) * eps 119 | covmean = sqrtm((sigma1 + offset).dot(sigma2 + offset)) 120 | 121 | # Numerical error might give slight imaginary component 122 | if np.iscomplexobj(covmean): 123 | if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): 124 | m = np.max(np.abs(covmean.imag)) 125 | raise ValueError('Imaginary component {}'.format(m)) 126 | covmean = covmean.real 127 | 128 | tr_covmean = np.trace(covmean) 129 | 130 | return (diff.dot(diff) + np.trace(sigma1) + 131 | np.trace(sigma2) - 2 * tr_covmean) 132 | 133 | 134 | def calculate_activation_statistics(pointclouds, model, batch_size=100, 135 | dims=1808, device=None, verbose=False): 136 | """Calculation of the statistics used by the FID. 137 | Params: 138 | -- pointcloud : pytorch Tensor of pointclouds. 139 | -- model : Instance of inception model 140 | -- batch_size : The images numpy array is split into batches with 141 | batch size batch_size. A reasonable batch size 142 | depends on the hardware. 143 | -- dims : Dimensionality of features returned by Inception 144 | -- device : If set to device, use GPU 145 | -- verbose : If set to True and parameter out_step is given, the 146 | number of calculated batches is reported. 147 | Returns: 148 | -- mu : The mean over samples of the activations of the pool_3 layer of 149 | the inception model. 150 | -- sigma : The covariance matrix of the activations of the pool_3 layer of 151 | the inception model. 152 | """ 153 | act = get_activations(pointclouds, model, batch_size, dims, device, verbose) 154 | mu = np.mean(act, axis=0) 155 | sigma = np.cov(act, rowvar=False) 156 | return mu, sigma 157 | 158 | 159 | def _compute_statistics_of_path(path, model, batch_size, dims, cuda): 160 | if path.endswith('.npz'): 161 | f = np.load(path) 162 | m, s = f['m'][:], f['s'][:] 163 | f.close() 164 | else: 165 | path = pathlib.Path(path) 166 | files = list(path.glob('*.jpg')) + list(path.glob('*.png')) 167 | m, s = calculate_activation_statistics(files, model, batch_size, 168 | dims, cuda) 169 | 170 | return m, s 171 | 172 | def save_statistics(real_pointclouds, path, model, batch_size, dims, cuda): 173 | m, s = calculate_activation_statistics(real_pointclouds, model, batch_size, 174 | dims, cuda) 175 | np.savez(path, m = m, s = s) 176 | print('save done !!!') 177 | 178 | def calculate_fpd(pointclouds1, pointclouds2=None, statistic_save_path=None, batch_size=100, dims=1808, device=None): 179 | """Calculates the FPD of two pointclouds""" 180 | 181 | PointNet_pretrained_path = './Frechet/cls_model_39.pth' 182 | if statistic_save_path is None: 183 | statistic_save_path = './Frechet/pre_statistics_all.npz' 184 | model = PointNetCls(k=16) 185 | model.load_state_dict(torch.load(PointNet_pretrained_path)) 186 | 187 | if device is not None: 188 | model.to(device) 189 | 190 | 191 | m1, s1 = calculate_activation_statistics(pointclouds1, model, batch_size, dims, device) 192 | if pointclouds2 is not None: 193 | m2, s2 = calculate_activation_statistics(pointclouds2, model, batch_size, dims, device) 194 | else: # Load saved statistics of real pointclouds. 195 | f = np.load(statistic_save_path) 196 | m2, s2 = f['m'][:], f['s'][:] 197 | f.close() 198 | 199 | 200 | fid_value = calculate_frechet_distance(m1, s1, m2, s2) 201 | 202 | return fid_value 203 | -------------------------------------------------------------------------------- /Frechet/cls_model_39.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/Frechet/cls_model_39.pth -------------------------------------------------------------------------------- /Frechet/pointnet.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import torch.nn as nn 4 | from torch.autograd import Variable 5 | import numpy as np 6 | import torch.nn.functional as F 7 | 8 | """PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation" 9 | github code : (https://github.com/charlesq34/pointnet) 10 | project page : (http://stanford.edu/~rqi/pointnet/) 11 | paper : (https://arxiv.org/abs/1612.00593) 12 | """ 13 | 14 | class STN3d(nn.Module): 15 | def __init__(self): 16 | super(STN3d, self).__init__() 17 | self.conv1 = torch.nn.Conv1d(3, 64, 1) 18 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 19 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 20 | self.fc1 = nn.Linear(1024, 512) 21 | self.fc2 = nn.Linear(512, 256) 22 | self.fc3 = nn.Linear(256, 9) 23 | self.relu = nn.ReLU() 24 | 25 | self.bn1 = nn.BatchNorm1d(64) 26 | self.bn2 = nn.BatchNorm1d(128) 27 | self.bn3 = nn.BatchNorm1d(1024) 28 | self.bn4 = nn.BatchNorm1d(512) 29 | self.bn5 = nn.BatchNorm1d(256) 30 | 31 | 32 | def forward(self, x): 33 | batchsize = x.size()[0] 34 | x = F.relu(self.bn1(self.conv1(x))) 35 | x = F.relu(self.bn2(self.conv2(x))) 36 | x = F.relu(self.bn3(self.conv3(x))) 37 | x = torch.max(x, 2, keepdim=True)[0] 38 | x = x.view(-1, 1024) 39 | 40 | x = F.relu(self.bn4(self.fc1(x))) 41 | x = F.relu(self.bn5(self.fc2(x))) 42 | x = self.fc3(x) 43 | 44 | iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1) 45 | if x.is_cuda: 46 | iden = iden.cuda() 47 | x = x + iden 48 | x = x.view(-1, 3, 3) 49 | return x 50 | 51 | 52 | class PointNetfeat(nn.Module): 53 | def __init__(self, global_feat = True): 54 | super(PointNetfeat, self).__init__() 55 | self.stn = STN3d() 56 | self.conv1 = torch.nn.Conv1d(3, 64, 1) 57 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 58 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 59 | self.bn1 = nn.BatchNorm1d(64) 60 | self.bn2 = nn.BatchNorm1d(128) 61 | self.bn3 = nn.BatchNorm1d(1024) 62 | self.global_feat = global_feat 63 | def forward(self, x): 64 | batchsize = x.size()[0] 65 | n_pts = x.size()[2] 66 | trans = self.stn(x) 67 | x = x.transpose(2,1) 68 | x = torch.bmm(x, trans) 69 | x = x.transpose(2,1) 70 | x = F.relu(self.bn1(self.conv1(x))) 71 | pointfeat = x 72 | x = F.relu(self.bn2(self.conv2(x))) 73 | x = self.bn3(self.conv3(x)) 74 | x = torch.max(x, 2, keepdim=True)[0] 75 | x = x.view(-1, 1024) 76 | if self.global_feat: 77 | return x, trans 78 | else: 79 | x = x.view(-1, 1024, 1).repeat(1, 1, n_pts) 80 | return torch.cat([x, pointfeat], 1), trans 81 | 82 | class PointNetCls(nn.Module): 83 | def __init__(self, k = 2): 84 | super(PointNetCls, self).__init__() 85 | self.feat = PointNetfeat(global_feat=True) 86 | self.fc1 = nn.Linear(1024, 512) 87 | self.fc2 = nn.Linear(512, 256) 88 | self.fc3 = nn.Linear(256, k) 89 | self.bn1 = nn.BatchNorm1d(512) 90 | self.bn2 = nn.BatchNorm1d(256) 91 | self.relu = nn.ReLU() 92 | def forward(self, x): 93 | x1, trans = self.feat(x) 94 | x2 = F.relu(self.bn1(self.fc1(x1))) 95 | x3 = F.relu(self.bn2(self.fc2(x2))) 96 | x4 = self.fc3(x3) 97 | 98 | actv = torch.cat((x1, x2, x3, x4), dim=1) 99 | 100 | return F.log_softmax(x4, dim=0), trans, actv 101 | 102 | class PointNetDenseCls(nn.Module): 103 | def __init__(self, k = 2): 104 | super(PointNetDenseCls, self).__init__() 105 | self.k = k 106 | self.feat = PointNetfeat(global_feat=False) 107 | self.conv1 = torch.nn.Conv1d(1088, 512, 1) 108 | self.conv2 = torch.nn.Conv1d(512, 256, 1) 109 | self.conv3 = torch.nn.Conv1d(256, 128, 1) 110 | self.conv4 = torch.nn.Conv1d(128, self.k, 1) 111 | self.bn1 = nn.BatchNorm1d(512) 112 | self.bn2 = nn.BatchNorm1d(256) 113 | self.bn3 = nn.BatchNorm1d(128) 114 | 115 | def forward(self, x): 116 | batchsize = x.size()[0] 117 | n_pts = x.size()[2] 118 | x, trans = self.feat(x) 119 | x = F.relu(self.bn1(self.conv1(x))) 120 | x = F.relu(self.bn2(self.conv2(x))) 121 | x = F.relu(self.bn3(self.conv3(x))) 122 | x = self.conv4(x) 123 | x = x.transpose(2,1).contiguous() 124 | x = F.log_softmax(x.view(-1,self.k), dim=-1) 125 | x = x.view(batchsize, n_pts, self.k) 126 | return x, trans 127 | 128 | 129 | if __name__ == '__main__': 130 | sim_data = Variable(torch.rand(32,3,2500)) 131 | trans = STN3d() 132 | out = trans(sim_data) 133 | print('stn', out.size()) 134 | 135 | pointfeat = PointNetfeat(global_feat=True) 136 | out, _ = pointfeat(sim_data) 137 | print('global feat', out.size()) 138 | 139 | pointfeat = PointNetfeat(global_feat=False) 140 | out, _ = pointfeat(sim_data) 141 | print('point feat', out.size()) 142 | 143 | cls = PointNetCls(k = 5) 144 | out, _ = cls(sim_data) 145 | print('class', out.size()) 146 | 147 | seg = PointNetDenseCls(k = 3) 148 | out, _ = seg(sim_data) 149 | print('seg', out.size()) 150 | -------------------------------------------------------------------------------- /Frechet/pre_statistics_all.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/Frechet/pre_statistics_all.npz -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) Microsoft Corporation. 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 | # Style-based Point Generator with Adversarial Rendering for Point Cloud Completion (CVPR 2021) 2 | 3 | An efficient PyTorch library for Point Cloud Completion. 4 | 5 | 6 |
7 | 8 |
9 | 10 | ### [Project page](https://alphapav.github.io/SpareNet/) | [Paper](https://arxiv.org/abs/2103.02535) | [Video]() 11 | 12 | [Chulin Xie*](https://github.com/AlphaPav), [Chuxin Wang*](https://chuxwa.github.io/), [Bo Zhang](https://bo-zhang.me/), [Hao Yang](https://www.microsoft.com/en-us/research/people/haya/), [Dong Chen](https://www.microsoft.com/en-us/research/people/doch/), and [Fang Wen](https://www.microsoft.com/en-us/research/people/fangwen/). (\*Equal contribution) 13 | 14 | ## Abstract 15 | >We proposed a novel Style-based Point Generator with Adversarial Rendering (SpareNet) for point cloud completion. Firstly, we present the channel-attentive EdgeConv to fully exploit the local structures as well as the global shape in point features. Secondly, we observe that the concatenation manner used by vanilla foldings limits its potential of generating a complex and faithful shape. Enlightened by the success of StyleGAN, we regard the shape feature as style code that modulates the normalization layers during the folding, which considerably enhances its capability. Thirdly, we realize that existing point supervisions, e.g., Chamfer Distance or Earth Mover’s Distance, cannot faithfully reflect the perceptual quality of the reconstructed points. To address this, we propose to project the completed points to depth maps with a differentiable renderer and apply adversarial training to advocate the perceptual realism under different viewpoints. Comprehensive experiments on ShapeNet and KITTI prove the effectiveness of our method, which achieves state-of-the-art quantitative performance while offering superior visual quality. 16 | 17 | 18 | ## Installation 19 | 20 | 1. Create a virtual environment via `conda`. 21 | 22 | ```shell 23 | conda create -n sparenet python=3.7 24 | conda activate sparenet 25 | ``` 26 | 27 | 2. Install `torch` and `torchvision`. 28 | 29 | ```shell 30 | conda install pytorch cudatoolkit=10.1 torchvision -c pytorch 31 | ``` 32 | 33 | 3. Install requirements. 34 | 35 | ```shell 36 | pip install -r requirements.txt 37 | ``` 38 | 39 | 4. Install cuda 40 | ```shell 41 | sh setup_env.sh 42 | ``` 43 | 44 | 45 | ## Dataset 46 | * Download [the processed ShapeNet dataset](https://gateway.infinitescript.com/?fileName=ShapeNetCompletion) (16384 points) generated by [GRNet](https://github.com/hzxie/GRNet), and the [KITTI dataset](https://drive.google.com/drive/folders/1fSu0_huWhticAlzLh3Ejpg8zxzqO1z-F). 47 | 48 | * Update the file path of the datasets in `configs/base_config.py`: 49 | 50 | ``` 51 | __C.DATASETS.shapenet.partial_points_path = "/path/to/datasets/ShapeNetCompletion/%s/partial/%s/%s/%02d.pcd" 52 | __C.DATASETS.shapenet.complete_points_path = "/path/to/datasets/ShapeNetCompletion/%s/complete/%s/%s.pcd" 53 | __C.DATASETS.kitti.partial_points_path = "/path/to/datasets/KITTI/cars/%s.pcd" 54 | __C.DATASETS.kitti.bounding_box_file_path = "/path/to/datasets/KITTI/bboxes/%s.txt" 55 | 56 | # Dataset Options: ShapeNet, ShapeNetCars, KITTI 57 | __C.DATASET.train_dataset = "ShapeNet" 58 | __C.DATASET.test_dataset = "ShapeNet" 59 | ``` 60 | 61 | 62 | ## Get Started 63 | 64 | ### Inference Using Pretrained Model 65 | 66 | The pretrained models: 67 | 68 | - [SpareNet for ShapeNet](https://drive.google.com/file/d/15PiH-bRlSlK4AUUnVwREzuAlMVJ9TfQG) 69 | - [PCN for ShapeNet](https://drive.google.com/drive/folders/1ruN16MlJm4OeRMd41C19HyWqYOIrNrNh) 70 | - [GRNet for ShapeNet](https://gateway.infinitescript.com/?fileName=GRNet-ShapeNet.pth) 71 | - [GRNet for KITTI](https://gateway.infinitescript.com/?fileName=GRNet-KITTI.pth) 72 | - [MSN for ShapeNet](https://drive.google.com/drive/folders/14UZXKqXIZ0gL3hhrV2ySll_pH2eLGFL5) (for 8192 points) 73 | 74 | 75 | - run 76 | 77 | ```shell 78 | python test.py --gpu ${GPUS}\ 79 | --workdir ${WORK_DIR} \ 80 | --model ${network} \ 81 | --weights ${path to checkpoint} \ 82 | --test_mode ${mode} 83 | ``` 84 | 85 | - example 86 | ```shell 87 | python test.py --gpu 0 --workdir /path/to/logfiles --model sparenet --weights /path/to/checkpoint --test_mode default 88 | ``` 89 | 90 | ### Train 91 | 92 | All log files in the training process, such as log message, checkpoints, etc, will be saved to the work directory. 93 | 94 | - run 95 | 96 | ```shell 97 | python train.py --gpu ${GPUS}\ 98 | --workdir ${WORK_DIR} \ 99 | --model ${network} \ 100 | --weights ${path to checkpoint} 101 | ``` 102 | - example 103 | ```shell 104 | python train.py --gpu 0,1,2,3 --workdir /path/to/logfiles --model sparenet --weights /path/to/checkpoint 105 | ``` 106 | 107 | 108 | 109 | ## Differentiable Renderer 110 | A fully differentiable point renderer that enables end-to-end rendering from 3D point cloud to 2D depth maps. See the paper for details. 111 | 112 | 113 |
114 | 115 |
116 | 117 | 118 | 119 | ### Usage of Renderer 120 | 121 | The inputs of renderer are pcd, views and radius, and the outputs of renderer are depth_maps. 122 | - example 123 | ```shell 124 | # `projection_mode`: a str with value "perspective" or "orthorgonal" 125 | # `eyepos_scale`: a float that defines the distance of eyes to (0, 0, 0) 126 | # `image_size`: an int defining the output image size 127 | renderer = ComputeDepthMaps(projection_mode, eyepos_scale, image_size) 128 | 129 | # `data`: a tensor with shape [batch_size, num_points, 3] 130 | # `view_id`: the index of selected view satisfying 0 <= view_id < 8 131 | # `radius_list`: a list of floats, defining the kernel radius to render each point 132 | depthmaps = renderer(data, view_id, radius_list) 133 | ``` 134 | 135 | ### Test FPD on ShapeNet Dataset 136 | * Run your model and save your results of test dataset 137 | 138 | * Update the file path of the results in `test_fpd.py` and run it: 139 | ``` 140 | parser.add_argument('--log_dir', default='/path/to/save/logs') 141 | parser.add_argument('--data_dir', default='/path/to/test/dataset/pcds') 142 | parser.add_argument('--fake_dir', default='/path/to/methods/pcds', 143 | help='/path/to/results/shapenet_fc/pcds/') 144 | ``` 145 | 146 | ## License 147 | 148 | The codes and the pretrained model in this repository are under the MIT license as specified by the LICENSE file. 149 | 150 | This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. 151 | 152 | 153 | ## BibTex 154 | 155 | If you like our work and use the codebase or models for your research, please cite our work as follows. 156 | 157 | ```bibtex 158 | @InProceedings{Xie_2021_CVPR, 159 | author = {Xie, Chulin and Wang, Chuxin and Zhang, Bo and Yang, Hao and Chen, Dong and Wen, Fang}, 160 | title = {Style-Based Point Generator With Adversarial Rendering for Point Cloud Completion}, 161 | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 162 | month = {June}, 163 | year = {2021}, 164 | pages = {4619-4628} 165 | } 166 | ``` 167 | 168 | 170 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | ## Security 4 | 5 | Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet), [Xamarin](https://github.com/xamarin), and [our GitHub organizations](https://opensource.microsoft.com/). 6 | 7 | If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://docs.microsoft.com/en-us/previous-versions/tn-archive/cc751383(v=technet.10)), please report it to us as described below. 8 | 9 | ## Reporting Security Issues 10 | 11 | **Please do not report security vulnerabilities through public GitHub issues.** 12 | 13 | Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://msrc.microsoft.com/create-report). 14 | 15 | If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://www.microsoft.com/en-us/msrc/pgp-key-msrc). 16 | 17 | You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc). 18 | 19 | Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue: 20 | 21 | * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.) 22 | * Full paths of source file(s) related to the manifestation of the issue 23 | * The location of the affected source code (tag/branch/commit or direct URL) 24 | * Any special configuration required to reproduce the issue 25 | * Step-by-step instructions to reproduce the issue 26 | * Proof-of-concept or exploit code (if possible) 27 | * Impact of the issue, including how an attacker might exploit the issue 28 | 29 | This information will help us triage your report more quickly. 30 | 31 | If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://microsoft.com/msrc/bounty) page for more details about our active programs. 32 | 33 | ## Preferred Languages 34 | 35 | We prefer all communications to be in English. 36 | 37 | ## Policy 38 | 39 | Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://www.microsoft.com/en-us/msrc/cvd). 40 | 41 | -------------------------------------------------------------------------------- /SUPPORT.md: -------------------------------------------------------------------------------- 1 | # TODO: The maintainer of this repo has not yet edited this file 2 | 3 | **REPO OWNER**: Do you want Customer Service & Support (CSS) support for this product/project? 4 | 5 | - **No CSS support:** Fill out this template with information about how to file issues and get help. 6 | - **Yes CSS support:** Fill out an intake form at [aka.ms/spot](https://aka.ms/spot). CSS will work with/help you to determine next steps. More details also available at [aka.ms/onboardsupport](https://aka.ms/onboardsupport). 7 | - **Not sure?** Fill out a SPOT intake as though the answer were "Yes". CSS will help you decide. 8 | 9 | *Then remove this first heading from this SUPPORT.MD file before publishing your repo.* 10 | 11 | # Support 12 | 13 | ## How to file issues and get help 14 | 15 | This project uses GitHub Issues to track bugs and feature requests. Please search the existing 16 | issues before filing new issues to avoid duplicates. For new issues, file your bug or 17 | feature request as a new Issue. 18 | 19 | For help and questions about using this project, please **REPO MAINTAINER: INSERT INSTRUCTIONS HERE 20 | FOR HOW TO ENGAGE REPO OWNERS OR COMMUNITY FOR HELP. COULD BE A STACK OVERFLOW TAG OR OTHER 21 | CHANNEL. WHERE WILL YOU HELP PEOPLE?**. 22 | 23 | ## Microsoft Support Policy 24 | 25 | Support for this **PROJECT or PRODUCT** is limited to the resources listed above. 26 | -------------------------------------------------------------------------------- /configs/atlasnet.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 2 | train_dataset: ShapeNet 3 | test_dataset: ShapeNet 4 | n_outpoints: 16384 5 | 6 | CONST: 7 | num_workers: 32 8 | n_input_points: 3000 9 | 10 | NETWORK: 11 | n_sampling_points: 2048 12 | n_primitives: 32 13 | model_type: "AtlasNet" 14 | metric: "emd" 15 | 16 | TRAIN: 17 | learning_rate: 0.0001 18 | batch_size: 32 19 | n_epochs: 150 20 | save_freq: 5 21 | log_freq: 1 22 | 23 | TEST: 24 | metric_name: "EMD" 25 | -------------------------------------------------------------------------------- /configs/base_config.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import yaml 6 | import datetime 7 | import platform 8 | import numpy as np 9 | from easydict import EasyDict as edict 10 | 11 | # create dict 12 | __C = edict() 13 | cfg = __C 14 | 15 | # Dataset 16 | __C.DATASET = edict() 17 | # Dataset Options: 'Completion3D', 'ShapeNet', 'ShapeNetCars', 'KITTI' 18 | __C.DATASET.train_dataset = "ShapeNet" 19 | __C.DATASET.test_dataset = "ShapeNet" 20 | __C.DATASET.n_outpoints = 16384 21 | __C.DATASET.num_class = 0 # overwrite by `train_MVGAN.py` 22 | 23 | # Constants 24 | __C.CONST = edict() 25 | __C.CONST.device = "0" 26 | __C.CONST.weights = None 27 | __C.CONST.num_workers = 32 28 | __C.CONST.n_input_points = 3000 29 | 30 | # Directories 31 | __C.DIR = edict() 32 | __C.DIR.out_path = "./output" 33 | __C.DIR.in_path = "./output/checkpoints" 34 | 35 | # Network 36 | __C.NETWORK = edict() 37 | __C.NETWORK.n_sampling_points = 2048 38 | __C.NETWORK.gridding_loss_scales = [128, 64] 39 | __C.NETWORK.gridding_loss_alphas = [0.1, 0.01] 40 | __C.NETWORK.n_primitives = 16 41 | __C.NETWORK.model_type = "SpareNet" 42 | __C.NETWORK.metric = "emd" 43 | __C.NETWORK.encode = "Residualnet" 44 | __C.NETWORK.use_adain = "share" 45 | __C.NETWORK.use_selayer = False 46 | __C.NETWORK.use_consist_loss = False 47 | 48 | # Apex training; BUG: currently cannot support multi-gpu training 49 | __C.APEX = edict() 50 | __C.APEX.flag = False 51 | __C.APEX.level = "O1" 52 | 53 | # RENDER 54 | __C.RENDER = edict() 55 | __C.RENDER.img_size = 256 56 | __C.RENDER.radius_list = [ 57 | 5.0, 58 | 7.0, 59 | 10.0, 60 | ] # [5.0, 7.0, 10.0] for shapenet; [10.0, 15.0, 20.0] for completion3d 61 | __C.RENDER.projection = "orthorgonal" # 'orthorgonal' or 'perspective' 62 | __C.RENDER.eyepos = 1.0 # not use 63 | __C.RENDER.n_views = 8 # number or rendered views 64 | 65 | # GAN training 66 | __C.GAN = edict() 67 | __C.GAN.use_im = True # image-levle loss 68 | __C.GAN.use_fm = True # discriminator feature matching loss 69 | __C.GAN.use_cgan = False # projection discriminator 70 | __C.GAN.weight_im = 1 # 1 71 | __C.GAN.weight_fm = 1 # 1 72 | __C.GAN.weight_l2 = 200 # 200 73 | __C.GAN.weight_gan = 0.1 # 0.1 74 | 75 | # Train 76 | __C.TRAIN = edict() 77 | __C.TRAIN.batch_size = 8 78 | __C.TRAIN.n_epochs = 150 79 | __C.TRAIN.save_freq = 5 80 | __C.TRAIN.log_freq = 1 81 | __C.TRAIN.learning_rate = 1e-4 82 | __C.TRAIN.lr_milestones = [1000] # don't schedule 83 | __C.TRAIN.gamma = 0.5 84 | __C.TRAIN.betas = (0.0, 0.9) 85 | __C.TRAIN.weight_decay = 0 86 | 87 | # Test config 88 | __C.TEST = edict() 89 | __C.TEST.mode = "default" 90 | __C.TEST.infer_freq = 25 91 | __C.TEST.metric_name = "EMD" # 'EMD' or 'ChamferDistance' 92 | 93 | # Dataset Config 94 | __C.DATASETS = edict() 95 | __C.DATASETS.shapenet = edict() 96 | __C.DATASETS.shapenet.n_renderings = 8 97 | __C.DATASETS.shapenet.n_points = 16384 98 | # 'GRnet' or 'ShapeNet' version dataset 99 | __C.DATASETS.shapenet.version = "GRnet" 100 | __C.DATASETS.shapenet.category_file_path = "./datasets/data/ShapeNet.json" 101 | __C.DATASETS.shapenet.partial_points_path = "/path/to/datasets/ShapeNetCompletion/%s/partial/%s/%s/%02d.pcd" 102 | __C.DATASETS.shapenet.complete_points_path = "/path/to/datasets/ShapeNetCompletion/%s/complete/%s/%s.pcd" 103 | __C.DATASETS.completion3d = edict() 104 | __C.DATASETS.completion3d.category_file_path = "/path/to/datasets/data/Completion3D.json" 105 | __C.DATASETS.completion3d.partial_points_path = "/path/to/datasets/completion3d/data/shapenet/%s/partial/%s/%s.h5" 106 | __C.DATASETS.completion3d.complete_points_path = "/path/to/datasets/completion3d/data/shapenet/%s/gt/%s/%s.h5" 107 | __C.DATASETS.kitti = edict() 108 | __C.DATASETS.kitti.category_file_path = "/path/to/datasets/data/KITTI.json" 109 | __C.DATASETS.kitti.partial_points_path = "/path/to/datasets/KITTI/cars/%s.pcd" 110 | __C.DATASETS.kitti.bounding_box_file_path = "/path/to/datasets/KITTI/bboxes/%s.txt" 111 | 112 | # Merge config dictionary 113 | 114 | 115 | def _merge_a_into_b(a, b): 116 | """Merge config dictionary a into config dictionary b, clobbering the options in b whenever they are also specified in a.""" 117 | if type(a) is not edict: 118 | return 119 | 120 | for k, v in a.items(): 121 | # a must specify keys that are in b 122 | if k not in b: 123 | raise KeyError("{} is not a valid config key".format(k)) 124 | 125 | # the types must match, too 126 | old_type = type(b[k]) 127 | if old_type is not type(v): 128 | if isinstance(b[k], np.ndarray): 129 | v = np.array(v, dtype=b[k].dtype) 130 | else: 131 | raise ValueError( 132 | ("Type mismatch ({} vs. {}) " "for config key: {}").format( 133 | type(b[k]), type(v), k 134 | ) 135 | ) 136 | 137 | # recursively merge dicts 138 | if type(v) is edict: 139 | try: 140 | _merge_a_into_b(a[k], b[k]) 141 | except: 142 | print(("Error under config key: {}".format(k))) 143 | raise 144 | else: 145 | b[k] = v 146 | 147 | 148 | # load yaml file 149 | def cfg_from_file(filename): 150 | """Load a config file and merge it into the default options.""" 151 | with open(filename, "r", encoding="utf-8") as f: 152 | yaml_cfg = edict(yaml.load(f)) 153 | 154 | _merge_a_into_b(yaml_cfg, __C) 155 | 156 | 157 | def cfg_update(args): 158 | """Overwrite the hyperparameters in cfg.""" 159 | # the path of model 160 | if args.weights is not None: 161 | cfg.CONST.weights = args.weights 162 | cfg.CONST.device = args.gpu_id 163 | if args.workdir is not None: 164 | cfg.DIR.out_path = args.workdir 165 | 166 | # set up folders for logs and checkpoints 167 | output_dir = os.path.join( 168 | cfg.DIR.out_path, "%s", datetime.datetime.now().isoformat() 169 | ) 170 | cfg.DIR.checkpoints = output_dir % "checkpoints" 171 | cfg.DIR.logs = output_dir % "logs" 172 | return output_dir 173 | -------------------------------------------------------------------------------- /configs/grnet.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 2 | train_dataset: ShapeNet 3 | test_dataset: ShapeNet 4 | n_outpoints: 16384 5 | 6 | CONST: 7 | num_workers: 32 8 | n_input_points: 3000 9 | 10 | NETWORK: 11 | n_sampling_points: 2048 12 | n_primitives: 32 13 | model_type: "GRNet" 14 | metric: "emd" 15 | 16 | TRAIN: 17 | learning_rate: 0.0001 18 | batch_size: 32 19 | n_epochs: 150 20 | save_freq: 5 21 | log_freq: 1 22 | 23 | TEST: 24 | metric_name: "EMD" 25 | -------------------------------------------------------------------------------- /configs/model_names.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | MODEL_MSN = "MSN" 5 | MODEL_ATLASNET = "AtlasNet" 6 | MODEL_GRNET = "GRNet" 7 | MODEL_SPARENET = "SpareNet" 8 | MODEL_PN2 = "PN2" 9 | MODEL_FCNET = "FCNet" 10 | MODEL_PCN = "PCN" 11 | MODEL_FOLDINGNET = "FoldingNet" 12 | MODEL_TOPNET = "TopNet" 13 | -------------------------------------------------------------------------------- /configs/msn.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 2 | train_dataset: ShapeNet 3 | test_dataset: ShapeNet 4 | n_outpoints: 16384 5 | 6 | CONST: 7 | num_workers: 32 8 | n_input_points: 3000 9 | 10 | NETWORK: 11 | n_sampling_points: 2048 12 | n_primitives: 32 13 | model_type: "MSN" 14 | metric: "emd" 15 | 16 | TRAIN: 17 | learning_rate: 0.0001 18 | batch_size: 32 19 | n_epochs: 150 20 | save_freq: 5 21 | log_freq: 1 22 | 23 | TEST: 24 | metric_name: "EMD" 25 | -------------------------------------------------------------------------------- /configs/sparenet.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 2 | train_dataset: ShapeNet 3 | test_dataset: ShapeNet 4 | n_outpoints: 16384 5 | 6 | CONST: 7 | num_workers: 32 8 | n_input_points: 3000 9 | 10 | NETWORK: 11 | n_sampling_points: 2048 12 | gridding_loss_scales: 13 | - 128 14 | - 64 15 | gridding_loss_alphas: 16 | - 0.1 17 | - 0.01 18 | n_primitives: 32 19 | model_type: "SpareNet" 20 | metric: "emd" 21 | encode: "Residualnet" 22 | use_adain: "share" 23 | use_selayer: true 24 | use_consist_loss: true 25 | 26 | RENDER: 27 | img_size: 256 28 | radius_list: 29 | - 5.0 30 | - 7.0 31 | - 10.0 32 | n_views: 8 33 | 34 | GAN: 35 | use_im: true 36 | use_fm: true 37 | use_cgan: true 38 | weight_gan: 0.1 39 | weight_l2: 200 40 | weight_im: 1 41 | weight_fm: 1 42 | 43 | TRAIN: 44 | learning_rate: 0.0001 45 | batch_size: 24 46 | n_epochs: 150 47 | save_freq: 5 48 | log_freq: 1 49 | 50 | TEST: 51 | metric_name: "EMD" 52 | -------------------------------------------------------------------------------- /configs/sparenet_gan.yaml: -------------------------------------------------------------------------------- 1 | DATASET: 2 | train_dataset: ShapeNet 3 | test_dataset: ShapeNet 4 | n_outpoints: 16384 5 | 6 | CONST: 7 | num_workers: 32 8 | n_input_points: 3000 9 | 10 | NETWORK: 11 | n_sampling_points: 2048 12 | gridding_loss_scales: 13 | - 128 14 | - 64 15 | gridding_loss_alphas: 16 | - 0.1 17 | - 0.01 18 | n_primitives: 32 19 | model_type: "SpareNet" 20 | metric: "emd" 21 | encode: "Residualnet" 22 | use_adain: "share" 23 | use_selayer: true 24 | use_consist_loss: true 25 | 26 | RENDER: 27 | img_size: 256 28 | radius_list: 29 | - 5.0 30 | - 7.0 31 | - 10.0 32 | n_views: 8 33 | 34 | GAN: 35 | use_im: true 36 | use_fm: true 37 | use_cgan: true 38 | weight_gan: 0.1 39 | weight_l2: 200 40 | weight_im: 1 41 | weight_fm: 1 42 | 43 | TRAIN: 44 | learning_rate: 0.0001 45 | batch_size: 32 46 | n_epochs: 150 47 | save_freq: 5 48 | log_freq: 1 49 | 50 | TEST: 51 | metric_name: "EMD" 52 | -------------------------------------------------------------------------------- /cuda/MDS/MDS.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | // NOTE: AT_CHECK has become TORCH_CHECK 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void minimum_density_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs, float *mean_mst_length); 14 | 15 | #define CUDA_CHECK_ERRORS() \ 16 | do \ 17 | { \ 18 | cudaError_t err = cudaGetLastError(); \ 19 | if (cudaSuccess != err) \ 20 | { \ 21 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 22 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 23 | __FILE__); \ 24 | exit(-1); \ 25 | } \ 26 | } while (0) 27 | 28 | #define CHECK_CUDA(x) \ 29 | do \ 30 | { \ 31 | TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 32 | } while (0) 33 | 34 | #define CHECK_CONTIGUOUS(x) \ 35 | do \ 36 | { \ 37 | TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 38 | } while (0) 39 | 40 | #define CHECK_IS_INT(x) \ 41 | do \ 42 | { \ 43 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ 44 | #x " must be an int tensor"); \ 45 | } while (0) 46 | 47 | #define CHECK_IS_FLOAT(x) \ 48 | do \ 49 | { \ 50 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ 51 | #x " must be a float tensor"); \ 52 | } while (0) 53 | 54 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) 55 | { 56 | CHECK_CONTIGUOUS(points); 57 | CHECK_CONTIGUOUS(idx); 58 | CHECK_IS_FLOAT(points); 59 | CHECK_IS_INT(idx); 60 | 61 | if (points.type().is_cuda()) 62 | { 63 | CHECK_CUDA(idx); 64 | } 65 | 66 | at::Tensor output = 67 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 68 | at::device(points.device()).dtype(at::ScalarType::Float)); 69 | 70 | if (points.type().is_cuda()) 71 | { 72 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 73 | idx.size(1), points.data(), 74 | idx.data(), output.data()); 75 | } 76 | else 77 | { 78 | TORCH_CHECK(false, "CPU not supported"); 79 | } 80 | 81 | return output; 82 | } 83 | 84 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 85 | const int n) 86 | { 87 | CHECK_CONTIGUOUS(grad_out); 88 | CHECK_CONTIGUOUS(idx); 89 | CHECK_IS_FLOAT(grad_out); 90 | CHECK_IS_INT(idx); 91 | 92 | if (grad_out.type().is_cuda()) 93 | { 94 | CHECK_CUDA(idx); 95 | } 96 | 97 | at::Tensor output = 98 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 99 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 100 | 101 | if (grad_out.type().is_cuda()) 102 | { 103 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 104 | idx.size(1), grad_out.data(), 105 | idx.data(), output.data()); 106 | } 107 | else 108 | { 109 | TORCH_CHECK(false, "CPU not supported"); 110 | } 111 | 112 | return output; 113 | } 114 | at::Tensor minimum_density_sampling(at::Tensor points, const int nsamples, at::Tensor mean_mst_length, at::Tensor output) 115 | { 116 | CHECK_CONTIGUOUS(points); 117 | CHECK_IS_FLOAT(points); 118 | 119 | at::Tensor tmp = 120 | torch::zeros({points.size(0), points.size(1)}, 121 | at::device(points.device()).dtype(at::ScalarType::Float)); 122 | 123 | if (points.type().is_cuda()) 124 | { 125 | minimum_density_sampling_kernel_wrapper( 126 | points.size(0), points.size(1), nsamples, points.data(), 127 | tmp.data(), output.data(), mean_mst_length.data()); 128 | } 129 | else 130 | { 131 | TORCH_CHECK(false, "CPU not supported"); 132 | } 133 | 134 | return output; 135 | } 136 | 137 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 138 | { 139 | m.def("minimum_density_sampling", &minimum_density_sampling, "minimum_density_sampling (CUDA)"); 140 | m.def("gather_points_grad", &gather_points_grad, "gather_points_grad (CUDA)"); 141 | m.def("gather_points", &gather_points, "gather_points (CUDA)"); 142 | } 143 | -------------------------------------------------------------------------------- /cuda/MDS/MDS.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: MDS 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/MDS/MDS.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | MDS.cpp 2 | MDS_cuda.cu 3 | setup.py 4 | MDS.egg-info/PKG-INFO 5 | MDS.egg-info/SOURCES.txt 6 | MDS.egg-info/dependency_links.txt 7 | MDS.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/MDS/MDS.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/MDS/MDS.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | MDS 2 | -------------------------------------------------------------------------------- /cuda/MDS/MDS_module.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.autograd import Function 4 | import MDS 5 | 6 | 7 | class MinimumDensitySampling(Function): 8 | @staticmethod 9 | def forward(ctx, xyz, npoint, mean_mst_length): 10 | # type: (Any, torch.Tensor, int) -> torch.Tensor 11 | r""" 12 | Uses iterative radius point sampling to select a set of npoint features that have the largest 13 | minimum distance 14 | 15 | Parameters 16 | ---------- 17 | xyz : torch.Tensor 18 | (B, N, 3) tensor where N > npoint 19 | npoint : int32 20 | number of features in the sampled set 21 | mean_mst_length : torch.Tensor 22 | (B) the average edge length from expansion penalty module 23 | 24 | Returns 25 | ------- 26 | torch.Tensor 27 | (B, npoint) tensor containing the set 28 | """ 29 | idx = torch.zeros( 30 | xyz.shape[0], npoint, requires_grad=False, device="cuda", dtype=torch.int32 31 | ).contiguous() 32 | MDS.minimum_density_sampling(xyz, npoint, mean_mst_length, idx) 33 | ctx.mark_non_differentiable(idx) 34 | return idx 35 | 36 | @staticmethod 37 | def backward(grad_idx, a=None): 38 | return None, None, None 39 | 40 | 41 | minimum_density_sample = MinimumDensitySampling.apply 42 | 43 | 44 | class GatherOperation(Function): 45 | @staticmethod 46 | def forward(ctx, features, idx): 47 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 48 | r""" 49 | 50 | Parameters 51 | ---------- 52 | features : torch.Tensor 53 | (B, C, N) tensor 54 | 55 | idx : torch.Tensor 56 | (B, npoint) tensor of the features to gather 57 | 58 | Returns 59 | ------- 60 | torch.Tensor 61 | (B, C, npoint) tensor 62 | """ 63 | 64 | _, C, N = features.size() 65 | 66 | ctx.for_backwards = (idx, C, N) 67 | 68 | return MDS.gather_points(features, idx) 69 | 70 | @staticmethod 71 | def backward(ctx, grad_out): 72 | idx, C, N = ctx.for_backwards 73 | 74 | grad_features = MDS.gather_points_grad(grad_out.contiguous(), idx, N) 75 | return grad_features, None 76 | 77 | 78 | gather_operation = GatherOperation.apply 79 | 80 | -------------------------------------------------------------------------------- /cuda/MDS/dist/MDS-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/MDS/dist/MDS-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/MDS/dist/MDS-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/MDS/dist/MDS-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/MDS/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name="MDS", 6 | ext_modules=[CUDAExtension("MDS", ["MDS_cuda.cu", "MDS.cpp",]),], 7 | cmdclass={"build_ext": BuildExtension}, 8 | ) 9 | 10 | -------------------------------------------------------------------------------- /cuda/chamfer_dist/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | import chamfer 4 | 5 | 6 | class ChamferFunction(torch.autograd.Function): 7 | @staticmethod 8 | def forward(ctx, xyz1, xyz2): 9 | dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) 10 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 11 | 12 | return dist1, dist2 13 | 14 | @staticmethod 15 | def backward(ctx, grad_dist1, grad_dist2): 16 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 17 | grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) 18 | return grad_xyz1, grad_xyz2 19 | 20 | 21 | class ChamferDistance(torch.nn.Module): 22 | def __init__(self, ignore_zeros=False): 23 | super(ChamferDistance, self).__init__() 24 | self.ignore_zeros = ignore_zeros 25 | 26 | def forward(self, xyz1, xyz2): 27 | batch_size = xyz1.size(0) 28 | if batch_size == 1 and self.ignore_zeros: 29 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 30 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 31 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 32 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 33 | 34 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 35 | return torch.mean(dist1) + torch.mean(dist2) 36 | 37 | 38 | class ChamferDistanceSeperate(torch.nn.Module): 39 | def __init__(self, ignore_zeros=False): 40 | super(ChamferDistanceSeperate, self).__init__() 41 | self.ignore_zeros = ignore_zeros 42 | 43 | def forward(self, xyz1, xyz2): 44 | batch_size = xyz1.size(0) 45 | if batch_size == 1 and self.ignore_zeros: 46 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 47 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 48 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 49 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 50 | 51 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 52 | return torch.mean(dist1), torch.mean(dist2) 53 | 54 | -------------------------------------------------------------------------------- /cuda/chamfer_dist/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-08-07 20:54:24 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2019-12-10 10:33:50 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | std::vector chamfer_cuda_forward(torch::Tensor xyz1, 13 | torch::Tensor xyz2); 14 | 15 | std::vector chamfer_cuda_backward(torch::Tensor xyz1, 16 | torch::Tensor xyz2, 17 | torch::Tensor idx1, 18 | torch::Tensor idx2, 19 | torch::Tensor grad_dist1, 20 | torch::Tensor grad_dist2); 21 | 22 | std::vector chamfer_forward(torch::Tensor xyz1, 23 | torch::Tensor xyz2) 24 | { 25 | return chamfer_cuda_forward(xyz1, xyz2); 26 | } 27 | 28 | std::vector chamfer_backward(torch::Tensor xyz1, 29 | torch::Tensor xyz2, 30 | torch::Tensor idx1, 31 | torch::Tensor idx2, 32 | torch::Tensor grad_dist1, 33 | torch::Tensor grad_dist2) 34 | { 35 | return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); 36 | } 37 | 38 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 39 | { 40 | m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); 41 | m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); 42 | } 43 | -------------------------------------------------------------------------------- /cuda/chamfer_dist/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-10 10:04:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup( 12 | name="chamfer", 13 | version="2.0.0", 14 | ext_modules=[CUDAExtension("chamfer", ["chamfer_cuda.cpp", "chamfer.cu",]),], 15 | cmdclass={"build_ext": BuildExtension}, 16 | ) 17 | -------------------------------------------------------------------------------- /cuda/chamfer_dist/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:38:01 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:21:36 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in chamfer.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.chamfer_dist import ChamferFunction 20 | 21 | 22 | class ChamferDistanceTestCase(unittest.TestCase): 23 | def test_chamfer_dist(self): 24 | x = torch.rand(4, 64, 3).double() 25 | y = torch.rand(4, 128, 3).double() 26 | x.requires_grad = True 27 | y.requires_grad = True 28 | print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) 29 | 30 | 31 | if __name__ == "__main__": 32 | unittest.main() 33 | -------------------------------------------------------------------------------- /cuda/chamfer_distance/__init__.py: -------------------------------------------------------------------------------- 1 | from .chamfer_distance import ChamferDistance, ChamferDistanceMean 2 | -------------------------------------------------------------------------------- /cuda/chamfer_distance/chamfer_distance.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | // CUDA forward declarations 4 | int ChamferDistanceKernelLauncher( 5 | const int b, const int n, 6 | const float *xyz, 7 | const int m, 8 | const float *xyz2, 9 | float *result, 10 | int *result_i, 11 | float *result2, 12 | int *result2_i); 13 | 14 | int ChamferDistanceGradKernelLauncher( 15 | const int b, const int n, 16 | const float *xyz1, 17 | const int m, 18 | const float *xyz2, 19 | const float *grad_dist1, 20 | const int *idx1, 21 | const float *grad_dist2, 22 | const int *idx2, 23 | float *grad_xyz1, 24 | float *grad_xyz2); 25 | 26 | void chamfer_distance_forward_cuda( 27 | const at::Tensor xyz1, 28 | const at::Tensor xyz2, 29 | const at::Tensor dist1, 30 | const at::Tensor dist2, 31 | const at::Tensor idx1, 32 | const at::Tensor idx2) 33 | { 34 | ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 35 | xyz2.size(1), xyz2.data(), 36 | dist1.data(), idx1.data(), 37 | dist2.data(), idx2.data()); 38 | } 39 | 40 | void chamfer_distance_backward_cuda( 41 | const at::Tensor xyz1, 42 | const at::Tensor xyz2, 43 | at::Tensor gradxyz1, 44 | at::Tensor gradxyz2, 45 | at::Tensor graddist1, 46 | at::Tensor graddist2, 47 | at::Tensor idx1, 48 | at::Tensor idx2) 49 | { 50 | ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 51 | xyz2.size(1), xyz2.data(), 52 | graddist1.data(), idx1.data(), 53 | graddist2.data(), idx2.data(), 54 | gradxyz1.data(), gradxyz2.data()); 55 | } 56 | 57 | void nnsearch( 58 | const int b, const int n, const int m, 59 | const float *xyz1, 60 | const float *xyz2, 61 | float *dist, 62 | int *idx) 63 | { 64 | for (int i = 0; i < b; i++) 65 | { 66 | for (int j = 0; j < n; j++) 67 | { 68 | const float x1 = xyz1[(i * n + j) * 3 + 0]; 69 | const float y1 = xyz1[(i * n + j) * 3 + 1]; 70 | const float z1 = xyz1[(i * n + j) * 3 + 2]; 71 | double best = 0; 72 | int besti = 0; 73 | for (int k = 0; k < m; k++) 74 | { 75 | const float x2 = xyz2[(i * m + k) * 3 + 0] - x1; 76 | const float y2 = xyz2[(i * m + k) * 3 + 1] - y1; 77 | const float z2 = xyz2[(i * m + k) * 3 + 2] - z1; 78 | const double d = x2 * x2 + y2 * y2 + z2 * z2; 79 | if (k == 0 || d < best) 80 | { 81 | best = d; 82 | besti = k; 83 | } 84 | } 85 | dist[i * n + j] = best; 86 | idx[i * n + j] = besti; 87 | } 88 | } 89 | } 90 | 91 | void chamfer_distance_forward( 92 | const at::Tensor xyz1, 93 | const at::Tensor xyz2, 94 | const at::Tensor dist1, 95 | const at::Tensor dist2, 96 | const at::Tensor idx1, 97 | const at::Tensor idx2) 98 | { 99 | const int batchsize = xyz1.size(0); 100 | const int n = xyz1.size(1); 101 | const int m = xyz2.size(1); 102 | 103 | const float *xyz1_data = xyz1.data(); 104 | const float *xyz2_data = xyz2.data(); 105 | float *dist1_data = dist1.data(); 106 | float *dist2_data = dist2.data(); 107 | int *idx1_data = idx1.data(); 108 | int *idx2_data = idx2.data(); 109 | 110 | nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); 111 | nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); 112 | } 113 | 114 | void chamfer_distance_backward( 115 | const at::Tensor xyz1, 116 | const at::Tensor xyz2, 117 | at::Tensor gradxyz1, 118 | at::Tensor gradxyz2, 119 | at::Tensor graddist1, 120 | at::Tensor graddist2, 121 | at::Tensor idx1, 122 | at::Tensor idx2) 123 | { 124 | const int b = xyz1.size(0); 125 | const int n = xyz1.size(1); 126 | const int m = xyz2.size(1); 127 | 128 | const float *xyz1_data = xyz1.data(); 129 | const float *xyz2_data = xyz2.data(); 130 | float *gradxyz1_data = gradxyz1.data(); 131 | float *gradxyz2_data = gradxyz2.data(); 132 | float *graddist1_data = graddist1.data(); 133 | float *graddist2_data = graddist2.data(); 134 | const int *idx1_data = idx1.data(); 135 | const int *idx2_data = idx2.data(); 136 | 137 | for (int i = 0; i < b * n * 3; i++) 138 | gradxyz1_data[i] = 0; 139 | for (int i = 0; i < b * m * 3; i++) 140 | gradxyz2_data[i] = 0; 141 | for (int i = 0; i < b; i++) 142 | { 143 | for (int j = 0; j < n; j++) 144 | { 145 | const float x1 = xyz1_data[(i * n + j) * 3 + 0]; 146 | const float y1 = xyz1_data[(i * n + j) * 3 + 1]; 147 | const float z1 = xyz1_data[(i * n + j) * 3 + 2]; 148 | const int j2 = idx1_data[i * n + j]; 149 | 150 | const float x2 = xyz2_data[(i * m + j2) * 3 + 0]; 151 | const float y2 = xyz2_data[(i * m + j2) * 3 + 1]; 152 | const float z2 = xyz2_data[(i * m + j2) * 3 + 2]; 153 | const float g = graddist1_data[i * n + j] * 2; 154 | 155 | gradxyz1_data[(i * n + j) * 3 + 0] += g * (x1 - x2); 156 | gradxyz1_data[(i * n + j) * 3 + 1] += g * (y1 - y2); 157 | gradxyz1_data[(i * n + j) * 3 + 2] += g * (z1 - z2); 158 | gradxyz2_data[(i * m + j2) * 3 + 0] -= (g * (x1 - x2)); 159 | gradxyz2_data[(i * m + j2) * 3 + 1] -= (g * (y1 - y2)); 160 | gradxyz2_data[(i * m + j2) * 3 + 2] -= (g * (z1 - z2)); 161 | } 162 | for (int j = 0; j < m; j++) 163 | { 164 | const float x1 = xyz2_data[(i * m + j) * 3 + 0]; 165 | const float y1 = xyz2_data[(i * m + j) * 3 + 1]; 166 | const float z1 = xyz2_data[(i * m + j) * 3 + 2]; 167 | const int j2 = idx2_data[i * m + j]; 168 | const float x2 = xyz1_data[(i * n + j2) * 3 + 0]; 169 | const float y2 = xyz1_data[(i * n + j2) * 3 + 1]; 170 | const float z2 = xyz1_data[(i * n + j2) * 3 + 2]; 171 | const float g = graddist2_data[i * m + j] * 2; 172 | gradxyz2_data[(i * m + j) * 3 + 0] += g * (x1 - x2); 173 | gradxyz2_data[(i * m + j) * 3 + 1] += g * (y1 - y2); 174 | gradxyz2_data[(i * m + j) * 3 + 2] += g * (z1 - z2); 175 | gradxyz1_data[(i * n + j2) * 3 + 0] -= (g * (x1 - x2)); 176 | gradxyz1_data[(i * n + j2) * 3 + 1] -= (g * (y1 - y2)); 177 | gradxyz1_data[(i * n + j2) * 3 + 2] -= (g * (z1 - z2)); 178 | } 179 | } 180 | } 181 | 182 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 183 | { 184 | m.def("forward", &chamfer_distance_forward, "ChamferDistance forward"); 185 | m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)"); 186 | m.def("backward", &chamfer_distance_backward, "ChamferDistance backward"); 187 | m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)"); 188 | } 189 | -------------------------------------------------------------------------------- /cuda/chamfer_distance/chamfer_distance.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | __global__ 7 | void ChamferDistanceKernel( 8 | int b, 9 | int n, 10 | const float* xyz, 11 | int m, 12 | const float* xyz2, 13 | float* result, 14 | int* result_i) 15 | { 16 | const int batch=512; 17 | __shared__ float buf[batch*3]; 18 | for (int i=blockIdx.x;ibest){ 130 | result[(i*n+j)]=best; 131 | result_i[(i*n+j)]=best_i; 132 | } 133 | } 134 | __syncthreads(); 135 | } 136 | } 137 | } 138 | 139 | void ChamferDistanceKernelLauncher( 140 | const int b, const int n, 141 | const float* xyz, 142 | const int m, 143 | const float* xyz2, 144 | float* result, 145 | int* result_i, 146 | float* result2, 147 | int* result2_i) 148 | { 149 | ChamferDistanceKernel<<>>(b, n, xyz, m, xyz2, result, result_i); 150 | ChamferDistanceKernel<<>>(b, m, xyz2, n, xyz, result2, result2_i); 151 | 152 | cudaError_t err = cudaGetLastError(); 153 | if (err != cudaSuccess) 154 | printf("error in chamfer distance updateOutput: %s\n", cudaGetErrorString(err)); 155 | } 156 | 157 | 158 | __global__ 159 | void ChamferDistanceGradKernel( 160 | int b, int n, 161 | const float* xyz1, 162 | int m, 163 | const float* xyz2, 164 | const float* grad_dist1, 165 | const int* idx1, 166 | float* grad_xyz1, 167 | float* grad_xyz2) 168 | { 169 | for (int i = blockIdx.x; i>>(b, n, xyz1, m, xyz2, grad_dist1, idx1, grad_xyz1, grad_xyz2); 204 | ChamferDistanceGradKernel<<>>(b, m, xyz2, n, xyz1, grad_dist2, idx2, grad_xyz2, grad_xyz1); 205 | 206 | cudaError_t err = cudaGetLastError(); 207 | if (err != cudaSuccess) 208 | printf("error in chamfer distance get grad: %s\n", cudaGetErrorString(err)); 209 | } 210 | -------------------------------------------------------------------------------- /cuda/chamfer_distance/chamfer_distance.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from torch.utils.cpp_extension import load 4 | import os 5 | 6 | module_path = os.path.dirname(os.path.abspath(__file__)) 7 | print(module_path) 8 | cd = load( 9 | name="cd", 10 | verbose=True, 11 | sources=[ 12 | os.path.join(module_path, "chamfer_distance.cpp"), 13 | os.path.join(module_path, "chamfer_distance.cu"), 14 | ], 15 | ) 16 | 17 | 18 | class ChamferDistanceFunction(torch.autograd.Function): 19 | @staticmethod 20 | def forward(ctx, xyz1, xyz2): 21 | batchsize, n, _ = xyz1.size() 22 | _, m, _ = xyz2.size() 23 | xyz1 = xyz1.contiguous() 24 | xyz2 = xyz2.contiguous() 25 | dist1 = torch.zeros(batchsize, n) 26 | dist2 = torch.zeros(batchsize, m) 27 | 28 | idx1 = torch.zeros(batchsize, n, dtype=torch.int) 29 | idx2 = torch.zeros(batchsize, m, dtype=torch.int) 30 | 31 | if not xyz1.is_cuda: 32 | cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 33 | else: 34 | dist1 = dist1.cuda() 35 | dist2 = dist2.cuda() 36 | idx1 = idx1.cuda() 37 | idx2 = idx2.cuda() 38 | cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2) 39 | 40 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 41 | 42 | return dist1, dist2 43 | 44 | @staticmethod 45 | def backward(ctx, graddist1, graddist2): 46 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 47 | 48 | graddist1 = graddist1.contiguous() 49 | graddist2 = graddist2.contiguous() 50 | 51 | gradxyz1 = torch.zeros(xyz1.size()) 52 | gradxyz2 = torch.zeros(xyz2.size()) 53 | 54 | if not graddist1.is_cuda: 55 | cd.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 56 | else: 57 | gradxyz1 = gradxyz1.cuda() 58 | gradxyz2 = gradxyz2.cuda() 59 | cd.backward_cuda(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 60 | 61 | return gradxyz1, gradxyz2 62 | 63 | 64 | class ChamferDistance(torch.nn.Module): 65 | def forward(self, xyz1, xyz2): 66 | return ChamferDistanceFunction.apply(xyz1, xyz2) 67 | 68 | 69 | class ChamferDistanceMean(torch.nn.Module): 70 | def forward(self, xyz1, xyz2): 71 | dist1, dist2 = ChamferDistanceFunction.apply(xyz1, xyz2) 72 | return (torch.mean(dist1)) + (torch.mean(dist2)) 73 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 16:55:15 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:15:14 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import cubic_feature_sampling 11 | 12 | 13 | class CubicFeatureSamplingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, ptcloud, cubic_features, neighborhood_size=1): 16 | scale = cubic_features.size(2) 17 | point_features, grid_pt_indexes = cubic_feature_sampling.forward( 18 | scale, neighborhood_size, ptcloud, cubic_features 19 | ) 20 | ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes) 21 | return point_features 22 | 23 | @staticmethod 24 | def backward(ctx, grad_point_features): 25 | scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors 26 | scale = int(scale.item()) 27 | neighborhood_size = int(neighborhood_size.item()) 28 | grad_point_features = grad_point_features.contiguous() 29 | grad_ptcloud, grad_cubic_features = cubic_feature_sampling.backward( 30 | scale, neighborhood_size, grad_point_features, grid_pt_indexes 31 | ) 32 | return grad_ptcloud, grad_cubic_features, None 33 | 34 | 35 | class CubicFeatureSampling(torch.nn.Module): 36 | def __init__(self): 37 | super(CubicFeatureSampling, self).__init__() 38 | 39 | def forward(self, ptcloud, cubic_features, neighborhood_size=1): 40 | h_scale = cubic_features.size(2) / 2 41 | ptcloud = ptcloud * h_scale + h_scale 42 | return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size) 43 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-19 20:36:36 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:55:41 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #define CUDA_NUM_THREADS 512 15 | 16 | // Computer the number of threads needed in GPU 17 | inline int get_n_threads(int n) { 18 | const int pow_2 = std::log(static_cast(n)) / std::log(2.0); 19 | return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); 20 | } 21 | 22 | __device__ int compute_index(int offset_x, 23 | int offset_y, 24 | int offset_z, 25 | int scale) { 26 | return offset_x * scale * scale + offset_y * scale + offset_z; 27 | } 28 | 29 | __global__ void cubic_feature_sampling_kernel( 30 | int scale, 31 | int neighborhood_size, 32 | int n_vertices, 33 | int n_pts, 34 | int n_cubic_channels, 35 | const float *__restrict__ ptcloud, 36 | const float *__restrict__ cubic_features, 37 | float *__restrict__ point_features, 38 | int *__restrict__ grid_pt_indexes) { 39 | int batch_index = blockIdx.x; 40 | int index = threadIdx.x; 41 | int stride = blockDim.x; 42 | int cub_scale = scale * scale * scale; 43 | 44 | ptcloud += batch_index * n_pts * 3; 45 | cubic_features += batch_index * n_cubic_channels * cub_scale; 46 | point_features += batch_index * n_pts * n_vertices * n_cubic_channels; 47 | grid_pt_indexes += batch_index * n_pts * n_vertices; 48 | 49 | for (int i = index; i < n_pts; i += stride) { 50 | float pt_x = ptcloud[i * 3 + 0]; 51 | float pt_y = ptcloud[i * 3 + 1]; 52 | float pt_z = ptcloud[i * 3 + 2]; 53 | 54 | int lower_x = std::floor(pt_x); 55 | int upper_x = std::ceil(pt_x); 56 | if (lower_x == upper_x) { 57 | upper_x += 1; 58 | } 59 | int lower_y = std::floor(pt_y); 60 | int upper_y = std::ceil(pt_y); 61 | if (lower_y == upper_y) { 62 | upper_y += 1; 63 | } 64 | int lower_z = std::floor(pt_z); 65 | int upper_z = std::ceil(pt_z); 66 | if (lower_z == upper_z) { 67 | upper_z += 1; 68 | } 69 | 70 | int ns = neighborhood_size - 1; 71 | int vertex_idx = 0; 72 | for (int j = lower_x - ns; j <= upper_x + ns; ++j) { 73 | for (int k = lower_y - ns; k <= upper_y + ns; ++k) { 74 | for (int m = lower_z - ns; m <= upper_z + ns; ++m) { 75 | if (j < 0 || j >= scale || k < 0 || k >= scale || m < 0 || 76 | m >= scale) { 77 | // Ignore points lies out of the grid 78 | grid_pt_indexes[i * n_vertices + vertex_idx++] = -1; 79 | } else { 80 | // Calcuating indexes for adjacent vertices 81 | grid_pt_indexes[i * n_vertices + vertex_idx++] = 82 | compute_index(j, k, m, scale); 83 | } 84 | } 85 | } 86 | } 87 | 88 | // Gather Features 89 | for (int j = 0; j < n_vertices; ++j) { 90 | for (int k = 0; k < n_cubic_channels; ++k) { 91 | int vertex_idx = grid_pt_indexes[i * n_vertices + j]; 92 | if (vertex_idx == -1) { 93 | continue; 94 | } 95 | int feature_idx = 96 | i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; 97 | float feature_val = cubic_features[k * cub_scale + vertex_idx]; 98 | point_features[feature_idx] = feature_val; 99 | } 100 | } 101 | } 102 | } 103 | 104 | std::vector cubic_feature_sampling_cuda_forward( 105 | int scale, 106 | int neighborhood_size, 107 | torch::Tensor ptcloud, 108 | torch::Tensor cubic_features, 109 | cudaStream_t stream) { 110 | int batch_size = ptcloud.size(0); 111 | int n_pts = ptcloud.size(1); 112 | int n_cubic_channels = cubic_features.size(1); 113 | 114 | int n_vertices = std::pow(neighborhood_size * 2, 3); 115 | torch::Tensor point_features = 116 | torch::zeros({batch_size, n_pts, n_vertices, n_cubic_channels}, 117 | torch::CUDA(torch::kFloat)); 118 | torch::Tensor grid_pt_indexes = 119 | torch::zeros({batch_size, n_pts, n_vertices}, torch::CUDA(torch::kInt)); 120 | 121 | cubic_feature_sampling_kernel<<>>( 123 | scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, 124 | ptcloud.data_ptr(), cubic_features.data_ptr(), 125 | point_features.data_ptr(), grid_pt_indexes.data_ptr()); 126 | 127 | cudaError_t err = cudaGetLastError(); 128 | if (err != cudaSuccess) { 129 | printf("Error in cubic_feature_sampling_cuda_forward: %s\n", 130 | cudaGetErrorString(err)); 131 | } 132 | return {point_features, grid_pt_indexes}; 133 | } 134 | 135 | __global__ void cubic_feature_sampling_grad_kernel( 136 | int scale, 137 | int neighborhood_size, 138 | int n_vertices, 139 | int n_pts, 140 | int n_cubic_channels, 141 | const float *__restrict__ grad_point_features, 142 | const int *__restrict__ grid_pt_indexes, 143 | float *__restrict__ grad_ptcloud, 144 | float *__restrict__ grad_cubic_features) { 145 | int batch_index = blockIdx.x; 146 | int index = threadIdx.x; 147 | int stride = blockDim.x; 148 | int cub_scale = scale * scale * scale; 149 | 150 | grad_point_features += batch_index * n_pts * n_vertices * n_cubic_channels; 151 | grid_pt_indexes += batch_index * n_pts * n_vertices; 152 | grad_ptcloud += batch_index * n_pts * 3; 153 | grad_cubic_features += batch_index * n_cubic_channels * cub_scale; 154 | 155 | for (int i = index; i < n_pts; i += stride) { 156 | for (int j = 0; j < n_vertices; ++j) { 157 | int vertex_idx = grid_pt_indexes[i * n_vertices + j]; 158 | if (vertex_idx == -1) { 159 | continue; 160 | } 161 | for (int k = 0; k < n_cubic_channels; ++k) { 162 | int grad_idx = 163 | i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; 164 | float grad_val = grad_point_features[grad_idx]; 165 | // Fix bugs: the gradients of ceil and floor functions are zeros. 166 | // Ref: https://github.com/tensorflow/tensorflow/issues/897 167 | // atomicAdd(&(grad_ptcloud[i * 3 + 0]), grad_val); 168 | // atomicAdd(&(grad_ptcloud[i * 3 + 1]), grad_val); 169 | // atomicAdd(&(grad_ptcloud[i * 3 + 2]), grad_val); 170 | atomicAdd(&(grad_cubic_features[k * cub_scale + vertex_idx]), grad_val); 171 | } 172 | } 173 | } 174 | } 175 | 176 | std::vector cubic_feature_sampling_cuda_backward( 177 | int scale, 178 | int neighborhood_size, 179 | torch::Tensor grad_point_features, 180 | torch::Tensor grid_pt_indexes, 181 | cudaStream_t stream) { 182 | int batch_size = grad_point_features.size(0); 183 | int n_cubic_channels = grad_point_features.size(3); 184 | int n_pts = grid_pt_indexes.size(1); 185 | int n_vertices = std::pow(neighborhood_size * 2, 3); 186 | 187 | torch::Tensor grad_ptcloud = 188 | torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); 189 | torch::Tensor grad_cubic_features = 190 | torch::zeros({batch_size, n_cubic_channels, scale, scale, scale}, 191 | torch::CUDA(torch::kFloat)); 192 | 193 | cubic_feature_sampling_grad_kernel<<>>( 195 | scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, 196 | grad_point_features.data_ptr(), grid_pt_indexes.data_ptr(), 197 | grad_ptcloud.data_ptr(), grad_cubic_features.data_ptr()); 198 | 199 | cudaError_t err = cudaGetLastError(); 200 | if (err != cudaSuccess) { 201 | printf("Error in cubic_feature_sampling_cuda_backward: %s\n", 202 | cudaGetErrorString(err)); 203 | } 204 | return {grad_ptcloud, grad_cubic_features}; 205 | } -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: cubic-feature-sampling 3 | Version: 1.1.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | cubic_feature_sampling.cu 2 | cubic_feature_sampling_cuda.cpp 3 | setup.py 4 | cubic_feature_sampling.egg-info/PKG-INFO 5 | cubic_feature_sampling.egg-info/SOURCES.txt 6 | cubic_feature_sampling.egg-info/dependency_links.txt 7 | cubic_feature_sampling.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | cubic_feature_sampling 2 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-19 17:04:38 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:50:22 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector cubic_feature_sampling_cuda_forward( 21 | int scale, 22 | int neighborhood_size, 23 | torch::Tensor ptcloud, 24 | torch::Tensor cubic_features, 25 | cudaStream_t stream); 26 | 27 | std::vector cubic_feature_sampling_cuda_backward( 28 | int scale, 29 | int neighborhood_size, 30 | torch::Tensor grad_point_features, 31 | torch::Tensor grid_pt_indexes, 32 | cudaStream_t stream); 33 | 34 | std::vector cubic_feature_sampling_forward( 35 | int scale, 36 | int neighborhood_size, 37 | torch::Tensor ptcloud, 38 | torch::Tensor cubic_features) 39 | { 40 | CHECK_INPUT(ptcloud); 41 | CHECK_INPUT(cubic_features); 42 | 43 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 44 | return cubic_feature_sampling_cuda_forward(scale, neighborhood_size, ptcloud, 45 | cubic_features, stream); 46 | } 47 | 48 | std::vector cubic_feature_sampling_backward( 49 | int scale, 50 | int neighborhood_size, 51 | torch::Tensor grad_point_features, 52 | torch::Tensor grid_pt_indexes) 53 | { 54 | CHECK_INPUT(grad_point_features); 55 | CHECK_INPUT(grid_pt_indexes); 56 | 57 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 58 | return cubic_feature_sampling_cuda_backward( 59 | scale, neighborhood_size, grad_point_features, grid_pt_indexes, stream); 60 | } 61 | 62 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 63 | { 64 | m.def("forward", &cubic_feature_sampling_forward, 65 | "Cubic Feature Sampling forward (CUDA)"); 66 | m.def("backward", &cubic_feature_sampling_backward, 67 | "Cubic Feature Sampling backward (CUDA)"); 68 | } -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 17:03:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:02:06 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup( 12 | name="cubic_feature_sampling", 13 | version="1.1.0", 14 | ext_modules=[ 15 | CUDAExtension( 16 | "cubic_feature_sampling", ["cubic_feature_sampling_cuda.cpp", "cubic_feature_sampling.cu"] 17 | ), 18 | ], 19 | cmdclass={"build_ext": BuildExtension}, 20 | ) 21 | -------------------------------------------------------------------------------- /cuda/cubic_feature_sampling/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-20 11:50:50 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:52:33 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in cubic_feature_sampling.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.cubic_feature_sampling import CubicFeatureSamplingFunction 20 | 21 | 22 | class CubicFeatureSamplingTestCase(unittest.TestCase): 23 | def test_neighborhood_size_1(self): 24 | ptcloud = torch.rand(2, 64, 3) * 2 - 1 25 | cubic_features = torch.rand(2, 4, 8, 8, 8) 26 | ptcloud.requires_grad = True 27 | cubic_features.requires_grad = True 28 | self.assertTrue( 29 | gradcheck( 30 | CubicFeatureSamplingFunction.apply, [ptcloud.double().cuda(), cubic_features.double().cuda()] 31 | ) 32 | ) 33 | 34 | def test_neighborhood_size_2(self): 35 | ptcloud = torch.rand(2, 32, 3) * 2 - 1 36 | cubic_features = torch.rand(2, 2, 8, 8, 8) 37 | ptcloud.requires_grad = True 38 | cubic_features.requires_grad = True 39 | self.assertTrue( 40 | gradcheck( 41 | CubicFeatureSamplingFunction.apply, 42 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 2], 43 | ) 44 | ) 45 | 46 | def test_neighborhood_size_3(self): 47 | ptcloud = torch.rand(1, 32, 3) * 2 - 1 48 | cubic_features = torch.rand(1, 2, 16, 16, 16) 49 | ptcloud.requires_grad = True 50 | cubic_features.requires_grad = True 51 | self.assertTrue( 52 | gradcheck( 53 | CubicFeatureSamplingFunction.apply, 54 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 3], 55 | ) 56 | ) 57 | 58 | 59 | if __name__ == "__main__": 60 | unittest.main() 61 | -------------------------------------------------------------------------------- /cuda/emd/CDEMD.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/emd/CDEMD.png -------------------------------------------------------------------------------- /cuda/emd/README.md: -------------------------------------------------------------------------------- 1 | ## Earth Mover's Distance of point clouds 2 | 3 | ![](/emd/CDEMD.png) 4 | 5 | Compared to the Chamfer Distance (CD), the Earth Mover's Distance (EMD) is more reliable to distinguish the visual quality of the point clouds. See our [paper](http://cseweb.ucsd.edu/~mil070/projects/AAAI2020/paper.pdf) for more details. 6 | 7 | We provide an EMD implementation for point cloud comparison, which only needs $O(n)$ memory and thus enables dense point clouds (with 10,000 points or over) and large batch size. It is based on an approximated algorithm (auction algorithm) and cannot guarantee a (but near) bijection assignment. It employs a parameter $\epsilon$ to balance the error rate and the speed of convergence. Smaller $\epsilon$ achieves more accurate results, but needs a longer time for convergence. The time complexity is $O(n^2k)$, where $k$ is the number of iterations. We set a $\epsilon = 0.005, k = 50$ during training and a $\epsilon = 0.002, k = 10000$ during testing. 8 | 9 | ### Compile 10 | 11 | Run `python3 setup.py install` to compile. 12 | 13 | ### Example 14 | 15 | See `emd_module.py/test_emd()` for examples. 16 | 17 | ### Input 18 | 19 | - **xyz1, xyz2**: float tensors with shape `[#batch, #points, 3]`. xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud. Two point clouds should have same size and be normalized to [0, 1]. The number of points should be a multiple of 1024. The batch size should be no greater than 512. Since we only calculate gradients for xyz1, please do not swap xyz1 and xyz2. 20 | - **eps**: a float tensor, the parameter balances the error rate and the speed of convergence. 21 | - **iters**: a int tensor, the number of iterations. 22 | 23 | ### Output 24 | 25 | - **dist**: a float tensor with shape `[#batch, #points]`. sqrt(dist) are the L2 distances between the pairs of points. 26 | - **assignment**: a int tensor with shape `[#batch, #points]`. The index of the matched point in the ground truth point cloud. 27 | -------------------------------------------------------------------------------- /cuda/emd/dist/emd-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/emd/dist/emd-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/emd/dist/emd-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/emd/dist/emd-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/emd/dist/emd-0.0.0-py3.8-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/emd/dist/emd-0.0.0-py3.8-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/emd/emd.cpp: -------------------------------------------------------------------------------- 1 | // EMD approximation module (based on auction algorithm) 2 | // author: Minghua Liu 3 | #include 4 | #include 5 | 6 | int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 7 | at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, 8 | at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters); 9 | 10 | int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx); 11 | 12 | int emd_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 13 | at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, 14 | at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) 15 | { 16 | return emd_cuda_forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters); 17 | } 18 | 19 | int emd_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) 20 | { 21 | 22 | return emd_cuda_backward(xyz1, xyz2, gradxyz, graddist, idx); 23 | } 24 | 25 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 26 | { 27 | m.def("forward", &emd_forward, "emd forward (CUDA)"); 28 | m.def("backward", &emd_backward, "emd backward (CUDA)"); 29 | } -------------------------------------------------------------------------------- /cuda/emd/emd.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: emd 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/emd/emd.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | README.md 2 | emd.cpp 3 | emd_cuda.cu 4 | setup.py 5 | emd.egg-info/PKG-INFO 6 | emd.egg-info/SOURCES.txt 7 | emd.egg-info/dependency_links.txt 8 | emd.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/emd/emd.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/emd/emd.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | emd 2 | -------------------------------------------------------------------------------- /cuda/emd/emd_module.py: -------------------------------------------------------------------------------- 1 | # EMD approximation module (based on auction algorithm) 2 | # memory complexity: O(n) 3 | # time complexity: O(n^2 * iter) 4 | # author: Minghua Liu 5 | 6 | # Input: 7 | # xyz1, xyz2: [#batch, #points, 3] 8 | # where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud 9 | # two point clouds should have same size and be normalized to [0, 1] 10 | # #points should be a multiple of 1024 11 | # #batch should be no greater than 512 12 | # eps is a parameter which balances the error rate and the speed of convergence 13 | # iters is the number of iteration 14 | # we only calculate gradient for xyz1 15 | 16 | # Output: 17 | # dist: [#batch, #points], sqrt(dist) -> L2 distance 18 | # assignment: [#batch, #points], index of the matched point in the ground truth point cloud 19 | # the result is an approximation and the assignment is not guranteed to be a bijection 20 | 21 | import time 22 | import numpy as np 23 | import torch 24 | from torch import nn 25 | from torch.autograd import Function 26 | import emd 27 | 28 | 29 | class emdFunction(Function): 30 | @staticmethod 31 | def forward(ctx, xyz1, xyz2, eps, iters): 32 | 33 | batchsize, n, _ = xyz1.size() 34 | _, m, _ = xyz2.size() 35 | 36 | assert n == m 37 | assert xyz1.size()[0] == xyz2.size()[0] 38 | assert n % 1024 == 0 39 | assert batchsize <= 512 40 | 41 | xyz1 = xyz1.contiguous().float().cuda() 42 | xyz2 = xyz2.contiguous().float().cuda() 43 | dist = torch.zeros(batchsize, n, device="cuda").contiguous() 44 | assignment = torch.zeros(batchsize, n, device="cuda", dtype=torch.int32).contiguous() - 1 45 | assignment_inv = torch.zeros(batchsize, m, device="cuda", dtype=torch.int32).contiguous() - 1 46 | price = torch.zeros(batchsize, m, device="cuda").contiguous() 47 | bid = torch.zeros(batchsize, n, device="cuda", dtype=torch.int32).contiguous() 48 | bid_increments = torch.zeros(batchsize, n, device="cuda").contiguous() 49 | max_increments = torch.zeros(batchsize, m, device="cuda").contiguous() 50 | unass_idx = torch.zeros(batchsize * n, device="cuda", dtype=torch.int32).contiguous() 51 | max_idx = torch.zeros(batchsize * m, device="cuda", dtype=torch.int32).contiguous() 52 | unass_cnt = torch.zeros(512, dtype=torch.int32, device="cuda").contiguous() 53 | unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device="cuda").contiguous() 54 | cnt_tmp = torch.zeros(512, dtype=torch.int32, device="cuda").contiguous() 55 | 56 | emd.forward( 57 | xyz1, 58 | xyz2, 59 | dist, 60 | assignment, 61 | price, 62 | assignment_inv, 63 | bid, 64 | bid_increments, 65 | max_increments, 66 | unass_idx, 67 | unass_cnt, 68 | unass_cnt_sum, 69 | cnt_tmp, 70 | max_idx, 71 | eps, 72 | iters, 73 | ) 74 | 75 | ctx.save_for_backward(xyz1, xyz2, assignment) 76 | return dist, assignment 77 | 78 | @staticmethod 79 | def backward(ctx, graddist, gradidx): 80 | xyz1, xyz2, assignment = ctx.saved_tensors 81 | graddist = graddist.contiguous() 82 | 83 | gradxyz1 = torch.zeros(xyz1.size(), device="cuda").contiguous() 84 | gradxyz2 = torch.zeros(xyz2.size(), device="cuda").contiguous() 85 | 86 | emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment) 87 | return gradxyz1, gradxyz2, None, None 88 | 89 | 90 | class emdModule(nn.Module): 91 | def __init__(self): 92 | super(emdModule, self).__init__() 93 | 94 | def forward(self, input1, input2, eps, iters): 95 | return emdFunction.apply(input1, input2, eps, iters) 96 | 97 | 98 | def test_emd(): 99 | x1 = torch.rand(20, 2048, 3).cuda() 100 | x2 = torch.rand(20, 2048, 3).cuda() 101 | emd = emdModule() 102 | start_time = time.perf_counter() 103 | dis, assigment = emd(x1, x2, 0.05, 3000) 104 | 105 | mean1 = torch.sqrt(dis).mean(1) 106 | mean0 = torch.sqrt(dis).mean(0) 107 | print("dis", dis, dis.shape) # torch.Size([bs, num_points]) 108 | print("mean1", mean1, mean1.shape) # torch.Size([bs]) 109 | print("mean0", mean0, mean0.shape) # torch.Size([num_points]) -- wrong 110 | print("Input_size: ", x1.shape) 111 | print("Runtime: %lfs" % (time.perf_counter() - start_time)) 112 | print("EMD: %lf" % np.sqrt(dis.cpu()).mean()) 113 | print("|set(assignment)|: %d" % assigment.unique().numel()) 114 | assigment = assigment.cpu().numpy() 115 | assigment = np.expand_dims(assigment, -1) 116 | x2 = np.take_along_axis(x2, assigment, axis=1) 117 | d = (x1 - x2) * (x1 - x2) 118 | print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean()) 119 | 120 | 121 | # test_emd() 122 | 123 | -------------------------------------------------------------------------------- /cuda/emd/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name="emd", 6 | ext_modules=[CUDAExtension("emd", ["emd.cpp", "emd_cuda.cu",]),], 7 | cmdclass={"build_ext": BuildExtension}, 8 | ) 9 | 10 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.8-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/expansion_penalty/dist/expansion_penalty-0.0.0-py3.8-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | int expansion_penalty_cuda_forward(at::Tensor xyz, int primitive_size, at::Tensor father, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length); 5 | 6 | int expansion_penalty_cuda_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx); 7 | 8 | int expansion_penalty_forward(at::Tensor xyz, int primitive_size, at::Tensor father, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length) 9 | { 10 | return expansion_penalty_cuda_forward(xyz, primitive_size, father, dist, alpha, neighbor, cost, mean_mst_length); 11 | } 12 | 13 | int expansion_penalty_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) 14 | { 15 | 16 | return expansion_penalty_cuda_backward(xyz, gradxyz, graddist, idx); 17 | } 18 | 19 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 20 | { 21 | m.def("forward", &expansion_penalty_forward, "expansion_penalty forward (CUDA)"); 22 | m.def("backward", &expansion_penalty_backward, "expansion_penalty backward (CUDA)"); 23 | } -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: expansion-penalty 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | expansion_penalty.cpp 2 | expansion_penalty_cuda.cu 3 | setup.py 4 | expansion_penalty.egg-info/PKG-INFO 5 | expansion_penalty.egg-info/SOURCES.txt 6 | expansion_penalty.egg-info/dependency_links.txt 7 | expansion_penalty.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | expansion_penalty 2 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | 7 | __global__ void calc_penalty(int b, int n, int primitive_size, const float * xyz, int * idx, float * dist, float alpha, int * neighbor, float * cost, float * mean_mst_length) { 8 | const int batch = 512; // primitive_size should be less than 512 9 | __shared__ float xyz_buf[batch * 3]; 10 | __shared__ bool vis[batch]; 11 | __shared__ float cur_dis[batch]; 12 | __shared__ int cur_idx[batch]; 13 | __shared__ float min_dis[batch]; 14 | __shared__ int min_idx[batch]; 15 | __shared__ float sum_dis[batch]; 16 | __shared__ int cnt[batch]; 17 | __shared__ int degree[batch]; 18 | 19 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 20 | vis[threadIdx.x] = false; 21 | cur_dis[threadIdx.x] = 1e9; 22 | cnt[threadIdx.x] = 0; 23 | degree[threadIdx.x] = 0; 24 | 25 | for (int j = threadIdx.x; j < primitive_size * 3; j += blockDim.x) { 26 | xyz_buf[j] = xyz[(i * n + blockIdx.y * primitive_size) * 3 + j]; 27 | } 28 | __syncthreads(); 29 | 30 | __shared__ int last; 31 | __shared__ float x_last; 32 | __shared__ float y_last; 33 | __shared__ float z_last; 34 | 35 | if (threadIdx.x == 0) { 36 | vis[0] = true; 37 | sum_dis[0] = 0; 38 | last = 0; 39 | x_last = xyz_buf[0]; 40 | y_last = xyz_buf[1]; 41 | z_last = xyz_buf[2]; 42 | } 43 | __syncthreads(); 44 | 45 | for (int j = 0; j < primitive_size - 1; j++) { 46 | if (vis[threadIdx.x] == false) { 47 | float delta_x = xyz_buf[threadIdx.x * 3 + 0] - x_last; 48 | float delta_y = xyz_buf[threadIdx.x * 3 + 1] - y_last; 49 | float delta_z = xyz_buf[threadIdx.x * 3 + 2] - z_last; 50 | float d = sqrtf(delta_x * delta_x + delta_y * delta_y + delta_z * delta_z); 51 | 52 | if (d < cur_dis[threadIdx.x]) { 53 | cur_dis[threadIdx.x] = d; 54 | cur_idx[threadIdx.x] = last; 55 | } 56 | min_dis[threadIdx.x] = cur_dis[threadIdx.x]; 57 | } 58 | else { 59 | min_dis[threadIdx.x] = 1e9; 60 | } 61 | min_idx[threadIdx.x] = threadIdx.x; 62 | __syncthreads(); 63 | 64 | int stride = 1; 65 | while(stride <= primitive_size / 2) { 66 | int index = (threadIdx.x + 1) * stride * 2 - 1; 67 | if(index < primitive_size && min_dis[index - stride] < min_dis[index]) { 68 | min_dis[index] = min_dis[index - stride]; 69 | min_idx[index] = min_idx[index - stride]; 70 | } 71 | stride = stride * 2; 72 | __syncthreads(); 73 | } 74 | __syncthreads(); 75 | 76 | if (threadIdx.x == primitive_size - 1) { 77 | last = min_idx[threadIdx.x]; 78 | int u = cur_idx[last]; 79 | vis[last] = true; 80 | x_last = xyz_buf[last * 3 + 0]; 81 | y_last = xyz_buf[last * 3 + 1]; 82 | z_last = xyz_buf[last * 3 + 2]; 83 | 84 | cnt[last] += 1; 85 | degree[last] += 1; 86 | neighbor[(i * n + blockIdx.y * primitive_size + last) * 512 + cnt[last]] = u; 87 | cost[(i * n + blockIdx.y * primitive_size + last) * 512 + cnt[last]] = cur_dis[last]; 88 | cnt[u] += 1; 89 | degree[u] += 1; 90 | neighbor[(i * n + blockIdx.y * primitive_size + u) * 512 + cnt[u]] = last; 91 | cost[(i * n + blockIdx.y * primitive_size + u) * 512 + cnt[u]] = cur_dis[last]; 92 | 93 | if (cnt[last] >= 512 || cnt[u] >= 512) { 94 | printf("MST Error: Too many neighbors! %d %d %d %d\n", cnt[last], cnt[u], last, u); 95 | } 96 | 97 | sum_dis[last] = cur_dis[last]; 98 | } 99 | __syncthreads(); 100 | } 101 | 102 | __syncthreads(); 103 | int stride = 1; 104 | while(stride <= primitive_size / 2) { 105 | int index = (threadIdx.x + 1) * stride * 2 - 1; 106 | if (index < primitive_size) 107 | sum_dis[index] += sum_dis[index - stride]; 108 | stride = stride * 2; 109 | __syncthreads(); 110 | } 111 | __syncthreads(); 112 | 113 | __shared__ float mean_dis; 114 | if (threadIdx.x == 0) { 115 | mean_dis = sum_dis[primitive_size - 1] / (primitive_size - 1); 116 | atomicAdd(&mean_mst_length[i], mean_dis); 117 | } 118 | 119 | dist[i * n + blockIdx.y * primitive_size + threadIdx.x] = 0; 120 | idx[i * n + blockIdx.y * primitive_size + threadIdx.x] = -1; 121 | __syncthreads(); 122 | 123 | while (true) { 124 | __shared__ int flag; 125 | flag = 0; 126 | int tmp = cnt[threadIdx.x]; 127 | __syncthreads(); 128 | if (tmp == 1) { 129 | atomicAdd(&flag, 1); 130 | for (int j = 1; j <= degree[threadIdx.x]; j++) { 131 | int u = neighbor[(i * n + blockIdx.y * primitive_size + threadIdx.x) * 512 + j]; 132 | if (cnt[u] > 1 || (cnt[u] == 1 && threadIdx.x > u)) { 133 | float c = cost[(i * n + blockIdx.y * primitive_size + threadIdx.x) * 512 + j]; 134 | atomicAdd(&cnt[threadIdx.x], -1); 135 | atomicAdd(&cnt[u], -1); 136 | if (c > mean_dis * alpha) { 137 | dist[i * n + blockIdx.y * primitive_size + threadIdx.x] = c; 138 | idx[i * n + blockIdx.y * primitive_size + threadIdx.x] = blockIdx.y * primitive_size + u; 139 | } 140 | } 141 | } 142 | } 143 | __syncthreads(); 144 | if (flag == 0) break; 145 | __syncthreads(); 146 | } 147 | __syncthreads(); 148 | } 149 | } 150 | 151 | int expansion_penalty_cuda_forward(at::Tensor xyz, int primitive_size, at::Tensor idx, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length) { 152 | 153 | const auto batch_size = xyz.size(0); 154 | const auto n = xyz.size(1); 155 | 156 | calc_penalty<<>>(batch_size, n, primitive_size, xyz.data(), idx.data(), dist.data(), 157 | alpha, neighbor.data(), cost.data(), mean_mst_length.data()); 158 | 159 | cudaError_t err = cudaGetLastError(); 160 | if (err != cudaSuccess) { 161 | printf("error in nnd Output: %s\n", cudaGetErrorString(err)); 162 | return 0; 163 | } 164 | return 1; 165 | } 166 | 167 | __global__ void calc_grad(int b, int n, const float * xyz, const float * grad_dist, const int * idx, float * grad_xyz) { 168 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 169 | for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x * gridDim.y) 170 | if (idx[i * n + j] != -1) { 171 | float x1 = xyz[(i * n + j) * 3 + 0]; 172 | float y1 = xyz[(i * n + j) * 3 + 1]; 173 | float z1 = xyz[(i * n + j) * 3 + 2]; 174 | int j2 = idx[i * n + j]; 175 | float x2 = xyz[(i * n + j2) * 3 + 0]; 176 | float y2 = xyz[(i * n + j2) * 3 + 1]; 177 | float z2 = xyz[(i * n + j2) * 3 + 2]; 178 | float g = grad_dist[i * n + j] * 2; 179 | atomicAdd(&(grad_xyz[(i * n + j) * 3 + 0]), g * (x1 - x2)); 180 | atomicAdd(&(grad_xyz[(i * n + j) * 3 + 1]), g * (y1 - y2)); 181 | atomicAdd(&(grad_xyz[(i * n + j) * 3 + 2]), g * (z1 - z2)); 182 | } 183 | } 184 | } 185 | 186 | int expansion_penalty_cuda_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) { 187 | const auto batch_size = xyz.size(0); 188 | const auto n = xyz.size(1); 189 | 190 | calc_grad<<>>(batch_size, n, xyz.data(), graddist.data(), idx.data(), gradxyz.data()); 191 | 192 | cudaError_t err = cudaGetLastError(); 193 | if (err != cudaSuccess) { 194 | printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); 195 | return 0; 196 | } 197 | return 1; 198 | } 199 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/expansion_penalty_module.py: -------------------------------------------------------------------------------- 1 | # Expansion penalty module (based on minimum spanning tree) 2 | # author: Minghua Liu 3 | 4 | # Input: 5 | # xyz: [#batch, #point] 6 | # primitive_size: int, the number of points of sampled points for each surface elements, which should be no greater than 512 7 | # in each point cloud, the points from the same surface element should be successive 8 | # alpha: float, > 1, only penalize those edges whose length are greater than (alpha * mean_length) 9 | 10 | # Output: 11 | # dist: [#batch, #point], if the point u is penalized then dist[u] is the distance between u and its neighbor in the MST, otherwise dist[u] is 0 12 | # assignment: [#batch, #point], if the point u is penalized then assignment[u] is its neighbor in the MST, otherwise assignment[u] is -1 13 | # mean_mst_length: [#batch], the average length of the edeges in each point clouds 14 | 15 | 16 | import time 17 | import numpy as np 18 | import torch 19 | from torch import nn 20 | from torch.autograd import Function 21 | import expansion_penalty 22 | 23 | # GPU tensors only 24 | class expansionPenaltyFunction(Function): 25 | @staticmethod 26 | def forward(ctx, xyz, primitive_size, alpha): 27 | assert primitive_size <= 512 28 | batchsize, n, _ = xyz.size() 29 | assert n % primitive_size == 0 30 | xyz = xyz.contiguous().float().cuda() 31 | dist = torch.zeros(batchsize, n, device="cuda").contiguous() 32 | assignment = torch.zeros(batchsize, n, device="cuda", dtype=torch.int32).contiguous() - 1 33 | neighbor = torch.zeros(batchsize, n * 512, device="cuda", dtype=torch.int32).contiguous() 34 | cost = torch.zeros(batchsize, n * 512, device="cuda").contiguous() 35 | mean_mst_length = torch.zeros(batchsize, device="cuda").contiguous() 36 | expansion_penalty.forward( 37 | xyz, primitive_size, assignment, dist, alpha, neighbor, cost, mean_mst_length 38 | ) 39 | ctx.save_for_backward(xyz, assignment) 40 | return dist, assignment, mean_mst_length / (n / primitive_size) 41 | 42 | @staticmethod 43 | def backward(ctx, grad_dist, grad_idx, grad_mml): 44 | xyz, assignment = ctx.saved_tensors 45 | grad_dist = grad_dist.contiguous() 46 | grad_xyz = torch.zeros(xyz.size(), device="cuda").contiguous() 47 | expansion_penalty.backward(xyz, grad_xyz, grad_dist, assignment) 48 | return grad_xyz, None, None 49 | 50 | 51 | class expansionPenaltyModule(nn.Module): 52 | def __init__(self): 53 | super(expansionPenaltyModule, self).__init__() 54 | 55 | def forward(self, input, primitive_size, alpha): 56 | return expansionPenaltyFunction.apply(input, primitive_size, alpha) 57 | 58 | 59 | def test_expansion_penalty(): 60 | x = torch.rand(20, 8192, 3).cuda() 61 | print("Input_size: ", x.shape) 62 | expansion = expansionPenaltyModule() 63 | start_time = time.perf_counter() 64 | dis, ass, mean_length = expansion(x, 512, 1.5) 65 | print("Runtime: %lfs" % (time.perf_counter() - start_time)) 66 | 67 | 68 | # test_expansion_penalty() 69 | 70 | -------------------------------------------------------------------------------- /cuda/expansion_penalty/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name="expansion_penalty", 6 | ext_modules=[ 7 | CUDAExtension("expansion_penalty", ["expansion_penalty.cpp", "expansion_penalty_cuda.cu",]), 8 | ], 9 | cmdclass={"build_ext": BuildExtension}, 10 | ) 11 | 12 | -------------------------------------------------------------------------------- /cuda/gridding/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-15 20:33:52 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 09:55:53 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding 11 | 12 | 13 | class GriddingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, scale, ptcloud): 16 | grid, grid_pt_weights, grid_pt_indexes = gridding.forward( 17 | -scale, scale - 1, -scale, scale - 1, -scale, scale - 1, ptcloud 18 | ) 19 | # print(grid.size()) # torch.Size(batch_size, n_grid_vertices) 20 | # print(grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 21 | # print(grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 22 | ctx.save_for_backward(grid_pt_weights, grid_pt_indexes) 23 | 24 | return grid 25 | 26 | @staticmethod 27 | def backward(ctx, grad_grid): 28 | grid_pt_weights, grid_pt_indexes = ctx.saved_tensors 29 | grad_ptcloud = gridding.backward(grid_pt_weights, grid_pt_indexes, grad_grid) 30 | # print(grad_ptcloud.size()) # torch.Size(batch_size, n_pts, 3) 31 | 32 | return None, grad_ptcloud 33 | 34 | 35 | class Gridding(torch.nn.Module): 36 | def __init__(self, scale=1): 37 | super(Gridding, self).__init__() 38 | self.scale = scale // 2 39 | 40 | def forward(self, ptcloud): 41 | ptcloud = ptcloud * self.scale 42 | _ptcloud = torch.split(ptcloud, 1, dim=0) 43 | grids = [] 44 | for p in _ptcloud: 45 | non_zeros = torch.sum(p, dim=2).ne(0) 46 | p = p[non_zeros].unsqueeze(dim=0) 47 | grids.append(GriddingFunction.apply(self.scale, p)) 48 | 49 | return torch.cat(grids, dim=0).contiguous() 50 | 51 | 52 | class GriddingReverseFunction(torch.autograd.Function): 53 | @staticmethod 54 | def forward(ctx, scale, grid): 55 | ptcloud = gridding.rev_forward(scale, grid) 56 | ctx.save_for_backward(torch.Tensor([scale]), grid, ptcloud) 57 | return ptcloud 58 | 59 | @staticmethod 60 | def backward(ctx, grad_ptcloud): 61 | scale, grid, ptcloud = ctx.saved_tensors 62 | scale = int(scale.item()) 63 | grad_grid = gridding.rev_backward(ptcloud, grid, grad_ptcloud) 64 | grad_grid = grad_grid.view(-1, scale, scale, scale) 65 | return None, grad_grid 66 | 67 | 68 | class GriddingReverse(torch.nn.Module): 69 | def __init__(self, scale=1): 70 | super(GriddingReverse, self).__init__() 71 | self.scale = scale 72 | 73 | def forward(self, grid): 74 | ptcloud = GriddingReverseFunction.apply(self.scale, grid) 75 | return ptcloud / self.scale * 2 76 | -------------------------------------------------------------------------------- /cuda/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/gridding/gridding.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: gridding 3 | Version: 2.1.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/gridding/gridding.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | gridding.cu 2 | gridding_cuda.cpp 3 | gridding_reverse.cu 4 | setup.py 5 | gridding.egg-info/PKG-INFO 6 | gridding.egg-info/SOURCES.txt 7 | gridding.egg-info/dependency_links.txt 8 | gridding.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/gridding/gridding.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/gridding/gridding.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | gridding 2 | -------------------------------------------------------------------------------- /cuda/gridding/gridding_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-11-13 10:52:53 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:32 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | torch::Tensor gridding_reverse_cuda_forward(int scale, 35 | torch::Tensor grid, 36 | cudaStream_t stream); 37 | 38 | torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, 39 | torch::Tensor grid, 40 | torch::Tensor grad_ptcloud, 41 | cudaStream_t stream); 42 | 43 | std::vector gridding_forward(float min_x, 44 | float max_x, 45 | float min_y, 46 | float max_y, 47 | float min_z, 48 | float max_z, 49 | torch::Tensor ptcloud) 50 | { 51 | CHECK_INPUT(ptcloud); 52 | 53 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 54 | return gridding_cuda_forward(min_x, max_x, min_y, max_y, min_z, max_z, 55 | ptcloud, stream); 56 | } 57 | 58 | torch::Tensor gridding_backward(torch::Tensor grid_pt_weights, 59 | torch::Tensor grid_pt_indexes, 60 | torch::Tensor grad_grid) 61 | { 62 | CHECK_INPUT(grid_pt_weights); 63 | CHECK_INPUT(grid_pt_indexes); 64 | CHECK_INPUT(grad_grid); 65 | 66 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 67 | return gridding_cuda_backward(grid_pt_weights, grid_pt_indexes, grad_grid, 68 | stream); 69 | } 70 | 71 | torch::Tensor gridding_reverse_forward(int scale, torch::Tensor grid) 72 | { 73 | CHECK_INPUT(grid); 74 | 75 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 76 | return gridding_reverse_cuda_forward(scale, grid, stream); 77 | } 78 | 79 | torch::Tensor gridding_reverse_backward(torch::Tensor ptcloud, 80 | torch::Tensor grid, 81 | torch::Tensor grad_ptcloud) 82 | { 83 | CHECK_INPUT(ptcloud); 84 | CHECK_INPUT(grid); 85 | CHECK_INPUT(grad_ptcloud); 86 | 87 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 88 | return gridding_reverse_cuda_backward(ptcloud, grid, grad_ptcloud, stream); 89 | } 90 | 91 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 92 | { 93 | m.def("forward", &gridding_forward, "Gridding forward (CUDA)"); 94 | m.def("backward", &gridding_backward, "Gridding backward (CUDA)"); 95 | m.def("rev_forward", &gridding_reverse_forward, 96 | "Gridding Reverse forward (CUDA)"); 97 | m.def("rev_backward", &gridding_reverse_backward, 98 | "Gridding Reverse backward (CUDA)"); 99 | } 100 | -------------------------------------------------------------------------------- /cuda/gridding/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-13 10:51:33 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-02 17:02:16 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup( 12 | name="gridding", 13 | version="2.1.0", 14 | ext_modules=[CUDAExtension("gridding", ["gridding_cuda.cpp", "gridding.cu", "gridding_reverse.cu"]),], 15 | cmdclass={"build_ext": BuildExtension}, 16 | ) 17 | -------------------------------------------------------------------------------- /cuda/gridding/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:48:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:20:42 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in gridding.cu and gridding_reverse.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.gridding import GriddingFunction, GriddingReverseFunction 20 | 21 | 22 | class GriddingTestCase(unittest.TestCase): 23 | def test_gridding_reverse_function_4(self): 24 | x = torch.rand(2, 4, 4, 4) 25 | x.requires_grad = True 26 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [4, x.double().cuda()])) 27 | 28 | def test_gridding_reverse_function_8(self): 29 | x = torch.rand(4, 8, 8, 8) 30 | x.requires_grad = True 31 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [8, x.double().cuda()])) 32 | 33 | def test_gridding_reverse_function_16(self): 34 | x = torch.rand(1, 16, 16, 16) 35 | x.requires_grad = True 36 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [16, x.double().cuda()])) 37 | 38 | def test_gridding_function_32pts(self): 39 | x = torch.rand(1, 32, 3) 40 | x.requires_grad = True 41 | self.assertTrue(gradcheck(GriddingFunction.apply, [x.double().cuda()])) 42 | 43 | 44 | if __name__ == "__main__": 45 | unittest.main() 46 | -------------------------------------------------------------------------------- /cuda/gridding_loss/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 09:56:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-02-22 19:19:43 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding_distance 11 | 12 | 13 | class GriddingDistanceFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud, gt_cloud): 16 | pred_grid, pred_grid_pt_weights, pred_grid_pt_indexes = gridding_distance.forward( 17 | min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud 18 | ) 19 | # print(pred_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 20 | # print(pred_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 21 | # print(pred_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 22 | gt_grid, gt_grid_pt_weights, gt_grid_pt_indexes = gridding_distance.forward( 23 | min_x, max_x, min_y, max_y, min_z, max_z, gt_cloud 24 | ) 25 | # print(gt_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 26 | # print(gt_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 27 | # print(gt_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 28 | 29 | ctx.save_for_backward( 30 | pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes 31 | ) 32 | return pred_grid, gt_grid 33 | 34 | @staticmethod 35 | def backward(ctx, grad_pred_grid, grad_gt_grid): 36 | pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes = ctx.saved_tensors 37 | 38 | grad_pred_cloud = gridding_distance.backward( 39 | pred_grid_pt_weights, pred_grid_pt_indexes, grad_pred_grid 40 | ) 41 | # print(grad_pred_cloud.size()) # torch.Size(batch_size, n_pts, 3) 42 | grad_gt_cloud = gridding_distance.backward(gt_grid_pt_weights, gt_grid_pt_indexes, grad_gt_grid) 43 | # print(grad_gt_cloud.size()) # torch.Size(batch_size, n_pts, 3) 44 | 45 | return None, None, None, None, None, None, grad_pred_cloud, grad_gt_cloud 46 | 47 | 48 | class GriddingDistance(torch.nn.Module): 49 | def __init__(self, scale=1): 50 | super(GriddingDistance, self).__init__() 51 | self.scale = scale 52 | 53 | def forward(self, pred_cloud, gt_cloud): 54 | """ 55 | pred_cloud(b, n_pts1, 3) 56 | gt_cloud(b, n_pts2, 3) 57 | """ 58 | pred_cloud = pred_cloud * self.scale / 2 59 | gt_cloud = gt_cloud * self.scale / 2 60 | 61 | min_pred_x = torch.min(pred_cloud[:, :, 0]) 62 | max_pred_x = torch.max(pred_cloud[:, :, 0]) 63 | min_pred_y = torch.min(pred_cloud[:, :, 1]) 64 | max_pred_y = torch.max(pred_cloud[:, :, 1]) 65 | min_pred_z = torch.min(pred_cloud[:, :, 2]) 66 | max_pred_z = torch.max(pred_cloud[:, :, 2]) 67 | 68 | min_gt_x = torch.min(gt_cloud[:, :, 0]) 69 | max_gt_x = torch.max(gt_cloud[:, :, 0]) 70 | min_gt_y = torch.min(gt_cloud[:, :, 1]) 71 | max_gt_y = torch.max(gt_cloud[:, :, 1]) 72 | min_gt_z = torch.min(gt_cloud[:, :, 2]) 73 | max_gt_z = torch.max(gt_cloud[:, :, 2]) 74 | 75 | min_x = torch.floor(torch.min(min_pred_x, min_gt_x)) - 1 76 | max_x = torch.ceil(torch.max(max_pred_x, max_gt_x)) + 1 77 | min_y = torch.floor(torch.min(min_pred_y, min_gt_y)) - 1 78 | max_y = torch.ceil(torch.max(max_pred_y, max_gt_y)) + 1 79 | min_z = torch.floor(torch.min(min_pred_z, min_gt_z)) - 1 80 | max_z = torch.ceil(torch.max(max_pred_z, max_gt_z)) + 1 81 | 82 | _pred_clouds = torch.split(pred_cloud, 1, dim=0) 83 | _gt_clouds = torch.split(gt_cloud, 1, dim=0) 84 | pred_grids = [] 85 | gt_grids = [] 86 | for pc, gc in zip(_pred_clouds, _gt_clouds): 87 | non_zeros = torch.sum(pc, dim=2).ne(0) 88 | pc = pc[non_zeros].unsqueeze(dim=0) 89 | non_zeros = torch.sum(gc, dim=2).ne(0) 90 | gc = gc[non_zeros].unsqueeze(dim=0) 91 | pred_grid, gt_grid = GriddingDistanceFunction.apply( 92 | min_x, max_x, min_y, max_y, min_z, max_z, pc, gc 93 | ) 94 | pred_grids.append(pred_grid) 95 | gt_grids.append(gt_grid) 96 | 97 | return torch.cat(pred_grids, dim=0).contiguous(), torch.cat(gt_grids, dim=0).contiguous() 98 | 99 | 100 | class GriddingLoss(torch.nn.Module): 101 | def __init__(self, scales=[], alphas=[]): 102 | super(GriddingLoss, self).__init__() 103 | self.scales = scales 104 | self.alphas = alphas 105 | self.gridding_dists = [GriddingDistance(scale=s) for s in scales] 106 | self.l1_loss = torch.nn.L1Loss() 107 | 108 | def forward(self, pred_cloud, gt_cloud): 109 | gridding_loss = None 110 | n_dists = len(self.scales) 111 | 112 | for i in range(n_dists): 113 | alpha = self.alphas[i] 114 | gdist = self.gridding_dists[i] 115 | pred_grid, gt_grid = gdist(pred_cloud, gt_cloud) 116 | 117 | if gridding_loss is None: 118 | gridding_loss = alpha * self.l1_loss(pred_grid, gt_grid) 119 | else: 120 | gridding_loss += alpha * self.l1_loss(pred_grid, gt_grid) 121 | 122 | return gridding_loss 123 | -------------------------------------------------------------------------------- /cuda/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/cuda/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /cuda/gridding_loss/gridding_distance.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: gridding-distance 3 | Version: 1.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /cuda/gridding_loss/gridding_distance.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | gridding_distance.cu 2 | gridding_distance_cuda.cpp 3 | setup.py 4 | gridding_distance.egg-info/PKG-INFO 5 | gridding_distance.egg-info/SOURCES.txt 6 | gridding_distance.egg-info/dependency_links.txt 7 | gridding_distance.egg-info/top_level.txt -------------------------------------------------------------------------------- /cuda/gridding_loss/gridding_distance.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /cuda/gridding_loss/gridding_distance.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | gridding_distance 2 | -------------------------------------------------------------------------------- /cuda/gridding_loss/gridding_distance_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-30 10:59:31 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:52 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_distance_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | std::vector gridding_distance_forward(float min_x, 35 | float max_x, 36 | float min_y, 37 | float max_y, 38 | float min_z, 39 | float max_z, 40 | torch::Tensor ptcloud) 41 | { 42 | CHECK_INPUT(ptcloud); 43 | 44 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 45 | return gridding_distance_cuda_forward(min_x, max_x, min_y, max_y, min_z, 46 | max_z, ptcloud, stream); 47 | } 48 | 49 | torch::Tensor gridding_distance_backward(torch::Tensor grid_pt_weights, 50 | torch::Tensor grid_pt_indexes, 51 | torch::Tensor grad_grid) 52 | { 53 | CHECK_INPUT(grid_pt_weights); 54 | CHECK_INPUT(grid_pt_indexes); 55 | CHECK_INPUT(grad_grid); 56 | 57 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 58 | return gridding_distance_cuda_backward(grid_pt_weights, grid_pt_indexes, 59 | grad_grid, stream); 60 | } 61 | 62 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 63 | { 64 | m.def("forward", &gridding_distance_forward, 65 | "Gridding Distance Forward (CUDA)"); 66 | m.def("backward", &gridding_distance_backward, 67 | "Gridding Distance Backward (CUDA)"); 68 | } 69 | -------------------------------------------------------------------------------- /cuda/gridding_loss/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 11:03:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 11:13:39 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup( 12 | name="gridding_distance", 13 | version="1.0.0", 14 | ext_modules=[CUDAExtension("gridding_distance", ["gridding_distance_cuda.cpp", "gridding_distance.cu"]),], 15 | cmdclass={"build_ext": BuildExtension}, 16 | ) 17 | -------------------------------------------------------------------------------- /cuda/p2i_op/.clang-format: -------------------------------------------------------------------------------- 1 | Language: Cpp 2 | BasedOnStyle: LLVM 3 | -------------------------------------------------------------------------------- /cuda/p2i_op/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | from torch import nn 5 | from torch.autograd import Function 6 | from torch.utils.cpp_extension import load 7 | 8 | __all__ = ["p2i"] 9 | 10 | module_path = os.path.dirname(os.path.abspath(__file__)) 11 | ext = load( 12 | "ext", 13 | sources=[ 14 | os.path.join(module_path, "ext.cpp"), 15 | os.path.join(module_path, "p2i_sum.cu"), 16 | os.path.join(module_path, "p2i_max.cu"), 17 | ], 18 | extra_cuda_cflags=["--expt-extended-lambda", "-O3"], 19 | ) 20 | 21 | 22 | class P2ISumFunction(Function): 23 | @staticmethod 24 | def forward(ctx, points, point_features, batch_inds, background, kernel_kind, kernel_radius): 25 | ctx.save_for_backward(points, point_features, batch_inds) 26 | ctx.kernel_kind = kernel_kind 27 | ctx.kernel_radius = kernel_radius 28 | 29 | out = ext.p2i_sum_forward_gpu( 30 | points.contiguous(), 31 | point_features.contiguous(), 32 | batch_inds.contiguous(), 33 | background.contiguous(), 34 | kernel_kind, 35 | kernel_radius, 36 | ) 37 | 38 | return out 39 | 40 | @staticmethod 41 | def backward(ctx, out_grad): 42 | points, point_features, batch_inds = ctx.saved_tensors 43 | kernel_kind = ctx.kernel_kind 44 | kernel_radius = ctx.kernel_radius 45 | 46 | points_grad, point_features_grad = ext.p2i_sum_backward_gpu( 47 | out_grad.contiguous(), 48 | points.contiguous(), 49 | point_features.contiguous(), 50 | batch_inds.contiguous(), 51 | kernel_kind, 52 | kernel_radius, 53 | ) 54 | 55 | background_grad = out_grad 56 | return points_grad, point_features_grad, None, background_grad, None, None 57 | 58 | 59 | class P2IMaxFunction(Function): 60 | @staticmethod 61 | def forward(ctx, points, point_features, batch_inds, background, kernel_kind, kernel_radius): 62 | 63 | out, out_point_ids = ext.p2i_max_forward_gpu( 64 | points.contiguous(), 65 | point_features.contiguous(), 66 | batch_inds.contiguous(), 67 | background.contiguous(), 68 | kernel_kind, 69 | kernel_radius, 70 | ) 71 | 72 | ctx.save_for_backward(points, point_features, out_point_ids) 73 | ctx.kernel_kind = kernel_kind 74 | ctx.kernel_radius = kernel_radius 75 | 76 | return out 77 | 78 | @staticmethod 79 | def backward(ctx, out_grad): 80 | points, point_features, out_point_ids = ctx.saved_tensors 81 | kernel_kind = ctx.kernel_kind 82 | kernel_radius = ctx.kernel_radius 83 | 84 | points_grad, point_features_grad, background_grad = ext.p2i_max_backward_gpu( 85 | out_grad.contiguous(), 86 | out_point_ids, 87 | points.contiguous(), 88 | point_features.contiguous(), 89 | kernel_kind, 90 | kernel_radius, 91 | ) 92 | 93 | return points_grad, point_features_grad, None, background_grad, None, None 94 | 95 | 96 | _kernel_kind_dict = {"cos": 0} 97 | 98 | 99 | def p2i(points, point_features, batch_inds, background, kernel_radius, kernel_kind_str="cos", reduce="sum"): 100 | r"""p2i 101 | 102 | Paint point cloud features on to a 2D feature map. 103 | 104 | inputs: 105 | - points: float, [npoints x 2], (+/-1, +/-1) represents the image corners 106 | - point_features: float, [npoints x channels] 107 | - batch_inds: int32, [npoints] 108 | - background: float, [batch x channels x out_h x out_w] 109 | - kernel_radius: float 110 | - kernel_kind_str: str, {'cos'} 111 | - reduce: str, {'sum', 'max'} 112 | returns: 113 | - output: float, [batch x channels x out_h x out_w] 114 | """ 115 | kernel_kind = _kernel_kind_dict[kernel_kind_str] 116 | out_h, out_w = background.shape[2:] 117 | points = ( 118 | (points + 1) 119 | / 2 120 | * torch.tensor([out_h - 1, out_w - 1], dtype=points.dtype, device=points.device).view(1, 2) 121 | ) 122 | 123 | if reduce == "sum": 124 | return P2ISumFunction.apply( 125 | points, point_features, batch_inds, background, kernel_kind, kernel_radius 126 | ) 127 | elif reduce == "max": 128 | return P2IMaxFunction.apply( 129 | points, point_features, batch_inds, background, kernel_kind, kernel_radius 130 | ) 131 | raise RuntimeError(f"Invalid reduce value: {reduce}") 132 | 133 | custom_fun = P2ISumFunction.apply 134 | -------------------------------------------------------------------------------- /cuda/p2i_op/common.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #ifdef _OPENMP 6 | #include 7 | #endif 8 | 9 | #ifdef __CUDACC__ 10 | #include 11 | #include 12 | #endif 13 | 14 | #ifdef _MSC_VER 15 | #define FORCE_INLINE __forceinline 16 | #pragma warning(disable : 4068) 17 | #else 18 | #define FORCE_INLINE __attribute__((always_inline)) 19 | #endif 20 | 21 | #ifdef __CUDACC__ 22 | #define XINLINE __device__ __host__ 23 | #define XGLOBAL __global__ 24 | #define XDEVICE __device__ 25 | #define XSHARED __shared__ 26 | #else 27 | #define XINLINE 28 | #define XGLOBAL 29 | #define XDEVICE 30 | #define XSHARED 31 | #endif 32 | 33 | namespace haya_ext { 34 | #ifdef __CUDACC__ 35 | // Macro for checking cuda errors following a cuda launch or api call 36 | #define cudaCheckError() \ 37 | do { \ 38 | cudaError_t e = cudaGetLastError(); \ 39 | if (e != cudaSuccess) { \ 40 | char buffer[512] = {'\0'}; \ 41 | sprintf(buffer, "Cuda failure %s:%d: '%s(%s)'", __FILE__, __LINE__, \ 42 | cudaGetErrorName(e), cudaGetErrorString(e)); \ 43 | AT_ERROR(buffer); \ 44 | } \ 45 | } while (0) 46 | #else 47 | #define cudaCheckError() 48 | #endif 49 | 50 | struct cpu_device {}; 51 | struct gpu_device {}; 52 | 53 | template struct kernel; 54 | 55 | template <> struct kernel { 56 | template 57 | inline static FORCE_INLINE void launch(OP op, const int N, Args... args) { 58 | #ifdef _OPENMP 59 | const int omp_cores = omp_get_thread_num(); 60 | if (omp_cores <= 1) { 61 | // Zero means not to use OMP, but don't interfere with external OMP 62 | // behavior 63 | for (int i = 0; i < N; ++i) { 64 | op(i, args...); 65 | } 66 | } else { 67 | #pragma omp parallel for num_threads(omp_cores) 68 | for (int i = 0; i < N; ++i) { 69 | op(i, args...); 70 | } 71 | } 72 | #else 73 | for (int i = 0; i < N; ++i) { 74 | op(i, args...); 75 | } 76 | #endif 77 | } 78 | }; 79 | 80 | #if defined(NO_CUDA) // try launching gpu kernel from a no cuda build 81 | template <> struct kernel { 82 | template 83 | inline static FORCE_INLINE void launch(OP op, const int N, Args... args) { 84 | AT_ERROR("failed to launch cuda kernel in a NO CUDA build"); 85 | } 86 | }; 87 | #elif defined(__CUDACC__) // launching gpu kernel within nvcc compilation 88 | namespace detail { 89 | constexpr int kMaxThreadsPerBlock = 1024; 90 | constexpr int kMaxGridNum = 65535; 91 | constexpr int kBaseThreadBits = 8; 92 | constexpr int kBaseThreadNum = 1 << kBaseThreadBits; 93 | constexpr int kBaseGridNum = 1024; 94 | 95 | template 96 | XGLOBAL void _generic_kernel(OP op, int N, Args... args) { 97 | for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; 98 | i += blockDim.x * gridDim.x) { 99 | op(i, args...); 100 | } 101 | } 102 | } // namespace detail 103 | 104 | template <> struct kernel { 105 | template 106 | inline static FORCE_INLINE void launch(OP op, const int N, Args... args) { 107 | static_assert(std::is_class::value, 108 | "You should pass a functor (including lambda) to " 109 | "kernel::launch. Passing a function pointer " 110 | "will cause cuda error in runtime."); 111 | const dim3 blocks = 112 | (N + detail::kBaseThreadNum - 1) / detail::kBaseThreadNum; 113 | detail::_generic_kernel 114 | <<>>(op, N, args...); 115 | } 116 | template 117 | inline static FORCE_INLINE void launch_max_threads(OP op, const int N, 118 | Args... args) { 119 | static_assert(std::is_class::value, 120 | "You should pass a functor (including lambda) to " 121 | "kernel::launch. Passing a function pointer " 122 | "will cause cuda error in runtime."); 123 | const dim3 blocks = 124 | (N + detail::kMaxThreadsPerBlock - 1) / detail::kMaxThreadsPerBlock; 125 | detail::_generic_kernel 126 | <<>>(op, N, args...); 127 | } 128 | }; 129 | #else // try launching gpu kernel without nvcc compilation, this should not 130 | // compile 131 | namespace detail { 132 | template struct always_false { 133 | static constexpr bool value = false; 134 | }; 135 | } // namespace detail 136 | template <> struct kernel { 137 | template 138 | inline static FORCE_INLINE void launch(OP op, const int N, Args... args) { 139 | static_assert(detail::always_false::value, 140 | "trying to instantiate gpu kernel under non cuda context"); 141 | } 142 | }; 143 | #endif 144 | } // namespace haya_ext 145 | -------------------------------------------------------------------------------- /cuda/p2i_op/ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "p2i_max.h" 4 | #include "p2i_sum.h" 5 | 6 | using namespace haya_ext; 7 | 8 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 9 | m.def("p2i_sum_forward_gpu", &p2i_sum_forward_gpu, "p2i sum forward (CUDA)"); 10 | m.def("p2i_sum_backward_gpu", &p2i_sum_backward_gpu, 11 | "p2i sum backward (CUDA)"); 12 | 13 | m.def("p2i_max_forward_gpu", &p2i_max_forward_gpu, "p2i max forward (CUDA)"); 14 | m.def("p2i_max_backward_gpu", &p2i_max_backward_gpu, 15 | "p2i max backward (CUDA)"); 16 | } -------------------------------------------------------------------------------- /cuda/p2i_op/p2i_max.cu: -------------------------------------------------------------------------------- 1 | #include "p2i_max.h" 2 | 3 | namespace haya_ext { 4 | std::vector 5 | p2i_max_forward_gpu(const at::Tensor &points, const at::Tensor &point_features, 6 | const at::Tensor &batch_inds, const at::Tensor &background, 7 | int kernel_kind, double kernel_radius) { 8 | return p2i_max_op::forward(points, point_features, batch_inds, 9 | background, kernel_kind, 10 | kernel_radius); 11 | } 12 | 13 | std::vector p2i_max_backward_gpu(const at::Tensor &out_grad, 14 | const at::Tensor &out_point_ids, 15 | const at::Tensor &points, 16 | const at::Tensor &point_features, 17 | int kernel_kind, 18 | double kernel_radius) { 19 | return p2i_max_op::backward(out_grad, out_point_ids, points, 20 | point_features, kernel_kind, 21 | kernel_radius); 22 | } 23 | } // namespace haya_ext -------------------------------------------------------------------------------- /cuda/p2i_op/p2i_sum.cu: -------------------------------------------------------------------------------- 1 | #include "p2i_sum.h" 2 | 3 | namespace haya_ext { 4 | at::Tensor p2i_sum_forward_gpu(const at::Tensor &points, 5 | const at::Tensor &point_features, 6 | const at::Tensor &batch_inds, 7 | const at::Tensor &background, int kernel_kind, 8 | double kernel_radius) { 9 | return p2i_sum_op::forward(points, point_features, batch_inds, 10 | background, kernel_kind, 11 | kernel_radius); 12 | } 13 | 14 | std::vector p2i_sum_backward_gpu(const at::Tensor &out_grad, 15 | const at::Tensor &points, 16 | const at::Tensor &point_features, 17 | const at::Tensor &batch_inds, 18 | int kernel_kind, 19 | double kernel_radius) { 20 | return p2i_sum_op::backward( 21 | out_grad, points, point_features, batch_inds, kernel_kind, kernel_radius); 22 | } 23 | } // namespace haya_ext -------------------------------------------------------------------------------- /cuda/p2i_op/p2i_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.autograd import gradcheck, grad 4 | 5 | import time 6 | 7 | from . import p2i 8 | 9 | 10 | def test1(): 11 | points = torch.zeros(1, 2, dtype=torch.float64, device="cuda") 12 | point_features = torch.ones(1, 3, dtype=torch.float64, device="cuda") 13 | batch_inds = torch.arange(1, dtype=torch.int32, device="cuda") 14 | background = torch.zeros(1, 3, 8, 8, dtype=torch.float64, device="cuda") 15 | 16 | out = p2i(points, point_features, batch_inds, background, 2, "cos", "sum") 17 | print(out) 18 | 19 | out = p2i(points, point_features, batch_inds, background, 2, "cos", "max") 20 | print(out) 21 | 22 | 23 | def test2(): 24 | for i in range(10): 25 | print(i) 26 | points = torch.randn(2, 2, dtype=torch.float64, device="cuda") 27 | point_features = torch.randn(2, 3, dtype=torch.float64, device="cuda") 28 | batch_inds = torch.zeros(2, dtype=torch.int32, device="cuda") 29 | background = torch.randn(1, 3, 8, 8, dtype=torch.float64, device="cuda") 30 | 31 | points.requires_grad = True 32 | point_features.requires_grad = True 33 | background.requires_grad = True 34 | gradcheck(p2i, inputs=(points, point_features, batch_inds, background, 2, "cos", "sum")) 35 | gradcheck(p2i, inputs=(points, point_features, batch_inds, background, 2, "cos", "max")) 36 | 37 | 38 | if __name__ == "__main__": 39 | test2() 40 | -------------------------------------------------------------------------------- /cuda/p2i_op/utility.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #ifndef __CUDACC__ 4 | #include 5 | #include 6 | using std::max; 7 | using std::min; 8 | #endif 9 | 10 | #include "common.h" 11 | 12 | namespace haya_ext { 13 | 14 | template XDEVICE T atomic_add(T *addr, T v) { 15 | #ifdef __CUDACC__ 16 | return atomicAdd(addr, v); 17 | #else 18 | return *addr += v; 19 | #endif 20 | } 21 | 22 | #if defined(__CUDACC__) && __CUDA_ARCH__ < 600 // old cuda 23 | __forceinline__ __device__ double atomic_add(double *addr, double v) { 24 | unsigned long long int *address_as_ull = (unsigned long long int *)addr; 25 | unsigned long long int old = *address_as_ull, assumed; 26 | do { 27 | assumed = old; 28 | old = atomicCAS(address_as_ull, assumed, 29 | __double_as_longlong(v + __longlong_as_double(assumed))); 30 | } while (assumed != old); 31 | return __longlong_as_double(old); 32 | } 33 | #endif 34 | 35 | template XDEVICE T atomic_cas(T *addr, T compare, T val) { 36 | #ifdef __CUDACC__ 37 | return atomicCAS(addr, compare, val); 38 | #else 39 | std::atomic this_val(*addr); 40 | this_val.compare_exchange_weak(compare, val); 41 | return *addr = this_val.load(); 42 | #endif 43 | } 44 | 45 | template XDEVICE T atomic_exch(T *addr, T v) { 46 | #ifdef __CUDACC__ 47 | return atomicExch(addr, v); 48 | #else 49 | T rd = *addr; 50 | *addr = v; 51 | return rd; 52 | #endif 53 | } 54 | 55 | #ifdef __CUDACC__ 56 | __forceinline__ __device__ double atomic_exch(double *addr, double v) { 57 | return atomicExch((unsigned long long int *)addr, __double_as_longlong(v)); 58 | } 59 | #endif 60 | 61 | template XINLINE T square(T v) { return v * v; } 62 | 63 | template XINLINE T clamp(T v, T min_v, T max_v) { 64 | v = v < min_v ? min_v : v; 65 | v = v > max_v ? max_v : v; 66 | return v; 67 | } 68 | 69 | template XINLINE void swap(T &a, T &b) { 70 | T tmp = a; 71 | a = b; 72 | b = tmp; 73 | } 74 | 75 | template XINLINE T min3(T a, T b, T c) { 76 | return min(min(a, b), c); 77 | } 78 | template XINLINE T max3(T a, T b, T c) { 79 | return max(max(a, b), c); 80 | } 81 | 82 | template 83 | XDEVICE void for_each_pixel_near_point(scalar_t py, scalar_t px, int out_h, 84 | int out_w, scalar_t radius, 85 | FunT callback) { 86 | int min_x = clamp(floor(px - radius), 0, out_w - 1); 87 | int max_x = clamp(ceil(px + radius), 0, out_w - 1); 88 | int min_y = clamp(floor(py - radius), 0, out_h - 1); 89 | int max_y = clamp(ceil(py + radius), 0, out_h - 1); 90 | for (int x = min_x; x <= max_x; x++) { 91 | for (int y = min_y; y <= max_y; y++) { 92 | scalar_t dx = x - px; 93 | scalar_t dy = y - py; 94 | scalar_t r = sqrt(dx * dx + dy * dy); 95 | if (r <= radius) { 96 | callback(y, x, dy, dx, r); 97 | } 98 | } 99 | } 100 | } 101 | } // namespace haya_ext 102 | -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 | ## Data 2 | 3 | ### Shapenet Dataset (16284 points) 4 | 5 | - Download the processed dataset from [GRNet](https://github.com/hzxie/GRNet): https://gateway.infinitescript.com/?fileName=ShapeNetCompletion 6 | - dataset structure: category; model-id; 8 views (00.pcd, 01.pcd ....... 07.pcd) 7 | - Or you can preprocess the [ShapeNet](https://drive.google.com/drive/folders/1P_W1tz5Q4ZLapUifuOE4rFAZp6L1XTJz) dataset from [PCN](https://github.com/wentaoyuan/pcn/tree/master/data) 8 | 9 | ```bash 10 | cd /utils 11 | python lmdb_serializer.py /path/to/shapenet/train.lmdb /path/to/output/shapenet/train 12 | python lmdb_serializer.py /path/to/shapenet/valid.lmdb /path/to/output/shapenet/val 13 | ``` 14 | where `lmdb_serializer.py` can be obtained [here](https://github.com/hzxie/GRNet/blob/master/utils/lmdb_serializer.py). 15 | 16 | ### KITTI Dataset 17 | - Download the dataset: https://drive.google.com/drive/folders/1fSu0_huWhticAlzLh3Ejpg8zxzqO1z-F -------------------------------------------------------------------------------- /datasets/io.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import sys 6 | from io import BytesIO 7 | 8 | import cv2 9 | import h5py 10 | import numpy as np 11 | import open3d 12 | 13 | from configs.base_config import cfg 14 | 15 | 16 | class IO: 17 | @classmethod 18 | def get(cls, file_path): 19 | _, file_extension = os.path.splitext(file_path) 20 | 21 | if file_extension in [".png", ".jpg"]: 22 | return cls._read_img(file_path) 23 | elif file_extension in [".npy"]: 24 | return cls._read_npy(file_path) 25 | elif file_extension in [".pcd"]: 26 | return cls._read_pcd(file_path) 27 | elif file_extension in [".h5"]: 28 | return cls._read_h5(file_path) 29 | elif file_extension in [".txt"]: 30 | return cls._read_txt(file_path) 31 | else: 32 | raise Exception("Unsupported file extension: %s" % file_extension) 33 | 34 | @classmethod 35 | def put(cls, file_path, file_content): 36 | _, file_extension = os.path.splitext(file_path) 37 | 38 | if file_extension in [".pcd"]: 39 | return cls._write_pcd(file_path, file_content) 40 | elif file_extension in [".h5"]: 41 | return cls._write_h5(file_path, file_content) 42 | else: 43 | raise Exception("Unsupported file extension: %s" % file_extension) 44 | 45 | @classmethod 46 | def _read_img(cls, file_path): 47 | return cv2.imread(file_path, cv2.IMREAD_UNCHANGED) / 255.0 48 | 49 | # References: https://github.com/numpy/numpy/blob/master/numpy/lib/format.py 50 | @classmethod 51 | def _read_npy(cls, file_path): 52 | return np.load(file_path) 53 | 54 | # References: https://github.com/dimatura/pypcd/blob/master/pypcd/pypcd.py#L275 55 | # NOTE: Support PCD files without compression ONLY! 56 | @classmethod 57 | def _read_pcd(cls, file_path): 58 | pc = open3d.io.read_point_cloud(file_path) 59 | return np.array(pc.points) 60 | 61 | @classmethod 62 | def _read_h5(cls, file_path): 63 | f = h5py.File(file_path, "r") 64 | # Avoid overflow while gridding 65 | return f["data"][()] * 0.9 66 | 67 | @classmethod 68 | def _read_txt(cls, file_path): 69 | return np.loadtxt(file_path) 70 | 71 | @classmethod 72 | def _write_pcd(cls, file_path, file_content): 73 | pc = open3d.geometry.PointCloud() 74 | pc.points = open3d.utility.Vector3dVector(file_content) 75 | open3d.io.write_point_cloud(file_path, pc) 76 | 77 | @classmethod 78 | def _write_h5(cls, file_path, file_content): 79 | with h5py.File(file_path, "w") as f: 80 | f.create_dataset("data", data=file_content) 81 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/models/__init__.py -------------------------------------------------------------------------------- /models/atlasnet_generator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.parallel 7 | import torch.utils.data 8 | from torch.autograd import Variable 9 | import numpy as np 10 | import torch.nn.functional as F 11 | from models.sparenet_generator import PointGenCon, PointNetfeat 12 | 13 | 14 | class AtlasNet(nn.Module): 15 | def __init__( 16 | self, 17 | num_points: int = 16382, 18 | bottleneck_size: int = 1024, 19 | n_primitives: int = 32, 20 | ): 21 | super(AtlasNet, self).__init__() 22 | self.num_points = num_points 23 | self.bottleneck_size = bottleneck_size 24 | self.n_primitives = n_primitives 25 | 26 | self.encoder = PointEncoder(bottleneck_size=self.bottleneck_size) 27 | self.decoder = nn.ModuleList( 28 | [ 29 | PointGenCon( 30 | input_dim=2 + self.bottleneck_size, 31 | bottleneck_size=2 + self.bottleneck_size, 32 | use_SElayer=False, 33 | ) 34 | for i in range(self.n_primitives) 35 | ] 36 | ) 37 | 38 | def forward(self, data): 39 | x = data["partial_cloud"] 40 | x = x.transpose(1, 2).contiguous() 41 | style = self.encoder(x) 42 | outs = [] 43 | for i in range(self.n_primitives): 44 | rand_grid = Variable( 45 | torch.cuda.FloatTensor( 46 | x.size(0), 2, self.num_points // self.n_primitives 47 | ) 48 | ) 49 | rand_grid.data.uniform_(0, 1) 50 | y = ( 51 | style.unsqueeze(2) 52 | .expand(style.size(0), style.size(1), rand_grid.size(2)) 53 | .contiguous() 54 | ) 55 | y = torch.cat((rand_grid, y), 1).contiguous() 56 | outs.append(self.decoder[i](y)) 57 | outs = torch.cat(outs, 2).contiguous() 58 | return outs.transpose(1, 2).contiguous() 59 | 60 | 61 | class PointEncoder(nn.Module): 62 | def __init__(self, bottleneck_size=1024, hide_size=1024): 63 | super(PointEncoder, self).__init__() 64 | self.feat_extractor = PointNetfeat( 65 | global_feat=True, use_SElayer=False, hide_size=hide_size 66 | ) 67 | self.linear = nn.Linear(hide_size, bottleneck_size) 68 | self.bn = nn.BatchNorm1d(bottleneck_size) 69 | self.relu = nn.ReLU() 70 | 71 | def forward(self, x): 72 | x = self.feat_extractor(x) 73 | x = self.linear(x) 74 | x = self.bn(x) 75 | x = self.relu(x) 76 | return x 77 | -------------------------------------------------------------------------------- /models/grnet_generator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import torch 5 | from cuda.gridding import Gridding, GriddingReverse 6 | from cuda.cubic_feature_sampling import CubicFeatureSampling 7 | 8 | class GRNet(torch.nn.Module): 9 | def __init__(self, cfg): 10 | super(GRNet, self).__init__() 11 | self.gridding = Gridding(scale=64) 12 | self.conv1 = torch.nn.Sequential( 13 | torch.nn.Conv3d(1, 32, kernel_size=4, padding=2), 14 | torch.nn.BatchNorm3d(32), 15 | torch.nn.LeakyReLU(0.2), 16 | torch.nn.MaxPool3d(kernel_size=2), 17 | ) 18 | self.conv2 = torch.nn.Sequential( 19 | torch.nn.Conv3d(32, 64, kernel_size=4, padding=2), 20 | torch.nn.BatchNorm3d(64), 21 | torch.nn.LeakyReLU(0.2), 22 | torch.nn.MaxPool3d(kernel_size=2), 23 | ) 24 | self.conv3 = torch.nn.Sequential( 25 | torch.nn.Conv3d(64, 128, kernel_size=4, padding=2), 26 | torch.nn.BatchNorm3d(128), 27 | torch.nn.LeakyReLU(0.2), 28 | torch.nn.MaxPool3d(kernel_size=2), 29 | ) 30 | self.conv4 = torch.nn.Sequential( 31 | torch.nn.Conv3d(128, 256, kernel_size=4, padding=2), 32 | torch.nn.BatchNorm3d(256), 33 | torch.nn.LeakyReLU(0.2), 34 | torch.nn.MaxPool3d(kernel_size=2), 35 | ) 36 | self.fc5 = torch.nn.Sequential(torch.nn.Linear(16384, 2048), torch.nn.ReLU()) 37 | self.fc6 = torch.nn.Sequential(torch.nn.Linear(2048, 16384), torch.nn.ReLU()) 38 | self.dconv7 = torch.nn.Sequential( 39 | torch.nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, bias=False, padding=1), 40 | torch.nn.BatchNorm3d(128), 41 | torch.nn.ReLU(), 42 | ) 43 | self.dconv8 = torch.nn.Sequential( 44 | torch.nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, bias=False, padding=1), 45 | torch.nn.BatchNorm3d(64), 46 | torch.nn.ReLU(), 47 | ) 48 | self.dconv9 = torch.nn.Sequential( 49 | torch.nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, bias=False, padding=1), 50 | torch.nn.BatchNorm3d(32), 51 | torch.nn.ReLU(), 52 | ) 53 | self.dconv10 = torch.nn.Sequential( 54 | torch.nn.ConvTranspose3d(32, 1, kernel_size=4, stride=2, bias=False, padding=1), 55 | torch.nn.BatchNorm3d(1), 56 | torch.nn.ReLU(), 57 | ) 58 | self.gridding_rev = GriddingReverse(scale=64) 59 | self.point_sampling = RandomPointSampling(n_points=2048) 60 | self.feature_sampling = CubicFeatureSampling() 61 | self.fc11 = torch.nn.Sequential(torch.nn.Linear(1792, 1792), torch.nn.ReLU()) 62 | self.fc12 = torch.nn.Sequential(torch.nn.Linear(1792, 448), torch.nn.ReLU()) 63 | self.fc13 = torch.nn.Sequential(torch.nn.Linear(448, 112), torch.nn.ReLU()) 64 | self.fc14 = torch.nn.Linear(112, 24) 65 | 66 | def forward(self, data): 67 | partial_cloud = data["partial_cloud"] 68 | # print(partial_cloud.size()) # torch.Size([batch_size, 2048, 3]) 69 | pt_features_64_l = self.gridding(partial_cloud).view(-1, 1, 64, 64, 64) 70 | # print(pt_features_64_l.size()) # torch.Size([batch_size, 1, 64, 64, 64]) 71 | pt_features_32_l = self.conv1(pt_features_64_l) 72 | # print(pt_features_32_l.size()) # torch.Size([batch_size, 32, 32, 32, 32]) 73 | pt_features_16_l = self.conv2(pt_features_32_l) 74 | # print(pt_features_16_l.size()) # torch.Size([batch_size, 64, 16, 16, 16]) 75 | pt_features_8_l = self.conv3(pt_features_16_l) 76 | # print(pt_features_8_l.size()) # torch.Size([batch_size, 128, 8, 8, 8]) 77 | pt_features_4_l = self.conv4(pt_features_8_l) 78 | # print(pt_features_4_l.size()) # torch.Size([batch_size, 256, 4, 4, 4]) 79 | features = self.fc5(pt_features_4_l.view(-1, 16384)) 80 | # print(features.size()) # torch.Size([batch_size, 2048]) 81 | pt_features_4_r = self.fc6(features).view(-1, 256, 4, 4, 4) + pt_features_4_l 82 | # print(pt_features_4_r.size()) # torch.Size([batch_size, 256, 4, 4, 4]) 83 | pt_features_8_r = self.dconv7(pt_features_4_r) + pt_features_8_l 84 | # print(pt_features_8_r.size()) # torch.Size([batch_size, 128, 8, 8, 8]) 85 | pt_features_16_r = self.dconv8(pt_features_8_r) + pt_features_16_l 86 | # print(pt_features_16_r.size()) # torch.Size([batch_size, 64, 16, 16, 16]) 87 | pt_features_32_r = self.dconv9(pt_features_16_r) + pt_features_32_l 88 | # print(pt_features_32_r.size()) # torch.Size([batch_size, 32, 32, 32, 32]) 89 | pt_features_64_r = self.dconv10(pt_features_32_r) + pt_features_64_l 90 | # print(pt_features_64_r.size()) # torch.Size([batch_size, 1, 64, 64, 64]) 91 | sparse_cloud = self.gridding_rev(pt_features_64_r.squeeze(dim=1)) 92 | # print(sparse_cloud.size()) # torch.Size([batch_size, 262144, 3]) 93 | sparse_cloud = self.point_sampling(sparse_cloud, partial_cloud) 94 | # print(sparse_cloud.size()) # torch.Size([batch_size, 2048, 3]) 95 | point_features_32 = self.feature_sampling(sparse_cloud, pt_features_32_r).view(-1, 2048, 256) 96 | # print(point_features_32.size()) # torch.Size([batch_size, 2048, 256]) 97 | point_features_16 = self.feature_sampling(sparse_cloud, pt_features_16_r).view(-1, 2048, 512) 98 | # print(point_features_16.size()) # torch.Size([batch_size, 2048, 512]) 99 | point_features_8 = self.feature_sampling(sparse_cloud, pt_features_8_r).view(-1, 2048, 1024) 100 | # print(point_features_8.size()) # torch.Size([batch_size, 2048, 1024]) 101 | point_features = torch.cat([point_features_32, point_features_16, point_features_8], dim=2) 102 | # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) 103 | point_features = self.fc11(point_features) 104 | # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) 105 | point_features = self.fc12(point_features) 106 | # print(point_features.size()) # torch.Size([batch_size, 2048, 448]) 107 | point_features = self.fc13(point_features) 108 | # print(point_features.size()) # torch.Size([batch_size, 2048, 112]) 109 | point_offset = self.fc14(point_features).view(-1, 16384, 3) 110 | # print(point_features.size()) # torch.Size([batch_size, 16384, 3]) 111 | dense_cloud = sparse_cloud.unsqueeze(dim=2).repeat(1, 1, 8, 1).view(-1, 16384, 3) + point_offset 112 | # print(dense_cloud.size()) # torch.Size([batch_size, 16384, 3]) 113 | 114 | return sparse_cloud, dense_cloud 115 | class RandomPointSampling(torch.nn.Module): 116 | def __init__(self, n_points): 117 | super(RandomPointSampling, self).__init__() 118 | self.n_points = n_points 119 | 120 | def forward(self, pred_cloud, partial_cloud=None): 121 | if partial_cloud is not None: 122 | pred_cloud = torch.cat([partial_cloud, pred_cloud], dim=1) 123 | 124 | _ptcloud = torch.split(pred_cloud, 1, dim=0) 125 | ptclouds = [] 126 | for p in _ptcloud: 127 | non_zeros = torch.sum(p, dim=2).ne(0) 128 | p = p[non_zeros].unsqueeze(dim=0) 129 | n_pts = p.size(1) 130 | if n_pts < self.n_points: 131 | rnd_idx = torch.cat([torch.randint(0, n_pts, (self.n_points,))]) 132 | else: 133 | rnd_idx = torch.randperm(p.size(1))[: self.n_points] 134 | ptclouds.append(p[:, rnd_idx, :]) 135 | 136 | return torch.cat(ptclouds, dim=0).contiguous() -------------------------------------------------------------------------------- /models/msn_generator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.utils.data 7 | from torch.autograd import Variable 8 | import cuda.MDS.MDS_module as MDS_module 9 | import cuda.expansion_penalty.expansion_penalty_module as expansion 10 | from models.atlasnet_generator import PointEncoder 11 | from models.sparenet_generator import PointGenCon, PointNetRes 12 | 13 | 14 | class MSN(nn.Module): 15 | def __init__( 16 | self, 17 | num_points=16382, 18 | bottleneck_size=1024, 19 | n_primitives=32, 20 | ): 21 | super(MSN, self).__init__() 22 | self.num_points = num_points 23 | self.bottleneck_size = bottleneck_size 24 | self.n_primitives = n_primitives 25 | self.expansion = expansion.expansionPenaltyModule() 26 | 27 | self.encoder = PointEncoder(bottleneck_size=self.bottleneck_size) 28 | 29 | self.decoder = nn.ModuleList( 30 | [ 31 | PointGenCon( 32 | input_dim=2 + self.bottleneck_size, 33 | bottleneck_size=2 + self.bottleneck_size, 34 | use_SElayer=False, 35 | ) 36 | for i in range(self.n_primitives) 37 | ] 38 | ) 39 | 40 | self.res = PointNetRes(use_SElayer=False) 41 | 42 | def forward(self, data): 43 | x = data["partial_cloud"] 44 | x = x.transpose(1, 2).contiguous() 45 | partial = x 46 | style = self.encoder(x) 47 | 48 | outs = [] 49 | for i in range(self.n_primitives): 50 | rand_grid = Variable( 51 | torch.cuda.FloatTensor( 52 | x.size(0), 2, self.num_points // self.n_primitives 53 | ) 54 | ) 55 | rand_grid.data.uniform_(0, 1) 56 | y = ( 57 | style.unsqueeze(2) 58 | .expand(style.size(0), style.size(1), rand_grid.size(2)) 59 | .contiguous() 60 | ) 61 | y = torch.cat((rand_grid, y), 1).contiguous() 62 | outs.append(self.decoder[i](y)) 63 | 64 | outs = torch.cat(outs, 2).contiguous() 65 | coarse = outs.transpose(1, 2).contiguous() 66 | dist, _, mean_mst_dis = self.expansion( 67 | coarse, self.num_points // self.n_primitives, 1.5 68 | ) 69 | loss_mst = torch.mean(dist) 70 | id0 = torch.zeros(outs.shape[0], 1, outs.shape[2]).cuda().contiguous() 71 | outs = torch.cat((outs, id0), 1) 72 | id1 = torch.ones(partial.shape[0], 1, partial.shape[2]).cuda().contiguous() 73 | partial = torch.cat((partial, id1), 1) 74 | xx = torch.cat((outs, partial), 2) 75 | resampled_idx = MDS_module.minimum_density_sample( 76 | xx[:, 0:3, :].transpose(1, 2).contiguous(), coarse.shape[1], mean_mst_dis 77 | ) 78 | xx = MDS_module.gather_operation(xx, resampled_idx) 79 | delta = self.res(xx) 80 | xx = xx[:, 0:3, :] 81 | refine = (xx + delta).transpose(2, 1).contiguous() 82 | return coarse, refine, loss_mst 83 | -------------------------------------------------------------------------------- /models/sparenet_discriminator.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.utils.data 7 | import torch.nn.functional as F 8 | from torch.nn import init 9 | from torch.nn import utils 10 | from torch.nn import Parameter 11 | 12 | 13 | class PatchDiscriminator(nn.Module): 14 | """ 15 | inputs: 16 | - img: b x (2*views) x [img_size, img_size] 17 | - feat: bool 18 | - y(label): b 19 | 20 | outputs: 21 | - validity: b 22 | - feat: 23 | - feat_1: b x num_dims(16) x [img_size, img_size](128) 24 | - feat_2: b x num_dims(32) x [img_size, img_size](64) 25 | - feat_3: b x num_dims(64) x [img_size, img_size](32) 26 | - feat_4: b x num_dims(128) x [img_size, img_size](16) 27 | """ 28 | 29 | def __init__(self, img_shape: tuple = (2, 256, 256)): 30 | super(PatchDiscriminator, self).__init__() 31 | 32 | def discriminator_block( 33 | in_filters: int, out_filters: int, normalization: bool = True 34 | ): 35 | """Returns downsampling layers of each discriminator block""" 36 | layers = [ 37 | SpectralNorm(nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)) 38 | ] 39 | if normalization: 40 | layers.append(nn.BatchNorm2d(out_filters)) 41 | layers.append(nn.LeakyReLU(0.2, inplace=True)) 42 | return layers 43 | 44 | self.conv1 = nn.Sequential( 45 | *discriminator_block(img_shape[0], 16, normalization=False), 46 | ) 47 | self.conv2 = nn.Sequential( 48 | *discriminator_block(16, 32), 49 | ) 50 | self.conv3 = nn.Sequential( 51 | *discriminator_block(32, 64), 52 | ) 53 | self.conv4 = nn.Sequential( 54 | *discriminator_block(64, 128), 55 | ) 56 | self.conv5 = nn.Sequential( 57 | *discriminator_block(128, 256), 58 | ) 59 | self.conv6 = nn.Sequential( 60 | *discriminator_block(256, 512), 61 | ) 62 | self.adv_layer = SpectralNorm(nn.Conv2d(512, 1, 3, padding=1, bias=False)) 63 | 64 | def forward(self, img, feat=False, y=None): 65 | 66 | feat_1 = self.conv1(img) # [bs, 16, 128, 128] 67 | feat_2 = self.conv2(feat_1) # [bs, 32, 64, 64] 68 | feat_3 = self.conv3(feat_2) # [bs, 64, 32, 32] 69 | feat_4 = self.conv4(feat_3) # [bs, 128, 16, 16] 70 | feat_5 = self.conv5(feat_4) # [bs, 256, 8, 8] 71 | feat_6 = self.conv6(feat_5) # [bs, 512, 4, 4] 72 | validity = self.adv_layer(feat_6) # [bs, 1, 4, 4] 73 | 74 | validity = F.avg_pool2d(validity, validity.size()[2:]).view( 75 | validity.size()[0], -1 76 | ) 77 | 78 | if feat: 79 | return validity, [feat_1, feat_2, feat_3, feat_4] 80 | else: 81 | return validity 82 | 83 | 84 | class ProjectionD(nn.Module): 85 | """ 86 | inputs: 87 | - img: b x (2*views) x [img_size, img_size] 88 | - feat: bool 89 | - y(label): b 90 | 91 | outputs: 92 | - validity: b 93 | - feat: 94 | - feat_1: b x num_dims(16) x [img_size, img_size](128) 95 | - feat_2: b x num_dims(32) x [img_size, img_size](64) 96 | - feat_3: b x num_dims(64) x [img_size, img_size](32) 97 | - feat_4: b x num_dims(128) x [img_size, img_size](16) 98 | """ 99 | 100 | def __init__(self, num_classes: int = 0, img_shape: tuple = (2, 256, 256)): 101 | super(ProjectionD, self).__init__() 102 | 103 | def discriminator_block(in_filters: int, out_filters: int, bn: bool = True): 104 | block = [ 105 | SpectralNorm(nn.Conv2d(in_filters, out_filters, 3, 2, 1)), 106 | nn.LeakyReLU(0.2, inplace=True), 107 | nn.Dropout2d(0.25), 108 | ] 109 | if bn: 110 | block.append(nn.BatchNorm2d(out_filters, 0.8)) 111 | return block 112 | 113 | self.conv1 = nn.Sequential( 114 | *discriminator_block(img_shape[0], 16, bn=False), 115 | ) 116 | self.conv2 = nn.Sequential( 117 | *discriminator_block(16, 32), 118 | ) 119 | self.conv3 = nn.Sequential( 120 | *discriminator_block(32, 64), 121 | ) 122 | feat_num = 128 123 | self.conv4 = nn.Sequential( 124 | *discriminator_block(64, feat_num), 125 | ) 126 | 127 | ds_size = img_shape[1] // (2 ** 4) 128 | self.adv_layer = utils.spectral_norm(nn.Linear(feat_num * ds_size ** 2, 1)) 129 | if num_classes > 0: 130 | self.l_y = utils.spectral_norm( 131 | nn.Embedding(num_classes, feat_num * ds_size ** 2) 132 | ) 133 | self._initialize() 134 | 135 | def _initialize(self): 136 | init.xavier_uniform_(self.adv_layer.weight.data) 137 | optional_l_y = getattr(self, "l_y", None) 138 | if optional_l_y is not None: 139 | init.xavier_uniform_(optional_l_y.weight.data) 140 | 141 | def forward(self, img, feat=False, y=None): 142 | feat_1 = self.conv1(img) # [bs, 16, 128, 128] 143 | feat_2 = self.conv2(feat_1) # [bs, 32, 64, 64] 144 | feat_3 = self.conv3(feat_2) # [bs, 64, 32, 32] 145 | feat_4 = self.conv4(feat_3) # [bs, 128, 16, 16] 146 | out = feat_4.view(feat_4.shape[0], -1) 147 | validity = self.adv_layer(out) 148 | if y is not None: 149 | validity += torch.sum(self.l_y(y) * out, dim=1, keepdim=True) 150 | if feat: 151 | return validity, [feat_1, feat_2, feat_3, feat_4] 152 | else: 153 | return validity 154 | 155 | 156 | def l2normalize(v, eps=1e-12): 157 | return v / (v.norm() + eps) 158 | 159 | 160 | class SpectralNorm(nn.Module): 161 | def __init__(self, module, name="weight", power_iterations=1): 162 | super(SpectralNorm, self).__init__() 163 | self.module = module 164 | self.name = name 165 | self.power_iterations = power_iterations 166 | if not self._made_params(): 167 | self._make_params() 168 | 169 | def _update_u_v(self): 170 | u = getattr(self.module, self.name + "_u") 171 | v = getattr(self.module, self.name + "_v") 172 | w = getattr(self.module, self.name + "_bar") 173 | 174 | height = w.data.shape[0] 175 | for _ in range(self.power_iterations): 176 | v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data)) 177 | u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data)) 178 | 179 | sigma = u.dot(w.view(height, -1).mv(v)) 180 | setattr(self.module, self.name, w / sigma.expand_as(w)) 181 | 182 | def _made_params(self): 183 | try: 184 | u = getattr(self.module, self.name + "_u") 185 | v = getattr(self.module, self.name + "_v") 186 | w = getattr(self.module, self.name + "_bar") 187 | return True 188 | except AttributeError: 189 | return False 190 | 191 | def _make_params(self): 192 | w = getattr(self.module, self.name) 193 | 194 | height = w.data.shape[0] 195 | width = w.view(height, -1).data.shape[1] 196 | 197 | u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) 198 | v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) 199 | u.data = l2normalize(u.data) 200 | v.data = l2normalize(v.data) 201 | w_bar = Parameter(w.data) 202 | 203 | del self.module._parameters[self.name] 204 | 205 | self.module.register_parameter(self.name + "_u", u) 206 | self.module.register_parameter(self.name + "_v", v) 207 | self.module.register_parameter(self.name + "_bar", w_bar) 208 | 209 | def forward(self, *args): 210 | self._update_u_v() 211 | return self.module.forward(*args) 212 | -------------------------------------------------------------------------------- /renderer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/renderer.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | argparse 2 | easydict 3 | h5py 4 | matplotlib 5 | numpy 6 | pyyaml 7 | open3d==0.9.0.0 8 | open3d-python 9 | opencv-python 10 | scipy 11 | tensorboardX>=1.2 12 | torchvision>=0.4.0 13 | torch>=1.2.0 14 | transforms3d 15 | tqdm 16 | pillow>=6.2.0 17 | geomloss==0.2.1 -------------------------------------------------------------------------------- /runners/atlasnet_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import logging 7 | from time import time 8 | import utils.misc as um 9 | import cuda.emd.emd_module as emd 10 | from cuda.chamfer_distance import ChamferDistanceMean 11 | from runners.misc import AverageMeter 12 | from runners.base_runner import BaseRunner 13 | 14 | 15 | class atlasnetRunner(BaseRunner): 16 | """Define the AtlasNet runner class""" 17 | 18 | def __init__(self, config, logger): 19 | super().__init__(config, logger) 20 | self.losses = AverageMeter(["RefineLoss"]) 21 | self.test_losses = AverageMeter(["RefineLoss"]) 22 | self.test_metrics = AverageMeter(um.Metrics.names()) 23 | self.chamfer_dist_mean = None 24 | self.emd_dist = None 25 | 26 | def build_models(self): 27 | super().build_models() 28 | 29 | def build_train_loss(self): 30 | self.chamfer_dist_mean = torch.nn.DataParallel( 31 | ChamferDistanceMean().to(self.gpu_ids[0]), device_ids=self.gpu_ids 32 | ) 33 | self.emd_dist = torch.nn.DataParallel( 34 | emd.emdModule().to(self.gpu_ids[0]), device_ids=self.gpu_ids 35 | ) 36 | 37 | def build_val_loss(self): 38 | self.chamfer_dist_mean = ChamferDistanceMean().cuda() 39 | self.emd_dist = emd.emdModule().cuda() 40 | 41 | def train_step(self, items): 42 | _, (_, _, _, data) = items 43 | for k, v in data.items(): 44 | data[k] = v.float().to(self.gpu_ids[0]) 45 | 46 | _loss, _, refine_loss = self.completion(data) 47 | self.models.zero_grad() 48 | _loss.backward() 49 | self.optimizers.step() 50 | 51 | self.loss["refine_loss"] = refine_loss * 1000 52 | self.loss["rec_loss"] = _loss 53 | self.losses.update([refine_loss.item() * 1000]) 54 | 55 | def val_step(self, items): 56 | _, (_, _, _, data) = items 57 | for k, v in data.items(): 58 | data[k] = um.var_or_cuda(v) 59 | 60 | _, refine_ptcloud, refine_loss = self.completion(data) 61 | self.test_losses.update([refine_loss.item() * 1000]) 62 | self.metrics = um.Metrics.get(refine_ptcloud, data["gtcloud"]) 63 | self.ptcloud = refine_ptcloud 64 | 65 | def completion(self, data): 66 | """ 67 | inputs: 68 | cfg: EasyDict 69 | data: tensor 70 | -partical_cloud: b x npoints1 x num_dims 71 | -gtcloud: b x npoints2 x num_dims 72 | 73 | outputs: 74 | _loss 75 | refine_ptcloud: b x npoints2 x num_dims 76 | refine_loss: float32 77 | """ 78 | refine_ptcloud = self.models(data) 79 | 80 | if self.config.NETWORK.metric == "chamfer": 81 | refine_loss = self.chamfer_dist_mean(refine_ptcloud, data["gtcloud"]).mean() 82 | 83 | elif self.config.NETWORK.metric == "emd": 84 | emd_refine, _ = self.emd_dist( 85 | refine_ptcloud, data["gtcloud"], eps=0.005, iters=50 86 | ) 87 | refine_loss = torch.sqrt(emd_refine).mean(1).mean() 88 | 89 | else: 90 | raise Exception("unknown training metric") 91 | 92 | _loss = refine_loss 93 | 94 | return _loss, refine_ptcloud, _loss -------------------------------------------------------------------------------- /runners/grnet_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import logging 7 | from time import time 8 | import utils.misc as um 9 | import cuda.emd.emd_module as emd 10 | from cuda.chamfer_distance import ChamferDistanceMean 11 | from runners.misc import AverageMeter 12 | from runners.base_runner import BaseRunner 13 | 14 | 15 | class grnetRunner(BaseRunner): 16 | """Define the AtlasNet runner class""" 17 | 18 | def __init__(self, config, logger): 19 | super().__init__(config, logger) 20 | self.losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 21 | self.test_losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 22 | self.test_metrics = AverageMeter(um.Metrics.names()) 23 | self.chamfer_dist_mean = None 24 | self.emd_dist = None 25 | 26 | def build_models(self): 27 | super().build_models() 28 | 29 | def build_train_loss(self): 30 | self.chamfer_dist_mean = torch.nn.DataParallel( 31 | ChamferDistanceMean().to(self.gpu_ids[0]), device_ids=self.gpu_ids 32 | ) 33 | self.emd_dist = torch.nn.DataParallel( 34 | emd.emdModule().to(self.gpu_ids[0]), device_ids=self.gpu_ids 35 | ) 36 | 37 | def build_val_loss(self): 38 | self.chamfer_dist_mean = ChamferDistanceMean().cuda() 39 | self.emd_dist = emd.emdModule().cuda() 40 | 41 | def train_step(self, items): 42 | _, (_, _, _, data) = items 43 | for k, v in data.items(): 44 | data[k] = v.float().to(self.gpu_ids[0]) 45 | _loss, _, _, refine_loss, coarse_loss = self.completion(data) 46 | self.models.zero_grad() 47 | _loss.backward() 48 | self.optimizers.step() 49 | 50 | self.loss["coarse_loss"] = coarse_loss * 1000 51 | self.loss["refine_loss"] = refine_loss * 1000 52 | self.loss["rec_loss"] = _loss 53 | self.losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 54 | 55 | def val_step(self, items): 56 | _, (_, _, _, data) = items 57 | for k, v in data.items(): 58 | data[k] = um.var_or_cuda(v) 59 | _, _, refine_ptcloud, coarse_loss, refine_loss = self.completion(data) 60 | self.test_losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 61 | self.metrics = um.Metrics.get(refine_ptcloud, data["gtcloud"]) 62 | self.ptcloud = refine_ptcloud 63 | 64 | def completion(self, data): 65 | """ 66 | inputs: 67 | cfg: EasyDict 68 | data: tensor 69 | -partical_cloud: b x npoints1 x num_dims 70 | -gtcloud: b x npoints2 x num_dims 71 | 72 | outputs: 73 | _loss: float32 74 | refine_ptcloud: b x npoints2 x num_dims 75 | coarse_ptcloud: b x npoints2 x num_dims 76 | refine_loss: float32 77 | coarse_loss: float32 78 | """ 79 | (coarse_ptcloud, refine_ptcloud) = self.models(data) 80 | coarse_loss = self.chamfer_dist_mean(coarse_ptcloud, data["gtcloud"]).mean() 81 | 82 | if self.config.NETWORK.metric == "chamfer": 83 | refine_loss = self.chamfer_dist_mean(refine_ptcloud, data["gtcloud"]).mean() 84 | elif self.config.NETWORK.metric == "emd": 85 | emd_refine, _ = self.emd_dist( 86 | refine_ptcloud, data["gtcloud"], eps=0.005, iters=50 87 | ) 88 | refine_loss = torch.sqrt(emd_refine).mean(1).mean() 89 | else: 90 | raise Exception("unknown training metric") 91 | 92 | _loss = coarse_loss + refine_loss 93 | 94 | return _loss, coarse_ptcloud, refine_ptcloud, coarse_loss, refine_loss 95 | -------------------------------------------------------------------------------- /runners/misc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | class AverageMeter(object): 5 | """Computes and stores the average and current value""" 6 | 7 | def __init__(self, items=None): 8 | self.items = items 9 | self.n_items = 1 if items is None else len(items) 10 | self.reset() 11 | 12 | def reset(self): 13 | self._val = [0] * self.n_items 14 | self._sum = [0] * self.n_items 15 | self._count = [0] * self.n_items 16 | 17 | def update(self, values): 18 | if type(values).__name__ == "list": 19 | for idx, v in enumerate(values): 20 | self._val[idx] = v 21 | self._sum[idx] += v 22 | self._count[idx] += 1 23 | else: 24 | self._val[0] = values 25 | self._sum[0] += values 26 | self._count[0] += 1 27 | 28 | def val(self, idx=None): 29 | if idx is None: 30 | return self._val[0] if self.items is None else [self._val[i] for i in range(self.n_items)] 31 | else: 32 | return self._val[idx] 33 | 34 | def count(self, idx=None): 35 | if idx is None: 36 | return self._count[0] if self.items is None else [self._count[i] for i in range(self.n_items)] 37 | else: 38 | return self._count[idx] 39 | 40 | def avg(self, idx=None): 41 | if idx is None: 42 | return self._sum[0] / self._count[0] if self.items is None else [self._sum[i] / self._count[i] for i in range(self.n_items)] 43 | else: 44 | return self._sum[idx] / self._count[idx] -------------------------------------------------------------------------------- /runners/msn_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import logging 7 | from time import time 8 | import utils.misc as um 9 | import cuda.emd.emd_module as emd 10 | from cuda.chamfer_distance import ChamferDistanceMean 11 | from runners.misc import AverageMeter 12 | from runners.base_runner import BaseRunner 13 | 14 | class msnRunner(BaseRunner): 15 | """Define the AtlasNet runner class""" 16 | 17 | def __init__(self, config, logger): 18 | super().__init__(config, logger) 19 | self.losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 20 | self.test_losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 21 | self.test_metrics = AverageMeter(um.Metrics.names()) 22 | self.chamfer_dist_mean = None 23 | self.emd_dist = None 24 | 25 | def build_models(self): 26 | super().build_models() 27 | 28 | def build_train_loss(self): 29 | self.chamfer_dist_mean = torch.nn.DataParallel(ChamferDistanceMean().to(self.gpu_ids[0]), device_ids=self.gpu_ids) 30 | self.emd_dist = torch.nn.DataParallel(emd.emdModule().to(self.gpu_ids[0]), device_ids=self.gpu_ids) 31 | 32 | def build_val_loss(self): 33 | self.chamfer_dist_mean = ChamferDistanceMean().cuda() 34 | self.emd_dist = emd.emdModule().cuda() 35 | 36 | def train_step(self, items): 37 | _, (_, _, _, data) = items 38 | for k, v in data.items(): 39 | data[k] = v.float().to(self.gpu_ids[0]) 40 | 41 | _loss, _, _, refine_loss, coarse_loss = self.completion(data) 42 | self.models.zero_grad() 43 | _loss.backward() 44 | self.optimizers.step() 45 | 46 | self.loss["coarse_loss"] = coarse_loss * 1000 47 | self.loss["refine_loss"] = refine_loss * 1000 48 | self.loss["rec_loss"] = _loss 49 | self.losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 50 | 51 | def val_step(self, items): 52 | _, (_, _, _, data) = items 53 | for k, v in data.items(): 54 | data[k] = um.var_or_cuda(v) 55 | 56 | _, _, refine_ptcloud, coarse_loss, refine_loss = self.completion(data) 57 | self.test_losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 58 | self.metrics = um.Metrics.get(refine_ptcloud, data["gtcloud"]) 59 | self.ptcloud = refine_ptcloud 60 | 61 | def completion(self, data): 62 | """ 63 | inputs: 64 | cfg: EasyDict 65 | data: tensor 66 | -partical_cloud: b x npoints1 x num_dims 67 | -gtcloud: b x npoints2 x num_dims 68 | 69 | outputs: 70 | _loss: float32 71 | refine_ptcloud: b x npoints2 x num_dims 72 | coarse_ptcloud: b x npoints2 x num_dims 73 | refine_loss: float32 74 | coarse_loss: float32 75 | """ 76 | coarse_ptcloud, refine_ptcloud, expansion_penalty = self.models(data) 77 | 78 | if self.config.NETWORK.metric == "chamfer": 79 | coarse_loss = self.chamfer_dist_mean(coarse_ptcloud, data["gtcloud"]).mean() 80 | refine_loss = self.chamfer_dist_mean(refine_ptcloud, data["gtcloud"]).mean() 81 | 82 | elif self.config.NETWORK.metric == "emd": 83 | emd_coarse, _ = self.emd_dist(coarse_ptcloud, data["gtcloud"], eps=0.005, iters=50) 84 | emd_refine, _ = self.emd_dist(refine_ptcloud, data["gtcloud"], eps=0.005, iters=50) 85 | coarse_loss = torch.sqrt(emd_coarse).mean(1).mean() 86 | refine_loss = torch.sqrt(emd_refine).mean(1).mean() 87 | 88 | else: 89 | raise Exception("unknown training metric") 90 | 91 | _loss = coarse_loss + refine_loss + expansion_penalty.mean() * 0.1 92 | 93 | return _loss, coarse_ptcloud, refine_ptcloud, coarse_loss, refine_loss 94 | -------------------------------------------------------------------------------- /runners/sparenet_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import logging 7 | from time import time 8 | import utils.misc as um 9 | import cuda.emd.emd_module as emd 10 | from cuda.chamfer_distance import ChamferDistance, ChamferDistanceMean 11 | from runners.misc import AverageMeter 12 | from runners.base_runner import BaseRunner 13 | 14 | 15 | class sparenetRunner(BaseRunner): 16 | """Define the SpareNet runner class""" 17 | 18 | def __init__(self, config, logger): 19 | super().__init__(config, logger) 20 | self.losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 21 | self.test_losses = AverageMeter(["CoarseLoss", "RefineLoss"]) 22 | self.test_metrics = AverageMeter(um.Metrics.names()) 23 | self.chamfer_dist = None 24 | self.chamfer_dist_mean = None 25 | self.emd_dist = None 26 | 27 | def build_models(self): 28 | super().build_models() 29 | 30 | def build_train_loss(self): 31 | # Set up loss functions 32 | self.chamfer_dist = torch.nn.DataParallel(ChamferDistance().to(self.gpu_ids[0]), device_ids=self.gpu_ids) 33 | self.chamfer_dist_mean = torch.nn.DataParallel(ChamferDistanceMean().to(self.gpu_ids[0]), device_ids=self.gpu_ids) 34 | self.emd_dist = torch.nn.DataParallel(emd.emdModule().to(self.gpu_ids[0]), device_ids=self.gpu_ids) 35 | 36 | def build_val_loss(self): 37 | # Set up loss functions 38 | self.chamfer_dist = ChamferDistance().cuda() 39 | self.chamfer_dist_mean = ChamferDistanceMean().cuda() 40 | self.emd_dist = emd.emdModule().cuda() 41 | 42 | def train_step(self, items): 43 | _, (_, _, _, data) = items 44 | for k, v in data.items(): 45 | data[k] = v.float().to(self.gpu_ids[0]) 46 | 47 | _loss, _, _, _, refine_loss, coarse_loss = self.completion(data) 48 | self.models.zero_grad() 49 | _loss.backward() 50 | self.optimizers.step() 51 | 52 | self.loss["coarse_loss"] = coarse_loss * 1000 53 | self.loss["refine_loss"] = refine_loss * 1000 54 | self.loss["rec_loss"] = _loss 55 | self.losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 56 | 57 | def val_step(self, items): 58 | _, (_, _, _, data) = items 59 | for k, v in data.items(): 60 | data[k] = um.var_or_cuda(v) 61 | 62 | _, refine_ptcloud, _, _, refine_loss, coarse_loss = self.completion(data) 63 | self.test_losses.update([coarse_loss.item() * 1000, refine_loss.item() * 1000]) 64 | self.metrics = um.Metrics.get(refine_ptcloud, data["gtcloud"]) 65 | self.ptcloud = refine_ptcloud 66 | 67 | def completion(self, data): 68 | """ 69 | inputs: 70 | cfg: EasyDict 71 | data: tensor 72 | -partical_cloud: b x npoints1 x num_dims 73 | -gtcloud: b x npoints2 x num_dims 74 | 75 | outputs: 76 | _loss: float32 77 | refine_ptcloud: b x npoints2 x num_dims 78 | middle_ptcloud: b x npoints2 x num_dims 79 | coarse_ptcloud: b x npoints2 x num_dims 80 | refine_loss: float32 81 | coarse_loss: float32 82 | """ 83 | (coarse_ptcloud, middle_ptcloud, refine_ptcloud, expansion_penalty) = self.models(data) 84 | 85 | if self.config.NETWORK.metric == "chamfer": 86 | coarse_loss = self.chamfer_dist_mean(coarse_ptcloud, data["gtcloud"]).mean() 87 | middle_loss = self.chamfer_dist_mean(middle_ptcloud, data["gtcloud"]).mean() 88 | refine_loss = self.chamfer_dist_mean(refine_ptcloud, data["gtcloud"]).mean() 89 | 90 | elif self.config.NETWORK.metric == "emd": 91 | emd_coarse, _ = self.emd_dist(coarse_ptcloud, data["gtcloud"], eps=0.005, iters=50) 92 | emd_middle, _ = self.emd_dist(middle_ptcloud, data["gtcloud"], eps=0.005, iters=50) 93 | emd_refine, _ = self.emd_dist(refine_ptcloud, data["gtcloud"], eps=0.005, iters=50) 94 | coarse_loss = torch.sqrt(emd_coarse).mean(1).mean() 95 | refine_loss = torch.sqrt(emd_refine).mean(1).mean() 96 | middle_loss = torch.sqrt(emd_middle).mean(1).mean() 97 | 98 | else: 99 | raise Exception("unknown training metric") 100 | 101 | _loss = coarse_loss + middle_loss + refine_loss + expansion_penalty.mean() * 0.1 102 | 103 | if self.config.NETWORK.use_consist_loss: 104 | dist1, _ = self.chamfer_dist(refine_ptcloud, data["gtcloud"]) 105 | cd_input2fine = torch.mean(dist1).mean() 106 | _loss += cd_input2fine * 0.5 107 | 108 | return _loss, refine_ptcloud, middle_ptcloud, coarse_ptcloud, refine_loss, coarse_loss -------------------------------------------------------------------------------- /scripts/atlasnet_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python test.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model atlasnet \ 5 | --weights ${path to checkpoint} \ 6 | --test_mode ${mode} -------------------------------------------------------------------------------- /scripts/atlasnet_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python train.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model atlasnet \ 5 | --weights ${path to checkpoint} -------------------------------------------------------------------------------- /scripts/grnet_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python test.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model grnet \ 5 | --weights ${path to checkpoint} \ 6 | --test_mode ${mode} -------------------------------------------------------------------------------- /scripts/grnet_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python train.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model grnet \ 5 | --weights ${path to checkpoint} -------------------------------------------------------------------------------- /scripts/msn_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python test.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model msn \ 5 | --weights ${path to checkpoint} \ 6 | --test_mode ${mode} -------------------------------------------------------------------------------- /scripts/msn_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python train.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model msn \ 5 | --weights ${path to checkpoint} -------------------------------------------------------------------------------- /scripts/sparenet_gan_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python train.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model sparenet \ 5 | --weights ${path to checkpoint} \ 6 | --gan # use gan -------------------------------------------------------------------------------- /scripts/sparenet_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python test.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model sparenet \ 5 | --weights ${path to checkpoint} \ 6 | --test_mode ${mode} -------------------------------------------------------------------------------- /scripts/sparenet_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python train.py --gpu ${GPUS}\ 3 | --work_dir ${WORK_DIR} \ 4 | --model sparenet \ 5 | --weights ${path to checkpoint} -------------------------------------------------------------------------------- /setup_env.sh: -------------------------------------------------------------------------------- 1 | sudo apt-get update 2 | sudo apt-get --yes install vim 3 | sudo apt-get --yes install tmux 4 | 5 | pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl --user 6 | 7 | PROJ_HOME=`pwd` 8 | 9 | cd $PROJ_HOME/cuda/emd/ 10 | rm -rf build/* 11 | python setup.py install --user 12 | 13 | cd $PROJ_HOME/cuda/expansion_penalty/ 14 | rm -rf build/* 15 | python setup.py install --user 16 | 17 | cd $PROJ_HOME/cuda/MDS/ 18 | rm -rf build/* 19 | python setup.py install --user 20 | 21 | cd $PROJ_HOME/cuda/cubic_feature_sampling/ 22 | rm -rf build/* 23 | python setup.py install --user 24 | 25 | cd $PROJ_HOME/cuda/gridding/ 26 | rm -rf build/* 27 | python setup.py install --user 28 | 29 | cd $PROJ_HOME/cuda/gridding_loss/ 30 | rm -rf build/* 31 | python setup.py install --user 32 | cd $PROJ_HOME 33 | -------------------------------------------------------------------------------- /teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/teaser.png -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import sys 6 | import argparse 7 | 8 | 9 | def get_args_from_command_line(): 10 | """ 11 | config the parameter 12 | """ 13 | parser = argparse.ArgumentParser(description="The argument parser of R2Net runner") 14 | 15 | # choose model 16 | parser.add_argument("--model", type=str, default="sparenet", help="sparenet, atlasnet, msn, grnet") 17 | 18 | # choose test mode 19 | parser.add_argument("--test_mode", default="default", help="default, vis, render, kitti", type=str) 20 | 21 | # choose load model 22 | parser.add_argument("--weights", dest="weights", help="Initialize network from the weights file", default=None) 23 | 24 | # setup gpu 25 | parser.add_argument("--gpu", dest="gpu_id", help="GPU device to use", default="0", type=str) 26 | 27 | # setup workdir 28 | parser.add_argument("--workdir", dest="workdir", help="where to save files", default="./output", type=str) 29 | 30 | # choose train mode 31 | parser.add_argument("--gan", dest="gan", help="use gan", action="store_true", default=False) 32 | return parser.parse_args() 33 | 34 | 35 | def main(): 36 | args = get_args_from_command_line() 37 | 38 | # Set GPU to use 39 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 40 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id 41 | 42 | # update config 43 | from configs.base_config import cfg, cfg_from_file, cfg_update 44 | 45 | if args.gan: 46 | cfg_from_file("configs/" + args.model + "_gan.yaml") 47 | else: 48 | cfg_from_file("configs/" + args.model + ".yaml") 49 | if args.test_mode is not None: 50 | cfg.TEST.mode = args.test_mode 51 | output_dir = cfg_update(args) 52 | 53 | # Set up folders for logs and checkpoints 54 | if not os.path.exists(cfg.DIR.logs): 55 | os.makedirs(cfg.DIR.logs) 56 | from utils.misc import set_logger 57 | 58 | logger = set_logger(os.path.join(cfg.DIR.logs, "log.txt")) 59 | logger.info("save into dir: %s" % cfg.DIR.logs) 60 | 61 | if "weights" not in cfg.CONST or not os.path.exists(cfg.CONST.weights): 62 | logger.error("Please specify the file path of checkpoint.") 63 | sys.exit(2) 64 | 65 | # Start inference process 66 | if args.gan: 67 | runners = __import__("runners." + args.model + "_gan_runner") 68 | module = getattr(runners, args.model + "_gan_runner") 69 | model = getattr(module, args.model + "GANRunner")(cfg, logger) 70 | 71 | else: 72 | runners = __import__("runners." + args.model + "_runner") 73 | module = getattr(runners, args.model + "_runner") 74 | model = getattr(module, args.model + "Runner")(cfg, logger) 75 | 76 | model.test() 77 | 78 | 79 | if __name__ == "__main__": 80 | main() 81 | -------------------------------------------------------------------------------- /test_fpd.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import argparse 3 | import logging 4 | import torch 5 | import os 6 | import numpy as np 7 | from Frechet.FPD import calculate_fpd 8 | logger = logging.getLogger() 9 | 10 | def read_pcd(filename): 11 | pcd = o3d.io.read_point_cloud(filename) 12 | return torch.from_numpy(np.array(pcd.points)).float() 13 | 14 | 15 | def save_pcd(filename, points): 16 | pcd = o3d.geometry.PointCloud() 17 | pcd.points = o3d.utility.Vector3dVector(points) 18 | o3d.io.write_point_cloud(filename, pcd) 19 | 20 | def set_logger(filename): 21 | logger.setLevel(logging.DEBUG) 22 | formatter = logging.Formatter("%(levelname)s: - %(message)s") 23 | 24 | fh = logging.FileHandler(filename) 25 | fh.setLevel(logging.DEBUG) 26 | fh.setFormatter(formatter) 27 | 28 | ch = logging.StreamHandler() 29 | ch.setLevel(logging.DEBUG) 30 | ch.setFormatter(formatter) 31 | 32 | logger.addHandler(ch) 33 | logger.addHandler(fh) 34 | 35 | 36 | 37 | parser = argparse.ArgumentParser() 38 | parser.add_argument('--plot_freq', type=int, default=1) 39 | parser.add_argument('--save_pcd', action='store_true',default=False) 40 | parser.add_argument('--log_dir', default='/path/to/save/logs') 41 | parser.add_argument('--list_path', default='Frechet/test.list') 42 | parser.add_argument('--data_dir', default='/path/to/test/dataset/pcds') 43 | parser.add_argument('--fake_dir', default='/path/to/methods/pcds', 44 | help='/path/to/results/shapenet_fc/pcds/') 45 | parser.add_argument('--num_points', type=int, default=16384, help='number of points: 2048 or 8192') 46 | 47 | 48 | 49 | opt = parser.parse_args() 50 | os.makedirs(opt.log_dir, exist_ok=True) 51 | 52 | set_logger(os.path.join(opt.log_dir, "log.txt")) 53 | logger.info("save into dir: %s" % opt.log_dir) 54 | 55 | 56 | with open(opt.list_path) as file: 57 | model_list = file.read().splitlines() 58 | 59 | 60 | # Testing loop 61 | n_samples = len(model_list) 62 | logger.info("n_samples %s"% n_samples) 63 | 64 | 65 | taxonomy2label = { 66 | '02691156':0, 67 | '02933112':1, 68 | '02958343':2, 69 | '03001627':3, 70 | '03636649':4, 71 | '04256520':5, 72 | '04379243':6, 73 | '04530566':7 74 | } 75 | 76 | 77 | label2taxonomy = { 78 | 0:'02691156', 79 | 1:'02933112', 80 | 2:'02958343', 81 | 3:'03001627', 82 | 4:'03636649', 83 | 5:'04256520', 84 | 6:'04379243', 85 | 7:'04530566' 86 | } 87 | 88 | fpd_values= [] 89 | 90 | # test for each category 91 | for batch_idx in range(0,8): 92 | fake_pointclouds = torch.Tensor([]).cuda() 93 | real_pointclouds = torch.Tensor([]).cuda() 94 | # prepare the batch data 95 | for idx, model_id in enumerate(model_list): 96 | taxonomy_id, model_id_real = model_id.split('/') 97 | if taxonomy2label[taxonomy_id] == batch_idx: 98 | fake = torch.zeros((1, opt.num_points, 3), device='cuda') 99 | gt = torch.zeros((1, opt.num_points, 3), device='cuda') 100 | pcd = o3d.io.read_point_cloud(os.path.join(opt.fake_dir, '%s.pcd' % model_id)) 101 | fake[0, :, :] = torch.from_numpy(np.array(pcd.points)) 102 | pcd = o3d.io.read_point_cloud(os.path.join(opt.data_dir, 'complete', '%s.pcd' % model_id)) 103 | gt[0, :, :] = torch.from_numpy(np.array(pcd.points)) 104 | fake_pointclouds = torch.cat((fake_pointclouds, fake), dim=0) 105 | real_pointclouds = torch.cat((real_pointclouds, gt), dim=0) 106 | 107 | # 150 data samples per batch 108 | fpd = calculate_fpd(fake_pointclouds, real_pointclouds, statistic_save_path=None, batch_size=30, dims=1808, device= fake_pointclouds.device) 109 | fpd_values.append(fpd) 110 | logger.info("[ %s category] Frechet Pointcloud Distance <<< %f >>>"% (label2taxonomy[batch_idx], fpd)) 111 | 112 | del fake_pointclouds 113 | del real_pointclouds 114 | 115 | logger.info("avg FPD for 8 categories: %f" % np.mean(fpd_values)) 116 | 117 | # test the all test data 118 | 119 | 120 | fake_pointclouds = torch.Tensor([]).cuda() 121 | real_pointclouds = torch.Tensor([]).cuda() 122 | # prepare the batch data 123 | for idx, model_id in enumerate(model_list): 124 | taxonomy_id, model_id_real = model_id.split('/') 125 | fake = torch.zeros((1, opt.num_points, 3), device='cuda') 126 | gt = torch.zeros((1, opt.num_points, 3), device='cuda') 127 | pcd = o3d.io.read_point_cloud(os.path.join(opt.fake_dir, '%s.pcd' % model_id)) 128 | fake[0, :, :] = torch.from_numpy(np.array(pcd.points)) 129 | pcd = o3d.io.read_point_cloud(os.path.join(opt.data_dir, 'complete', '%s.pcd' % model_id)) 130 | gt[0, :, :] = torch.from_numpy(np.array(pcd.points)) 131 | fake_pointclouds = torch.cat((fake_pointclouds, fake), dim=0) 132 | real_pointclouds = torch.cat((real_pointclouds, gt), dim=0) 133 | 134 | # 150 data samples per batch 135 | fpd = calculate_fpd(fake_pointclouds, real_pointclouds, statistic_save_path=None, batch_size=30, dims=1808, device= fake_pointclouds.device) 136 | logger.info("[ all category] Frechet Pointcloud Distance <<< %f >>>"% fpd) 137 | 138 | 139 | 140 | 141 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import argparse 6 | 7 | 8 | def get_args_from_command_line(): 9 | """ 10 | config the parameter 11 | """ 12 | parser = argparse.ArgumentParser(description="The argument parser of R2Net runner") 13 | 14 | # choose model 15 | parser.add_argument("--model", type=str, default="sparenet", help="sparenet, atlasnet, msn, grnet") 16 | 17 | # choose train mode 18 | parser.add_argument("--gan", dest="gan", help="use gan", action="store_true", default=False) 19 | 20 | # choose load model 21 | parser.add_argument("--weights", dest="weights", help="Initialize network from the weights file", default=None) 22 | 23 | # setup gpu 24 | parser.add_argument("--gpu", dest="gpu_id", help="GPU device to use", default="0", type=str) 25 | 26 | # setup workdir 27 | parser.add_argument("--workdir", dest="workdir", help="where to save files", default=None) 28 | return parser.parse_args() 29 | 30 | 31 | def main(): 32 | args = get_args_from_command_line() 33 | 34 | # Set GPU to use 35 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 36 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id 37 | 38 | # update config 39 | from configs.base_config import cfg, cfg_from_file, cfg_update 40 | 41 | if args.gan: 42 | cfg_from_file("configs/" + args.model + "_gan.yaml") 43 | else: 44 | cfg_from_file("configs/" + args.model + ".yaml") 45 | output_dir = cfg_update(args) 46 | 47 | # Set up folders for logs and checkpoints 48 | if not os.path.exists(cfg.DIR.logs): 49 | os.makedirs(cfg.DIR.logs) 50 | from utils.misc import set_logger 51 | 52 | logger = set_logger(os.path.join(cfg.DIR.logs, "log.txt")) 53 | logger.info("save into dir: %s" % cfg.DIR.logs) 54 | 55 | # Start train/inference process 56 | if args.gan: 57 | runners = __import__("runners." + args.model + "_gan_runner") 58 | module = getattr(runners, args.model + "_gan_runner") 59 | model = getattr(module, args.model + "GANRunner")(cfg, logger) 60 | 61 | else: 62 | runners = __import__("runners." + args.model + "_runner") 63 | module = getattr(runners, args.model + "_runner") 64 | model = getattr(module, args.model + "Runner")(cfg, logger) 65 | 66 | model.runner() 67 | 68 | 69 | if __name__ == "__main__": 70 | main() 71 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/microsoft/SpareNet/2da5b09363fbe42499c99709752af94848ecfb93/utils/__init__.py -------------------------------------------------------------------------------- /utils/misc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import open3d 7 | import logging 8 | from tensorboardX import SummaryWriter 9 | import cuda.emd.emd_module as emd 10 | from cuda.chamfer_distance import ChamferDistanceMean 11 | 12 | logger = logging.getLogger() 13 | 14 | ##### type conversion ##### 15 | def var_or_cuda(x): 16 | if torch.cuda.is_available(): 17 | x = x.cuda(non_blocking=True) 18 | 19 | return x 20 | 21 | 22 | def gpu_init(cfg): 23 | """ 24 | input: 25 | cfg: EasyDict 26 | 27 | output: 28 | gup_ids: list 29 | """ 30 | # Set up folders for checkpoints 31 | if not os.path.exists(cfg.DIR.checkpoints): 32 | os.makedirs(cfg.DIR.checkpoints) 33 | # GPU setup 34 | torch.backends.cudnn.benchmark = True 35 | gup_ids = [int(x) for x in cfg.CONST.device.split(",")] 36 | return list(range(len(gup_ids))) 37 | 38 | 39 | def writer_init(cfg): 40 | """ 41 | input: 42 | cfg: EasyDict 43 | 44 | outputs: 45 | train_writer: SummaryWriter 46 | val_writer: SummaryWriter 47 | """ 48 | # Create tensorboard writers 49 | train_writer = SummaryWriter(os.path.join(cfg.DIR.logs, "train")) 50 | val_writer = SummaryWriter(os.path.join(cfg.DIR.logs, "test")) 51 | return train_writer, val_writer 52 | 53 | 54 | def model_load(cfg, net_G): 55 | """ 56 | load model 57 | 58 | inputs: 59 | cfg: EasyDict 60 | net_G: torch.nn.module 61 | 62 | outputs: 63 | init_epoch: int 64 | best_metrics: dic 65 | """ 66 | 67 | init_epoch = 0 68 | best_metrics = None 69 | # Load pretrained model if exists 70 | if cfg.CONST.weights: 71 | logger.info("Recovering from %s ..." % (cfg.CONST.weights)) 72 | checkpoint = torch.load(cfg.CONST.weights) 73 | best_metrics = Metrics(cfg.TEST.metric_name, checkpoint["best_metrics"]) 74 | init_epoch = checkpoint["epoch_index"] 75 | net_G.load_state_dict(checkpoint["net_G"]) # change into net_G!! 76 | logger.info("Recover complete. Current epoch = #%d; best metrics = %s." % (init_epoch, best_metrics)) 77 | return init_epoch, best_metrics 78 | 79 | 80 | def checkpoint_save(cfg, epoch_idx, metrics, best_metrics, net_G): 81 | """ 82 | save the model 83 | 84 | inputs: 85 | cfg: EasyDict 86 | epoch_idx: int 87 | metrics: dic 88 | best_metrics: dic 89 | net_G: torch.nn.module 90 | 91 | outputs: 92 | best_metrics: dic 93 | """ 94 | # save tbe best model 95 | if epoch_idx % cfg.TRAIN.save_freq == 0 or metrics.better_than(best_metrics): 96 | file_name = "ckpt-best.pth" if metrics.better_than(best_metrics) else "ckpt-epoch-%03d.pth" % epoch_idx 97 | output_path = os.path.join(cfg.DIR.checkpoints, file_name) 98 | # save the epoch and metrics and net_G 99 | state = { 100 | "epoch_index": epoch_idx, 101 | "best_metrics": metrics.state_dict(), 102 | "net_G": net_G.state_dict(), 103 | } 104 | torch.save(state, output_path) 105 | 106 | logger.info("Saved checkpoint to %s ..." % output_path) 107 | if metrics.better_than(best_metrics): 108 | best_metrics = metrics 109 | return best_metrics 110 | 111 | 112 | def set_logger(filename): 113 | """s 114 | set logger 115 | """ 116 | logger = logging.getLogger() 117 | logger.setLevel(logging.DEBUG) 118 | formatter = logging.Formatter("%(levelname)s: - %(message)s") 119 | 120 | fh = logging.FileHandler(filename) 121 | fh.setLevel(logging.DEBUG) 122 | fh.setFormatter(formatter) 123 | 124 | ch = logging.StreamHandler() 125 | ch.setLevel(logging.DEBUG) 126 | ch.setFormatter(formatter) 127 | 128 | logger.addHandler(ch) 129 | logger.addHandler(fh) 130 | return logger 131 | 132 | 133 | class Metrics(object): 134 | ITEMS = [ 135 | { 136 | "name": "F-Score", 137 | "enabled": True, 138 | "eval_func": "cls._get_f_score", 139 | "is_greater_better": True, 140 | "init_value": 0, 141 | }, 142 | { 143 | "name": "ChamferDistance", 144 | "enabled": True, 145 | "eval_func": "cls._get_chamfer_distance", 146 | "eval_object": ChamferDistanceMean(), 147 | "is_greater_better": False, 148 | "init_value": 32767, 149 | }, 150 | { 151 | "name": "EMD", 152 | "enabled": True, 153 | "eval_func": "cls._get_emd", 154 | "eval_object": emd.emdModule(), 155 | "is_greater_better": False, 156 | "init_value": 32767, 157 | }, 158 | ] 159 | 160 | @classmethod 161 | def get(cls, pred, gt): 162 | _items = cls.items() 163 | _values = [0] * len(_items) 164 | for i, item in enumerate(_items): 165 | eval_func = eval(item["eval_func"]) 166 | _values[i] = eval_func(pred, gt) 167 | 168 | return _values 169 | 170 | @classmethod 171 | def items(cls): 172 | return [i for i in cls.ITEMS if i["enabled"]] 173 | 174 | @classmethod 175 | def names(cls): 176 | _items = cls.items() 177 | return [i["name"] for i in _items] 178 | 179 | @classmethod 180 | def _get_f_score(cls, pred, gt, th=0.01): 181 | """References: https://github.com/lmb-freiburg/what3d/blob/master/util.py""" 182 | pred = cls._get_open3d_ptcloud(pred) 183 | gt = cls._get_open3d_ptcloud(gt) 184 | 185 | dist1 = pred.compute_point_cloud_distance(gt) 186 | dist2 = gt.compute_point_cloud_distance(pred) 187 | 188 | recall = float(sum(d < th for d in dist2)) / float(len(dist2)) 189 | precision = float(sum(d < th for d in dist1)) / float(len(dist1)) 190 | return 2 * recall * precision / (recall + precision) if recall + precision else 0 191 | 192 | @classmethod 193 | def _get_open3d_ptcloud(cls, tensor): 194 | tensor = tensor.squeeze().cpu().numpy() 195 | ptcloud = open3d.geometry.PointCloud() 196 | ptcloud.points = open3d.utility.Vector3dVector(tensor) 197 | 198 | return ptcloud 199 | 200 | @classmethod 201 | def _get_chamfer_distance(cls, pred, gt): 202 | chamfer_distance = cls.ITEMS[1]["eval_object"] 203 | return chamfer_distance(pred, gt).item() * 1000 204 | 205 | @classmethod 206 | def _get_emd(cls, pred, gt): 207 | EMD = cls.ITEMS[2]["eval_object"] 208 | dist, _ = EMD(pred, gt, eps=0.005, iters=50) # for val 209 | # dist, _ = EMD(pred, gt, 0.002, 10000) # final test ? 210 | emd = torch.sqrt(dist).mean(1).mean() 211 | return emd.item() * 100 212 | 213 | def __init__(self, metric_name, values): 214 | self._items = Metrics.items() 215 | self._values = [item["init_value"] for item in self._items] 216 | self.metric_name = metric_name 217 | 218 | if type(values).__name__ == "dict": 219 | metric_indexes = {} 220 | for idx, item in enumerate(self._items): 221 | item_name = item["name"] 222 | metric_indexes[item_name] = idx 223 | for k, v in values.items(): 224 | if k not in metric_indexes: 225 | logger.warn("Ignore Metric[Name=%s] due to disability." % k) 226 | continue 227 | self._values[metric_indexes[k]] = v 228 | elif type(values).__name__ == "list": 229 | self._values = values 230 | else: 231 | raise Exception("Unsupported value type: %s" % type(values)) 232 | 233 | def state_dict(self): 234 | _dict = {} 235 | for i in range(len(self._items)): 236 | item = self._items[i]["name"] 237 | value = self._values[i] 238 | _dict[item] = value 239 | 240 | return _dict 241 | 242 | def __repr__(self): 243 | return str(self.state_dict()) 244 | 245 | def better_than(self, other): 246 | if other is None: 247 | return True 248 | 249 | _index = -1 250 | for i, _item in enumerate(self._items): 251 | if _item["name"] == self.metric_name: 252 | _index = i 253 | break 254 | if _index == -1: 255 | raise Exception("Invalid metric name to compare.") 256 | 257 | _metric = self._items[i] 258 | _value = self._values[_index] 259 | other_value = other._values[_index] 260 | return _value > other_value if _metric["is_greater_better"] else _value < other_value 261 | -------------------------------------------------------------------------------- /utils/model_init.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | 4 | import os 5 | import torch 6 | import logging 7 | import numpy as np 8 | import configs.model_names as name 9 | from utils.p2i_utils import ComputeDepthMaps 10 | from models.sparenet_discriminator import ProjectionD, PatchDiscriminator 11 | 12 | logger = logging.getLogger() 13 | 14 | 15 | def generator_init(cfg): 16 | """ 17 | input: 18 | cfg: EasyDict 19 | 20 | outputs: 21 | net_G: torch.nn.module 22 | optimizerG: torch.optim.Adam 23 | lr_schedulerG: torch.optim.lr_scheduler.MultiStepLR 24 | """ 25 | # Create the networks 26 | net_G = define_G(cfg) 27 | net_G.apply(init_weights) 28 | logger.debug("Parameters in net_G: %d." % count_parameters(net_G)) 29 | 30 | # Optimizer and learning scheduler 31 | optimizerG = torch.optim.Adam( 32 | filter(lambda p: p.requires_grad, net_G.parameters()), 33 | lr=cfg.TRAIN.learning_rate, 34 | weight_decay=cfg.TRAIN.weight_decay, 35 | betas=cfg.TRAIN.betas, 36 | ) 37 | lr_schedulerG = torch.optim.lr_scheduler.MultiStepLR(optimizerG, milestones=cfg.TRAIN.lr_milestones, gamma=cfg.TRAIN.gamma) 38 | return net_G, optimizerG, lr_schedulerG 39 | 40 | 41 | def renderer_init(cfg): 42 | """ 43 | input: 44 | cfg: EasyDict 45 | 46 | output: 47 | renderer: torch.nn.module 48 | """ 49 | # Create the differentiable renderer 50 | renderer = ComputeDepthMaps( 51 | projection=cfg.RENDER.projection, 52 | eyepos_scale=cfg.RENDER.eyepos, 53 | image_size=cfg.RENDER.img_size, 54 | ).float() 55 | return renderer 56 | 57 | 58 | def discriminator_init(cfg): 59 | """ 60 | input: 61 | cfg: EasyDict 62 | 63 | outputs: 64 | net_D: torch.nn.module 65 | optimizerD: torch.optim.Adam 66 | lr_schedulerD: torch.optim.lr_scheduler.MultiStepLR 67 | """ 68 | # Create discriminator: projection discriminator or patchgan 69 | if cfg.GAN.use_cgan: 70 | net_D = ProjectionD( 71 | num_classes=cfg.DATASET.num_classes, 72 | img_shape=(cfg.RENDER.n_views * 2, cfg.RENDER.img_size, cfg.RENDER.img_size), 73 | ) 74 | else: 75 | net_D = PatchDiscriminator( 76 | img_shape=(cfg.RENDER.n_views * 2, cfg.RENDER.img_size, cfg.RENDER.img_size), 77 | ) 78 | 79 | net_D.apply(init_weights_D) 80 | logger.debug("Parameters in net_D: %d." % count_parameters(net_D)) 81 | 82 | # Optimizer and learning scheduler 83 | optimizerD = torch.optim.Adam( 84 | filter(lambda p: p.requires_grad, net_D.parameters()), 85 | lr=cfg.TRAIN.learning_rate, 86 | weight_decay=cfg.TRAIN.weight_decay, 87 | betas=cfg.TRAIN.betas, 88 | ) 89 | lr_schedulerD = torch.optim.lr_scheduler.MultiStepLR(optimizerD, milestones=cfg.TRAIN.lr_milestones, gamma=cfg.TRAIN.gamma) 90 | return net_D, optimizerD, lr_schedulerD 91 | 92 | 93 | 94 | def define_G(cfg): 95 | if cfg.NETWORK.model_type == name.MODEL_SPARENET: 96 | from models.sparenet_generator import SpareNetGenerator 97 | 98 | # from models.mvnet import MVNet 99 | network = SpareNetGenerator( 100 | num_points=cfg.DATASET.n_outpoints, 101 | bottleneck_size=4096, 102 | n_primitives=cfg.NETWORK.n_primitives, 103 | use_SElayer=cfg.NETWORK.use_selayer, 104 | use_AdaIn=cfg.NETWORK.use_adain, 105 | encode=cfg.NETWORK.encode, 106 | hide_size=4096, 107 | ) 108 | 109 | elif cfg.NETWORK.model_type == name.MODEL_ATLASNET: 110 | from models.atlasnet_generator import AtlasNet 111 | 112 | network = AtlasNet( 113 | num_points=cfg.DATASET.n_outpoints, 114 | bottleneck_size=1024, 115 | n_primitives=cfg.NETWORK.n_primitives, 116 | ) 117 | 118 | elif cfg.NETWORK.model_type == name.MODEL_MSN: 119 | from models.msn_generator import MSN 120 | 121 | network = MSN( 122 | num_points=cfg.DATASET.n_outpoints, 123 | bottleneck_size=1024, 124 | n_primitives=cfg.NETWORK.n_primitives, 125 | ) 126 | 127 | elif cfg.NETWORK.model_type == name.MODEL_GRNET: 128 | from models.grnet_generator import GRNet 129 | 130 | network = GRNet(cfg) 131 | 132 | else: 133 | raise Exception("Unknown model type") 134 | return network 135 | 136 | 137 | def init_weights(m): 138 | if type(m) in [ 139 | torch.nn.Conv2d, 140 | torch.nn.ConvTranspose2d, 141 | torch.nn.Conv3d, 142 | torch.nn.ConvTranspose3d, 143 | ] and hasattr(m, "weight"): 144 | torch.nn.init.kaiming_normal_(m.weight) 145 | if m.bias is not None: 146 | torch.nn.init.constant_(m.bias, 0) 147 | if type(m) == torch.nn.Conv1d: 148 | torch.nn.init.normal_(m.weight.data, 0.0, 0.02) 149 | elif type(m) in [torch.nn.BatchNorm2d, torch.nn.BatchNorm3d]: 150 | torch.nn.init.constant_(m.weight, 1) 151 | torch.nn.init.constant_(m.bias, 0) 152 | elif type(m) == torch.nn.BatchNorm1d: 153 | torch.nn.init.normal_(m.weight.data, 1.0, 0.02) 154 | torch.nn.init.constant_(m.bias.data, 0.0) 155 | elif type(m) == torch.nn.Linear: 156 | if hasattr(m, "weight"): 157 | torch.nn.init.normal_(m.weight, 0, 0.01) 158 | if m.bias is not None: 159 | torch.nn.init.constant_(m.bias, 0) 160 | 161 | 162 | def init_weights_D(m): 163 | classname = m.__class__.__name__ 164 | if ( 165 | classname.find("Conv2d") != -1 166 | and hasattr(m, "weight") 167 | or classname.find("Conv2d") == -1 168 | and classname.find("Conv1d") != -1 169 | and hasattr(m, "weight") 170 | ): 171 | torch.nn.init.normal_(m.weight.data, 0.0, 0.02) 172 | if ( 173 | classname.find("BatchNorm2d") != -1 174 | or classname.find("BatchNorm2d") == -1 175 | and classname.find("BatchNorm1d") != -1 176 | ): 177 | torch.nn.init.normal_(m.weight.data, 1.0, 0.02) 178 | torch.nn.init.constant_(m.bias.data, 0.0) 179 | 180 | 181 | def count_parameters(network): 182 | return sum(p.numel() for p in network.parameters()) 183 | 184 | --------------------------------------------------------------------------------