├── .gitignore ├── README.md ├── _ext ├── __init__.py └── my_lib │ └── __init__.py ├── build.py ├── eval_img2pc.py ├── functions ├── __init__.py └── nnd.py ├── models ├── AutoEncoder.py ├── GAN.py ├── ImageToShape.py ├── MRTDecoder.py ├── Model.py ├── VoxUNet.py └── __init__.py ├── modules ├── __init__.py └── nnd.py ├── mr_train_4k.py ├── nndistance ├── .gitignore ├── README.md ├── _ext │ ├── __init__.py │ └── my_lib │ │ └── __init__.py ├── build.py ├── functions │ ├── __init__.py │ └── nnd.py ├── modules │ ├── __init__.py │ └── nnd.py ├── src │ ├── make.sh │ ├── my_lib.c │ ├── my_lib.h │ ├── my_lib_cuda.c │ ├── my_lib_cuda.h │ ├── nnd_cuda.cu │ └── nnd_cuda.h └── test.py ├── run_img2pc.py ├── sampler ├── README ├── cotSpectral.java └── vecmath.jar ├── src ├── emd_cuda.cu ├── emd_cuda.cu.o ├── emd_cuda.h ├── make.sh ├── my_lib.c ├── my_lib.h ├── my_lib_cuda.c ├── my_lib_cuda.h ├── nnd_cuda.cu ├── nnd_cuda.cu.o └── nnd_cuda.h ├── test.py ├── testchamfer.py ├── tools ├── .ImageToPCDataset.py.swp ├── .PointCloudDataset.py.swp ├── DataVis.py ├── ImageToPCDataset.py ├── Ops.py ├── PointCloudDataset.py ├── Trainer.py └── __init__.py └── train_img2pc.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multiresolution Tree Networks for 3D Point Cloud Processing (ECCV 2018) 2 | 3 | 4 | This repository contains the source code for the ECCV 2018 paper Multiresolution Tree Networks for 3D Point Clout Processing. 5 | 6 | [Project page](http://mgadelha.me/mrt/) 7 | 8 | ![MRTNet reconstructions](http://mgadelha.me/mrt/fig/realrec2.png) 9 | 10 | ## Dependencies 11 | 12 | * numpy 13 | * pytorch 14 | * tensorboardX 15 | * fxia22/pointGAN (optional - if you want faster Chamfer Distance) Thanks to Fei Xia'a for making the code publicly available. We have a version in this repo already, but might not be up-to-date. 16 | 17 | 18 | ## Train 19 | 20 | First, you need to change the dataset path in the file `train_img2pc.py`. We are using the rendered images from https://github.com/chrischoy/3D-R2N2. You will also need a path for the point clouds in .npy format. Finally, you can train a model by using the following command: 21 | ``` 22 | python train_img2pc.py --name experiment_name 23 | ``` 24 | 25 | If you want to run the model, change the folder name indicated in `run_img2pc.py` and use the following command: 26 | ``` 27 | python run_img2pc.py --n experiment_name 28 | ``` 29 | Notice that`experiment_name` should match in both cases. Similarly, we also have evaluation code to reproduce the paper's numbers. 30 | 31 | 32 | ## Point Sampling 33 | 34 | This repository also contains code for sampling the point clouds in the sampler folder. It is a single .java file and it contains a README with specific instructions. 35 | This code automatically sorts the points according to a kd-tree structure. 36 | 37 | 38 | ## Citation 39 | 40 | If you use any part of this code or data, consider citing this work: 41 | ``` 42 | @inProceedings{mrt18, 43 | title={Multiresolution Tree Networks for 3D Point Cloud Processing}, 44 | author = {Matheus Gadelha and Rui Wang and Subhransu Maji}, 45 | booktitle={ECCV}, 46 | year={2018} 47 | } 48 | ``` 49 | -------------------------------------------------------------------------------- /_ext/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/_ext/__init__.py -------------------------------------------------------------------------------- /_ext/my_lib/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._my_lib import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.utils.ffi import create_extension 4 | 5 | this_file = os.path.dirname(__file__) 6 | 7 | sources = ['src/my_lib.c'] 8 | headers = ['src/my_lib.h'] 9 | defines = [] 10 | with_cuda = False 11 | 12 | if torch.cuda.is_available(): 13 | print('Including CUDA code.') 14 | sources += ['src/my_lib_cuda.c'] 15 | headers += ['src/my_lib_cuda.h'] 16 | defines += [('WITH_CUDA', None)] 17 | with_cuda = True 18 | 19 | this_file = os.path.dirname(os.path.realpath(__file__)) 20 | print(this_file) 21 | extra_objects = ['src/nnd_cuda.cu.o'] 22 | extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] 23 | 24 | ffi = create_extension( 25 | '_ext.my_lib', 26 | headers=headers, 27 | sources=sources, 28 | define_macros=defines, 29 | relative_to=__file__, 30 | with_cuda=with_cuda, 31 | extra_objects=extra_objects 32 | ) 33 | 34 | if __name__ == '__main__': 35 | ffi.build() 36 | -------------------------------------------------------------------------------- /eval_img2pc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.optim as optim 4 | import torch.nn as nn 5 | 6 | import argparse 7 | import os 8 | 9 | from tools.Trainer import ImageToPCTrainer 10 | from tools.PointCloudDataset import ImageToPointCloudDataset 11 | from models.AutoEncoder import PointCloudVAE 12 | from models.AutoEncoder import ChamferLoss 13 | from models.AutoEncoder import ChamferWithNormalLoss 14 | from models.AutoEncoder import L2WithNormalLoss 15 | from models.ImageToShape import MultiResImageToShape 16 | from models.ImageToShape import SingleResImageToShape 17 | from models.ImageToShape import FCImageToShape 18 | 19 | parser = argparse.ArgumentParser(description='MultiResolution image to shape model.') 20 | parser.add_argument("-n", "--name", type=str, help="Name of the experiment.", default="MRI2PC") 21 | parser.add_argument("-a", "--arch", type=str, help="Encoder architecture.", default="vgg") 22 | parser.add_argument("-pt", "--pretrained", type=str, help="Use pretrained net", default="True") 23 | parser.add_argument("-c", "--category", type=str, help="Category code (all is possible)", default="all") 24 | parser.add_argument("--train", dest='train', action='store_true') 25 | parser.set_defaults(train=False) 26 | 27 | image_datapath = "/media/mgadelha/hd2/ShapenetRenderings" 28 | pc_datapath = "/media/mgadelha/hd2/shapenet_4k" 29 | 30 | if __name__ == '__main__': 31 | args = parser.parse_args() 32 | 33 | ptrain = None 34 | if args.pretrained == "False": 35 | ptrain = False 36 | elif args.pretrained == "True": 37 | ptrain = True 38 | 39 | full_name = "{}_{}_{}_{}".format(args.name, args.category, args.arch, ptrain) 40 | #full_name = args.name 41 | print full_name 42 | 43 | mri2pc = MultiResImageToShape(size=4096, dim=3, batch_size=1, 44 | name=full_name, pretrained=ptrain, arch=args.arch) 45 | mri2pc.load('checkpoint') 46 | optimizer = optim.Adam(mri2pc.parameters(), lr=0.001) 47 | 48 | train_dataset = ImageToPointCloudDataset(image_datapath, pc_datapath, 49 | category=args.category, train_mode=True) 50 | test_dataset = ImageToPointCloudDataset(image_datapath, pc_datapath, 51 | category=args.category, train_mode=False) 52 | train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, 53 | shuffle=True, num_workers=2) 54 | test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, 55 | shuffle=True, num_workers=2) 56 | 57 | log_dir = os.path.join("log", full_name) 58 | if not os.path.exists(log_dir): 59 | os.makedirs(log_dir) 60 | 61 | trainer = ImageToPCTrainer(mri2pc, train_loader, test_loader, 62 | optimizer, ChamferLoss(), log_dir=log_dir) 63 | trainer.evaluate() 64 | 65 | -------------------------------------------------------------------------------- /functions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/functions/__init__.py -------------------------------------------------------------------------------- /functions/nnd.py: -------------------------------------------------------------------------------- 1 | # functions/add.py 2 | import torch 3 | from torch.autograd import Function 4 | from _ext import my_lib 5 | 6 | 7 | class NNDFunction(Function): 8 | def forward(self, xyz1, xyz2): 9 | batchsize, n, _ = xyz1.size() 10 | _, m, _ = xyz2.size() 11 | self.xyz1 = xyz1 12 | self.xyz2 = xyz2 13 | dist1 = torch.zeros(batchsize, n) 14 | dist2 = torch.zeros(batchsize, m) 15 | 16 | self.idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) 17 | self.idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) 18 | 19 | if not xyz1.is_cuda: 20 | my_lib.nnd_forward(xyz1, xyz2, dist1, dist2, self.idx1, self.idx2) 21 | else: 22 | dist1 = dist1.cuda() 23 | dist2 = dist2.cuda() 24 | self.idx1 = self.idx1.cuda() 25 | self.idx2 = self.idx2.cuda() 26 | my_lib.nnd_forward_cuda(xyz1, xyz2, dist1, dist2, self.idx1, self.idx2) 27 | 28 | self.dist1 = dist1 29 | self.dist2 = dist2 30 | 31 | #print(batchsize, n, m) 32 | 33 | return dist1, dist2 34 | 35 | def backward(self, graddist1, graddist2): 36 | #print(self.idx1, self.idx2) 37 | 38 | 39 | graddist1 = graddist1.contiguous() 40 | graddist2 = graddist2.contiguous() 41 | 42 | gradxyz1 = torch.zeros(self.xyz1.size()) 43 | gradxyz2 = torch.zeros(self.xyz2.size()) 44 | 45 | if not graddist1.is_cuda: 46 | my_lib.nnd_backward(self.xyz1, self.xyz2, gradxyz1, gradxyz2, graddist1, graddist2, self.idx1, self.idx2) 47 | else: 48 | gradxyz1 = gradxyz1.cuda() 49 | gradxyz2 = gradxyz2.cuda() 50 | my_lib.nnd_backward_cuda(self.xyz1, self.xyz2, gradxyz1, gradxyz2, graddist1, graddist2, self.idx1, self.idx2) 51 | 52 | return gradxyz1, gradxyz2 -------------------------------------------------------------------------------- /models/AutoEncoder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from torch.autograd import Variable 7 | 8 | from Model import Model 9 | import tools.DataVis as DataVis 10 | from tools.PointCloudDataset import save_objs 11 | from tools import Ops 12 | from modules.nnd import NNDModule 13 | 14 | 15 | class PointCloudEncoder(nn.Module): 16 | 17 | def __init__(self, size, dim, batch_size=64, enc_size=100, kernel_size=16, 18 | init_channels=16): 19 | super(PointCloudEncoder, self).__init__() 20 | self.size = size 21 | self.dim = dim 22 | self.batch_size = batch_size 23 | self.kernel_size = kernel_size 24 | self.enc_size = enc_size 25 | self.init_channels = init_channels 26 | 27 | conv_enc = nn.Sequential() 28 | 29 | current_size = self.size 30 | in_channels = self.dim 31 | out_channels = self.init_channels 32 | layer_num = 1 33 | padding = (self.kernel_size - 1)/2 34 | while current_size > 16: 35 | conv_enc.add_module('conv{}'.format(layer_num), 36 | nn.Conv1d(in_channels, out_channels, self.kernel_size, 37 | stride=2, 38 | padding=padding)) 39 | conv_enc.add_module('bn{}'.format(layer_num), 40 | nn.BatchNorm1d(out_channels)) 41 | conv_enc.add_module('lrelu{}'.format(layer_num), 42 | nn.LeakyReLU(0.2, inplace=True)) 43 | 44 | current_size /= 2 45 | in_channels = out_channels 46 | out_channels *= 2 47 | layer_num += 1 48 | 49 | self.conv_enc = conv_enc 50 | 51 | self.fc = nn.Linear(16*in_channels, self.enc_size) 52 | 53 | def forward(self, x): 54 | t = self.conv_enc(x).view(self.batch_size, -1) 55 | out = self.fc(t) 56 | return out 57 | 58 | 59 | class PointCloudDecoder(nn.Module): 60 | 61 | def __init__(self, size, dim, batch_size=64, enc_size=100, kernel_size=16, 62 | init_channels=1024): 63 | super(PointCloudDecoder, self).__init__() 64 | self.size = size 65 | self.dim = dim 66 | self.batch_size = batch_size 67 | self.kernel_size = kernel_size 68 | self.enc_size = enc_size 69 | self.init_channels = init_channels 70 | 71 | self.fc = nn.Linear(self.enc_size, 16*self.init_channels) 72 | 73 | conv_dec = nn.Sequential() 74 | 75 | current_size = 16*2 76 | in_channels = self.init_channels 77 | out_channels = in_channels/2 78 | layer_num = 1 79 | padding = (self.kernel_size - 1)/2 80 | while current_size < self.size: 81 | conv_dec.add_module('conv{}'.format(layer_num), 82 | nn.ConvTranspose1d(in_channels, out_channels, self.kernel_size, 83 | stride=2, 84 | padding=padding)) 85 | conv_dec.add_module('bn{}'.format(layer_num), 86 | nn.BatchNorm1d(out_channels)) 87 | conv_dec.add_module('lrelu{}'.format(layer_num), 88 | nn.LeakyReLU(0.2, inplace=True)) 89 | 90 | current_size *= 2 91 | in_channels = out_channels 92 | out_channels /= 2 93 | layer_num += 1 94 | 95 | conv_dec.add_module('conv{}'.format(layer_num), 96 | nn.ConvTranspose1d(in_channels, self.dim, self.kernel_size, 97 | stride=2, 98 | padding=padding)) 99 | conv_dec.add_module('lrelu{}'.format(layer_num), 100 | nn.Tanh()) 101 | 102 | self.conv_dec = conv_dec 103 | 104 | 105 | def forward(self, x): 106 | t = self.fc(x).view(self.batch_size, self.init_channels, 16) 107 | out = self.conv_dec(t) 108 | return out 109 | 110 | 111 | class MultiResBlock1d(nn.Module): 112 | 113 | def __init__(self, name, in_channels, out_channels, blocktype, activation): 114 | super(MultiResBlock1d, self).__init__() 115 | 116 | self.upsample = Ops.NNUpsample1d() 117 | self.pool = nn.MaxPool1d(kernel_size=4, stride=4) 118 | self.in_channels = in_channels 119 | self.out_channels = out_channels 120 | self.name = name 121 | 122 | self.conv0 = nn.Sequential() 123 | self.conv0.add_module('{}_conv0'.format(self.name), 124 | blocktype(self.in_channels*2, 125 | self.out_channels, 126 | kernel_size=2, 127 | stride=2, 128 | padding=0)) 129 | self.conv0.add_module('{}_bn0'.format(self.name), 130 | nn.BatchNorm1d(self.out_channels)) 131 | self.conv0.add_module('{}_activation0'.format(self.name), 132 | activation) 133 | 134 | self.conv1 = nn.Sequential() 135 | self.conv1.add_module('{}_conv1'.format(self.name), 136 | blocktype(self.in_channels*3, 137 | self.out_channels, 138 | kernel_size=2, 139 | stride=2, 140 | padding=0)) 141 | self.conv1.add_module('{}_bn1'.format(self.name), 142 | nn.BatchNorm1d(self.out_channels)) 143 | self.conv1.add_module('{}_activation1'.format(self.name), 144 | activation) 145 | 146 | self.conv2 = nn.Sequential() 147 | self.conv2.add_module('{}_conv2'.format(self.name), 148 | blocktype(self.in_channels*2, 149 | self.out_channels, 150 | kernel_size=2, 151 | stride=2, 152 | padding=0)) 153 | self.conv2.add_module('{}_bn2'.format(self.name), 154 | nn.BatchNorm1d(self.out_channels)) 155 | self.conv2.add_module('{}_activation2'.format(self.name), 156 | activation) 157 | 158 | def forward(self, x): 159 | x0 = x[0] 160 | x1 = x[1] 161 | x2 = x[2] 162 | 163 | in0 = torch.cat((x0, self.upsample(x1)), 1) 164 | in1 = torch.cat((self.pool(x0), x1, self.upsample(x2)), 1) 165 | in2 = torch.cat((self.pool(x1), x2), 1) 166 | 167 | out0 = self.conv0(in0) 168 | out1 = self.conv1(in1) 169 | out2 = self.conv2(in2) 170 | 171 | return [out0, out1, out2] 172 | 173 | 174 | class MultiResConv1d(MultiResBlock1d): 175 | 176 | def __init__(self, name, in_channels, out_channels, activation=nn.ReLU(inplace=True)): 177 | super(MultiResConv1d, self).__init__( 178 | name, in_channels, out_channels, nn.Conv1d, activation=activation) 179 | 180 | 181 | class MultiResConvTranspose1d(MultiResBlock1d): 182 | 183 | def __init__(self, name, in_channels, out_channels, activation=nn.ReLU(inplace=True)): 184 | super(MultiResConvTranspose1d, self).__init__( 185 | name, in_channels, out_channels, nn.ConvTranspose1d, activation=activation) 186 | 187 | 188 | class ChamferLoss(nn.Module): 189 | 190 | def __init__(self, n_samples=1024, cuda_opt=True): 191 | super(ChamferLoss, self).__init__() 192 | self.n_samples = n_samples 193 | self.dist = NNDModule() 194 | self.cuda_opt = cuda_opt 195 | 196 | 197 | def chamfer(self, a, b): 198 | pcsize = a.size()[1] 199 | 200 | a = torch.t(a) 201 | b = torch.t(b) 202 | mma = torch.stack([a]*pcsize) 203 | mmb = torch.stack([b]*pcsize).transpose(0,1) 204 | d = torch.sum((mma-mmb)**2,2).squeeze() 205 | 206 | return torch.min(d, 1)[0].sum() + torch.min(d, 0)[0].sum() 207 | 208 | 209 | def chamfer_batch(self, a, b): 210 | pcsize = a.size()[-1] 211 | 212 | if pcsize != self.n_samples: 213 | indices = np.arange(pcsize).astype(int) 214 | np.random.shuffle(indices) 215 | indices = torch.from_numpy(indices[:self.n_samples]).cuda() 216 | a = a[:, :, indices] 217 | b = b[:, :, indices] 218 | 219 | a = torch.transpose(a, 1, 2).contiguous() 220 | b = torch.transpose(b, 1, 2).contiguous() 221 | 222 | if self.cuda_opt: 223 | d1, d2 = self.dist(a, b) 224 | out = torch.sum(d1) + torch.sum(d2) 225 | return out 226 | else: 227 | d = Ops.batch_pairwise_dist(a, b) 228 | return torch.min(d, dim=2)[0].sum() + torch.min(d, dim=1)[0].sum() 229 | 230 | #mma = torch.stack([a]*self.n_samples, dim=1) 231 | #mmb = torch.stack([b]*self.n_samples, dim=1).transpose(1,2) 232 | #d = torch.sum((mma-mmb)**2,3).squeeze() 233 | #d = pd 234 | 235 | 236 | 237 | def forward(self, a, b): 238 | batch_size = a.size()[0] 239 | assert(batch_size == b.size()[0]) 240 | loss = self.chamfer_batch(a, b) 241 | return loss/(float(batch_size) * 1) 242 | 243 | 244 | class MultiResChamferLoss(nn.Module): 245 | 246 | def __init__(self, n_samples=1024): 247 | super(MultiResChamferLoss, self).__init__() 248 | self.n_samples = n_samples 249 | self.pool = nn.MaxPool1d(kernel_size=8, stride=8) 250 | 251 | def chamfer_batch(self, a, b): 252 | pcsize = a.size()[-1] 253 | 254 | if pcsize > self.n_samples: 255 | pcsize = self.n_samples 256 | 257 | indices = np.arange(pcsize) 258 | np.random.shuffle(indices) 259 | indices = indices[:self.n_samples] 260 | 261 | a = a[:, :, indices] 262 | b = b[:, :, indices] 263 | 264 | a = torch.transpose(a, 1, 2) 265 | b = torch.transpose(b, 1, 2) 266 | #d = Ops.batch_pairwise_dist(a, b) 267 | mma = torch.stack([a]*self.n_samples, dim=1) 268 | mmb = torch.stack([b]*self.n_samples, dim=1).transpose(1,2) 269 | d = torch.sum((mma-mmb)**2,3).squeeze() 270 | #d = pd 271 | 272 | return (torch.min(d, dim=2)[0].sum() + torch.min(d, dim=1)[0].sum())/float(pcsize) 273 | 274 | 275 | def forward(self, a, b): 276 | batch_size = a[0].size()[0] 277 | 278 | b_samples = [] 279 | b_samples.append(b) 280 | b_samples.append(self.pool(b_samples[-1])) 281 | b_samples.append(self.pool(b_samples[-1])) 282 | 283 | loss = 0.0 284 | for i in [0]: 285 | loss += self.chamfer_batch(a[i], b_samples[i]) 286 | 287 | return 1e3*loss/float(batch_size) 288 | 289 | 290 | 291 | class ChamferWithNormalLoss(nn.Module): 292 | 293 | def __init__(self, normal_weight=0.001, n_samples=1024): 294 | super(ChamferWithNormalLoss, self).__init__() 295 | self.normal_weight = normal_weight 296 | self.nlogger = DataVis.LossLogger("normal component") 297 | self.n_samples = n_samples 298 | 299 | 300 | def forward(self, a, b): 301 | pcsize = a.size()[-1] 302 | 303 | if pcsize != self.n_samples: 304 | indices = np.arange(pcsize) 305 | np.random.shuffle(indices) 306 | indices = indices[:self.n_samples] 307 | a = a[:, :, indices] 308 | b = b[:, :, indices] 309 | 310 | a_points = torch.transpose(a, 1, 2)[:, :, 0:3] 311 | b_points = torch.transpose(b, 1, 2)[:, :, 0:3] 312 | pd = Ops.batch_pairwise_dist(a_points, b_points) 313 | #mma = torch.stack([a_points]*self.n_samples, dim=1) 314 | #mmb = torch.stack([b_points]*self.n_samples, dim=1).transpose(1,2) 315 | d = pd 316 | 317 | a_normals = torch.transpose(a, 1, 2)[:, :, 3:6] 318 | b_normals = torch.transpose(b, 1, 2)[:, :, 3:6] 319 | mma = torch.stack([a_normals]*self.n_samples, dim=1) 320 | mmb = torch.stack([b_normals]*self.n_samples, dim=1).transpose(1,2) 321 | d_norm = 1 - torch.sum(mma*mmb,3).squeeze() 322 | d += self.normal_weight * d_norm 323 | 324 | normal_min_mean = torch.min(d_norm, dim=2)[0].mean() 325 | self.nlogger.update(normal_min_mean) 326 | 327 | chamfer_sym = torch.min(d, dim=2)[0].sum() + torch.min(d, dim=1)[0].sum() 328 | chamfer_sym /= a.size()[0] 329 | 330 | return chamfer_sym 331 | 332 | 333 | class SampleChamfer(nn.Module): 334 | 335 | def __init__(self, normal_weight=0.001, n_samples=1024): 336 | super(SampleChamfer, self).__init__() 337 | self.normal_weight = normal_weight 338 | self.nlogger = DataVis.LossLogger("normal component") 339 | self.n_samples = n_samples 340 | 341 | 342 | def chamfer(self, a, b): 343 | 344 | a_indices = np.arange(a.size()[-1]) 345 | b_indices = np.arange(b.size()[-1]) 346 | np.random.shuffle(a_indices) 347 | np.random.shuffle(b_indices) 348 | a_indices = a_indices[:self.n_samples] 349 | b_indices = b_indices[:self.n_samples] 350 | a = a[:, a_indices] 351 | b = b[:, b_indices] 352 | 353 | a = torch.t(a) 354 | b = torch.t(b) 355 | mma = torch.stack([a[:, 0:3]]*self.n_samples) 356 | mmb = torch.stack([b[:, 0:3]]*self.n_samples).transpose(0,1) 357 | d = torch.sum((mma-mmb)**2,2).squeeze() 358 | 359 | #return torch.min(d, 1)[0].sum() + torch.min(d, 0)[0].sum() 360 | return torch.min(d, 0)[0].sum() 361 | 362 | def forward(self, a, b): 363 | # pcsize = a.size()[-1] 364 | # 365 | # if pcsize != self.n_samples: 366 | # indices = np.arange(pcsize) 367 | # np.random.shuffle(indices) 368 | # indices = indices[:self.n_samples] 369 | # a = a[:, :, indices] 370 | # b = b[:, :, indices] 371 | # 372 | # a_points = torch.transpose(a, 1, 2)[:, :, 0:3] 373 | # b_points = torch.transpose(b, 1, 2)[:, :, 0:3] 374 | # mma = torch.stack([a_points]*self.n_samples, dim=1) 375 | # mmb = torch.stack([b_points]*self.n_samples, dim=1).transpose(1,2) 376 | # d = torch.sum((mma-mmb)**2,3).squeeze() 377 | # 378 | # a_normals = torch.transpose(a, 1, 2)[:, :, 3:6] 379 | # b_normals = torch.transpose(b, 1, 2)[:, :, 3:6] 380 | # mma = torch.stack([a_normals]*self.n_samples, dim=1) 381 | # mmb = torch.stack([b_normals]*self.n_samples, dim=1).transpose(1,2) 382 | # d_norm = 1 - torch.sum(mma*mmb,3).squeeze() 383 | # d += self.normal_weight * d_norm 384 | # 385 | # normal_min_mean = torch.min(d_norm, dim=2)[0].mean() 386 | # self.nlogger.update(normal_min_mean) 387 | # 388 | # chamfer_sym = torch.min(d, dim=2)[0].sum() + torch.min(d, dim=1)[0].sum() 389 | # chamfer_sym /= a.size()[0] 390 | 391 | return self.chamfer(a, b) 392 | 393 | 394 | 395 | class SinkhornLoss(nn.Module): 396 | 397 | def __init__(self, n_iter=20, eps=1.0, batch_size=64, enc_size=512): 398 | super(SinkhornLoss, self).__init__() 399 | self.eps = eps 400 | self.n_iter = n_iter 401 | self.batch_size = batch_size 402 | self.normal_noise = torch.FloatTensor(batch_size, enc_size) 403 | 404 | 405 | def forward(self, x): 406 | bsize = x.size()[0] 407 | assert bsize == self.batch_size 408 | 409 | self.normal_noise.normal_() 410 | y = Variable(self.normal_noise.cuda()) 411 | 412 | #Computes MSE cost 413 | mmx = torch.stack([x]*bsize) 414 | mmy = torch.stack([x]*bsize).transpose(0, 1) 415 | c = torch.sum((mmx-mmy)**2,2).squeeze() 416 | 417 | k = (-c/self.eps).exp() 418 | b = Variable(torch.ones((bsize, 1))).cuda() 419 | a = Variable(torch.ones((bsize, 1))).cuda() 420 | 421 | #Sinkhorn iterations 422 | for l in range(self.n_iter): 423 | a = Variable(torch.ones((bsize, 1))).cuda() / (torch.mm(k, b)) 424 | b = Variable(torch.ones((bsize, 1))).cuda() / (torch.mm(k.t(), a)) 425 | 426 | loss = torch.mm(k * c, b) 427 | loss = torch.sum(loss*a) 428 | return loss 429 | 430 | 431 | 432 | class GeodesicChamferLoss(nn.Module): 433 | 434 | def __init__(self): 435 | super(GeodesicChamferLoss, self).__init__() 436 | 437 | 438 | def forward(self, a, b): 439 | pass 440 | 441 | 442 | class L2WithNormalLoss(nn.Module): 443 | 444 | def __init__(self): 445 | super(L2WithNormalLoss, self).__init__() 446 | self.nlogger = DataVis.LossLogger("normal w/ L2") 447 | self.L1 = nn.L1Loss() 448 | 449 | def forward(self, a, b): 450 | position_loss = self.L1(a[:, 0:3, :], b[:, 0:3, :]) 451 | normal_loss = torch.mean(1 - Ops.cosine_similarity(a[:, 3:6, :], b[:, 3:6, :])) 452 | self.nlogger.update(normal_loss) 453 | 454 | return normal_loss 455 | 456 | 457 | class PointCloudAutoEncoder(Model): 458 | 459 | def __init__(self, size, dim, batch_size=64, enc_size=100, kernel_size=16, 460 | noise=0, 461 | name="PCAutoEncoder"): 462 | super(PointCloudAutoEncoder, self).__init__(name) 463 | 464 | self.size = size 465 | self.dim = dim 466 | self.batch_size = batch_size 467 | self.kernel_size = kernel_size 468 | self.enc_size = enc_size 469 | self.noise_factor = noise 470 | self.enc_noise = torch.FloatTensor(self.batch_size, self.enc_size) 471 | 472 | self.encoder = PointCloudEncoder(self.size, self.dim, 473 | batch_size = self.batch_size, 474 | enc_size = self.enc_size, 475 | kernel_size = self.kernel_size) 476 | 477 | self.decoder = PointCloudDecoder(self.size, self.dim, 478 | batch_size = self.batch_size, 479 | enc_size = self.enc_size, 480 | kernel_size = self.kernel_size) 481 | 482 | # if self.dim == 6: 483 | # self.normal_decoder = PointCloudDecoder(self.size, 3, 484 | # batch_size = self.batch_size, 485 | # enc_size = self.enc_size, 486 | # kernel_size = self.kernel_size) 487 | 488 | 489 | def forward(self, x): 490 | encoding = self.encoder(x) 491 | self.enc_noise.normal_() 492 | 493 | added_noise = Variable(self.noise_factor*self.enc_noise.cuda()) 494 | 495 | encoding += added_noise 496 | x_prime = self.decoder(encoding) 497 | if self.dim == 6: 498 | x_normal = x_prime[:, 3:6, :] 499 | x_normal = F.normalize(x_normal) 500 | result = torch.cat((x_prime[:, 0:3, :], x_normal), dim=1) 501 | else: 502 | result = x_prime 503 | 504 | return result 505 | 506 | 507 | def save_results(self, path, data): 508 | results = data.cpu().data.numpy() 509 | results = results.transpose(0, 2, 1) 510 | save_objs(results, path) 511 | print "Points saved." 512 | 513 | 514 | class NormalReg(nn.Module): 515 | 516 | def __init__(self): 517 | super(NormalReg, self).__init__() 518 | 519 | def forward(self, x): 520 | 521 | mean = torch.mean(x, dim=0).pow(2) 522 | cov = Ops.cov(x) 523 | 524 | cov_loss = torch.mean( 525 | (Variable(torch.eye(cov.size()[0]).cuda())-cov) 526 | .pow(2)) 527 | 528 | return torch.mean(mean) + cov_loss 529 | 530 | 531 | class PointCloudVAE(PointCloudAutoEncoder): 532 | 533 | def __init__(self, size, dim, batch_size=64, enc_size=100, kernel_size=16, 534 | reg_fn=NormalReg(), 535 | noise = 0, 536 | name="PCVAE"): 537 | super(PointCloudVAE, self).__init__(size, dim, batch_size, enc_size, kernel_size, 538 | noise=noise, name=name) 539 | self.reg_fn = reg_fn 540 | self.noise = torch.FloatTensor(self.batch_size, self.enc_size) 541 | 542 | 543 | def encoding_regularizer(self, x): 544 | return self.reg_fn(self.encoder(x)) 545 | 546 | 547 | def sample(self): 548 | self.noise.normal_() 549 | return self.decoder(Variable(self.noise.cuda())) 550 | 551 | 552 | class MultiResVAE(Model): 553 | 554 | def __init__(self, size, dim, batch_size=64, enc_size=100, kernel_size=2, 555 | reg_fn=NormalReg(), 556 | noise = 0, 557 | name="MLVAE"): 558 | super(MultiResVAE, self).__init__(name) 559 | 560 | self.reg_fn = reg_fn 561 | 562 | self.size = size 563 | self.dim = dim 564 | self.enc_size = enc_size 565 | self.batch_size = batch_size 566 | self.kernel_size = kernel_size 567 | self.enc_modules = nn.ModuleList() 568 | self.dec_modules = nn.ModuleList() 569 | self.upsample = Ops.NNUpsample1d() 570 | self.pool = nn.MaxPool1d(kernel_size=4, stride=4) 571 | self.noise_factor = noise 572 | 573 | self.enc_noise = torch.FloatTensor(self.batch_size, self.enc_size) 574 | 575 | custom_nfilters = [3, 32, 64, 128, 256, 512, 512, 1024, 1024, 1024] 576 | custom_nfilters = np.array(custom_nfilters) 577 | custom_nfilters[1:] /= 2 578 | self.last_size = 16 579 | 580 | self.noise = torch.FloatTensor(self.batch_size, self.enc_size) 581 | 582 | current_size = self.size 583 | layer_num = 1 584 | padding = (self.kernel_size - 1)/2 585 | n_channels = [] 586 | n_channels.append(custom_nfilters[layer_num-1]) 587 | while current_size > self.last_size: 588 | in_channels = custom_nfilters[layer_num-1] 589 | out_channels = custom_nfilters[layer_num] 590 | conv_enc = MultiResConv1d('down{}'.format(layer_num), 591 | in_channels, out_channels) 592 | current_size /= 2 593 | in_channels = out_channels 594 | n_channels.append(out_channels) 595 | layer_num += 1 596 | 597 | self.enc_modules.append(conv_enc) 598 | 599 | self.enc_fc = nn.Linear(3*self.last_size*in_channels, self.enc_size) 600 | #self.enc_fc_mean = nn.Linear(3*self.last_size*in_channels, self.enc_size) 601 | #self.enc_fc_var = nn.Linear(3*self.last_size*in_channels, self.enc_size) 602 | self.dec_fc = nn.Linear(self.enc_size, self.last_size*n_channels[-1]) 603 | 604 | self.final_feature = 128 605 | n_channels.reverse() 606 | n_channels[-1] = self.final_feature 607 | current_size = self.last_size 608 | layer_num = 1 609 | padding = (self.kernel_size - 1)/2 610 | while current_size < self.size: 611 | in_channels = n_channels[layer_num-1] 612 | out_channels = n_channels[layer_num] 613 | conv_dec = MultiResConvTranspose1d('up{}'.format(layer_num), 614 | in_channels, out_channels) 615 | current_size *= 2 616 | in_channels = out_channels 617 | layer_num += 1 618 | 619 | self.dec_modules.append(conv_dec) 620 | 621 | self.final_conv = nn.Sequential() 622 | self.final_conv.add_module('final_conv1', 623 | nn.ConvTranspose1d(self.final_feature*3, 128, 1, stride=1, padding=0)) 624 | self.final_conv.add_module('bn_final', 625 | nn.BatchNorm1d(128)) 626 | self.final_conv.add_module('relu_final', 627 | nn.ReLU(inplace=True)) 628 | self.final_conv.add_module('final_conv2', 629 | nn.ConvTranspose1d(128, 3, 1, stride=1, padding=0)) 630 | self.final_conv.add_module('tanh_final', 631 | nn.Tanh()) 632 | 633 | 634 | def enc_forward(self, x): 635 | x0 = x 636 | x1 = self.pool(x) 637 | x2 = self.pool(x1) 638 | 639 | enc_tensors = [] 640 | enc_tensors.append([x0, x1, x2]) 641 | 642 | for enc_op in self.enc_modules: 643 | enc_tensors.append(enc_op(enc_tensors[-1])) 644 | 645 | t0 = enc_tensors[-1][0] 646 | t1 = self.upsample(enc_tensors[-1][1]) 647 | t2 = self.upsample(self.upsample(enc_tensors[-1][2])) 648 | t = torch.cat((t0, t1, t2), 1).view(self.batch_size, -1) 649 | 650 | encoding = self.enc_fc(t) 651 | return encoding, enc_tensors 652 | #encoding_mean = self.enc_fc_mean(t) 653 | #encoding_var = self.enc_fc_var(t) 654 | #return (encoding_mean, encoding_var) 655 | 656 | 657 | def dec_forward(self, x): 658 | 659 | mr_enc0 = self.dec_fc(x).view(self.batch_size, -1, self.last_size) 660 | mr_enc1 = self.pool(mr_enc0) 661 | mr_enc2 = self.pool(mr_enc1) 662 | mr_enc = [mr_enc0, mr_enc1, mr_enc2] 663 | 664 | dec_tensors = [] 665 | dec_tensors.append(mr_enc) 666 | 667 | for i in xrange(0, len(self.dec_modules)-1): 668 | dec_tensors.append(self.dec_modules[i](dec_tensors[-1])) 669 | 670 | conv_out = self.dec_modules[-1](dec_tensors[-1]) 671 | out0 = conv_out[0] 672 | out1 = self.upsample(conv_out[1]) 673 | out2 = self.upsample(self.upsample(conv_out[2])) 674 | out = torch.cat((out0, out1, out2), 1) 675 | return self.final_conv(out) 676 | 677 | # 678 | # def reparameterize(self, mu, logvar): 679 | # if self.training: 680 | # std = logvar.mul(0.5).exp_() 681 | # eps = Variable(std.data.new(std.size()).normal_()) 682 | # return eps.mul(std).add_(mu) 683 | # else: 684 | # return mu 685 | # 686 | 687 | def forward(self, x): 688 | encoding = self.enc_forward(x)[0] 689 | self.enc_noise.normal_() 690 | 691 | added_noise = Variable(self.noise_factor*self.enc_noise.cuda()) 692 | 693 | encoding += added_noise 694 | return self.dec_forward(encoding) 695 | 696 | 697 | def encoding_regularizer(self, x): 698 | return self.reg_fn(self.enc_forward(x)[0]) 699 | 700 | 701 | def sample(self): 702 | self.noise.normal_() 703 | return self.dec_forward(Variable(self.noise.cuda())) 704 | 705 | 706 | def save_results(self, path, data, start_idx=0): 707 | results = data.cpu().data.numpy() 708 | results = results.transpose(0, 2, 1) 709 | save_objs(results, path, start_idx) 710 | print "Points saved." 711 | 712 | 713 | class EncodingSVM(Model): 714 | 715 | def __init__(self, enc_size, n_classes, ae_model, batch_size, name="EncSVM"): 716 | super(EncodingSVM, self).__init__(name) 717 | 718 | self.batch_size = batch_size 719 | self.enc_size = enc_size 720 | self.n_classes = n_classes 721 | self.ae_model = ae_model 722 | 723 | self.upsample = Ops.NNUpsample1d() 724 | 725 | alpha = 32 726 | self.pools = [] 727 | self.pools.append(nn.MaxPool1d(kernel_size=alpha, stride=alpha)) 728 | self.pools.append(nn.MaxPool1d(kernel_size=alpha/2, stride=alpha/2)) 729 | self.pools.append(nn.MaxPool1d(kernel_size=alpha/4, stride=alpha/4)) 730 | #self.pools.append(nn.MaxPool1d(kernel_size=alpha/8, stride=alpha/8)) 731 | 732 | self.fc = nn.Linear(self.enc_size, self.n_classes) 733 | 734 | def forward(self, x): 735 | enc, features = self.ae_model.enc_forward(x) 736 | descriptor = [] 737 | for i, p in enumerate(self.pools): 738 | t0 = p(features[i][0]) 739 | t1 = self.upsample(p(features[i][1])) 740 | t2 = self.upsample(self.upsample(p(features[i][2]))) 741 | descriptor.append(torch.cat((t0, t1, t2), 1)) 742 | 743 | descriptor = torch.cat(descriptor, 1) 744 | descriptor = descriptor.view(self.batch_size, -1) 745 | 746 | return self.fc(descriptor) 747 | -------------------------------------------------------------------------------- /models/GAN.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from Model import Model 6 | from tools.PointCloudDataset import save_objs 7 | from tools import Ops 8 | 9 | 10 | class LinearDiscriminator(Model): 11 | 12 | def __init__(self, input_size=256, n_layers=3, layer_size=256, 13 | name="LinearGenerator"): 14 | super(LinearDiscriminator, self).__init__(name) 15 | self.n_layers = n_layers 16 | self.layer_size = layer_size 17 | self.activation = nn.LeakyReLU(0.2, inplace=True) 18 | self.bn = nn.BatchNorm1d(layer_size) 19 | 20 | self.layers = [] 21 | self.bns = [] 22 | for i in range(n_layers): 23 | if i == 0: 24 | self.layers.append(nn.Linear(input_size, layer_size)) 25 | else: 26 | self.layers.append(nn.Linear(input_size, layer_size)) 27 | self.bns.append(nn.BatchNorm1d(layer_size)) 28 | 29 | self.out = nn.Linear(layer_size, 1) 30 | 31 | 32 | def cuda(self): 33 | super(LinearDiscriminator, self).cuda() 34 | for l in self.layers: 35 | l.cuda() 36 | for b in self.bns: 37 | b.cuda() 38 | 39 | 40 | def forward(self, x): 41 | 42 | features = [] 43 | for i, layer in enumerate(self.layers): 44 | if i == 0: 45 | features.append(self.activation(self.bns[i](layer(x)))) 46 | else: 47 | features.append(self.activation(self.bns[i](layer(features[-1])))) 48 | output = self.out(features[-1]) 49 | return (output, features) 50 | 51 | 52 | class LinearGenerator(Model): 53 | 54 | def __init__(self, enc_size=100, output_size=256, n_layers=3, layer_size=256, 55 | name="LinearGenerator"): 56 | super(LinearGenerator, self).__init__(name) 57 | self.n_layers = n_layers 58 | self.layer_size = layer_size 59 | self.enc_size = enc_size 60 | self.activation = nn.LeakyReLU(0.2, inplace=True) 61 | self.bn = nn.BatchNorm1d(layer_size) 62 | 63 | self.net = nn.Sequential() 64 | for i in range(n_layers): 65 | if i == 0: 66 | self.net.add_module("linear{}".format(i), 67 | nn.Linear(enc_size, layer_size)) 68 | else: 69 | self.net.add_module("linear{}".format(i), 70 | nn.Linear(layer_size, layer_size)) 71 | self.net.add_module("bn{}".format(i), 72 | nn.BatchNorm1d(layer_size)) 73 | self.net.add_module("relu{}".format(i), 74 | nn.ReLU()) 75 | 76 | self.out = nn.Linear(layer_size, output_size) 77 | 78 | 79 | def forward(self, x): 80 | return self.out(self.net(x)) 81 | 82 | 83 | def save_results(self, path, data): 84 | results = data.cpu().data.numpy() 85 | results = results.transpose(0, 2, 1) 86 | save_objs(results, path) 87 | print "Points saved." 88 | 89 | 90 | class DiscriminatorBCELoss(nn.Module): 91 | 92 | def __init__(self): 93 | super(DiscriminatorBCELoss, self).__init__() 94 | self.BCE = nn.BCELoss() 95 | 96 | 97 | def forward(self, x): 98 | input_data, target = x 99 | input_logits, _ = input_data 100 | 101 | return self.BCE(F.sigmoid(input_logits), target) 102 | 103 | 104 | class GeneratorFeatureLoss(nn.Module): 105 | 106 | def __init__(self): 107 | super(GeneratorFeatureLoss, self).__init__() 108 | 109 | 110 | def forward(self, x): 111 | real, fake = x 112 | _, real_features = real 113 | _, fake_features = fake 114 | 115 | real_features = torch.cat(real_features, dim=1) 116 | fake_features = torch.cat(fake_features, dim=1) 117 | 118 | real_mean = torch.mean(real_features, dim=0) 119 | fake_mean = torch.mean(fake_features, dim=0) 120 | 121 | real_cov = Ops.cov(real_features) 122 | fake_cov = Ops.cov(fake_features) 123 | 124 | mean_loss = torch.sum((real_mean-fake_mean).pow(2)) 125 | cov_loss = torch.sum((real_cov-fake_cov).pow(2)) 126 | 127 | return mean_loss + cov_loss 128 | 129 | 130 | -------------------------------------------------------------------------------- /models/ImageToShape.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from torch.autograd import Variable 7 | import torchvision 8 | 9 | from Model import Model 10 | from tools.PointCloudDataset import save_objs 11 | from tools import Ops 12 | import tools.DataVis as DataVis 13 | 14 | from AutoEncoder import MultiResConvTranspose1d 15 | 16 | 17 | class MultiResImageToShape(Model): 18 | 19 | def __init__(self, size, dim, batch_size=64, kernel_size=2, name="MRI2S", 20 | pretrained=False, arch=True): 21 | super(MultiResImageToShape, self).__init__(name) 22 | 23 | self.size = size 24 | self.dim = dim 25 | self.kernel_size = kernel_size 26 | self.batch_size = batch_size 27 | if arch == 'vgg': 28 | self.encoder = torchvision.models.vgg11(pretrained=pretrained) 29 | elif arch == 'alexnet': 30 | self.encoder = torchvision.models.alexnet(pretrained=pretrained) 31 | self.encoder.classifier._modules['6'] = nn.Linear(4096, 16*1024) 32 | self.dec_modules = nn.ModuleList() 33 | self.base_size = 16 34 | 35 | self.upsample = Ops.NNUpsample1d() 36 | self.pool = nn.MaxPool1d(kernel_size=4, stride=4) 37 | 38 | custom_nfilters = [128, 128, 128, 256, 512, 512, 1024, 1024, 1024] 39 | custom_nfilters.reverse() 40 | custom_nfilters = np.array(custom_nfilters) 41 | 42 | current_size = self.base_size 43 | layer_num = 1 44 | padding = (self.kernel_size - 1)/2 45 | while current_size < self.size: 46 | in_channels = custom_nfilters[layer_num-1] 47 | print in_channels 48 | out_channels = custom_nfilters[layer_num] 49 | conv_dec = MultiResConvTranspose1d('up{}'.format(layer_num), 50 | in_channels, out_channels) 51 | current_size *= 2 52 | in_channels = out_channels 53 | layer_num += 1 54 | 55 | self.dec_modules.append(conv_dec) 56 | 57 | self.final_conv = nn.Sequential() 58 | self.final_conv.add_module('final_conv1', 59 | nn.ConvTranspose1d(custom_nfilters[-1]*3, 128, 1, stride=1, padding=0)) 60 | self.final_conv.add_module('bn_final', 61 | nn.BatchNorm1d(128)) 62 | self.final_conv.add_module('relu_final', 63 | nn.ReLU(inplace=True)) 64 | self.final_conv.add_module('final_conv2', 65 | nn.ConvTranspose1d(128, 3, 1, stride=1, padding=0)) 66 | self.final_conv.add_module('tanh_final', 67 | nn.Tanh()) 68 | 69 | 70 | def forward(self, x): 71 | mr_enc0 = self.encoder(x).view(self.batch_size, -1, self.base_size) 72 | mr_enc1 = self.pool(mr_enc0) 73 | mr_enc2 = self.pool(mr_enc1) 74 | mr_enc = [mr_enc0, mr_enc1, mr_enc2] 75 | 76 | dec_tensors = [] 77 | dec_tensors.append(mr_enc) 78 | 79 | for i in xrange(0, len(self.dec_modules)-1): 80 | dec_tensors.append(self.dec_modules[i](dec_tensors[-1])) 81 | 82 | conv_out = self.dec_modules[-1](dec_tensors[-1]) 83 | out0 = conv_out[0] 84 | out1 = self.upsample(conv_out[1]) 85 | out2 = self.upsample(self.upsample(conv_out[2])) 86 | out = torch.cat((out0, out1, out2), 1) 87 | return self.final_conv(out) 88 | 89 | 90 | def save_results(self, path, data): 91 | results = data.cpu().data.numpy() 92 | results = results.transpose(0, 2, 1) 93 | save_objs(results, path) 94 | print "Points saved." 95 | 96 | 97 | -------------------------------------------------------------------------------- /models/MRTDecoder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from torch.autograd import Variable 7 | import torchvision 8 | 9 | from Model import Model 10 | from tools.PointCloudDataset import save_objs 11 | from tools import Ops 12 | import tools.DataVis as DataVis 13 | 14 | from AutoEncoder import MultiResConvTranspose1d 15 | from AutoEncoder import MultiResConv1d 16 | 17 | 18 | class FoldingNet(Model): 19 | def __init__(self, size, batch_size=64, name="FoldingNet"): 20 | super(FoldingNet, self).__init__(name) 21 | 22 | self.fold = nn.Sequential( 23 | nn.Conv1d(2+1024, 1024, 1, stride=1, padding=0), 24 | nn.BatchNorm1d(1024), 25 | nn.ReLU(inplace=True), 26 | 27 | nn.Conv1d(1024, 512, 1, stride=1, padding=0), 28 | nn.BatchNorm1d(512), 29 | nn.ReLU(inplace=True), 30 | 31 | nn.Conv1d(512, 256, 1, stride=1, padding=0), 32 | nn.BatchNorm1d(256), 33 | nn.ReLU(inplace=True), 34 | 35 | nn.Conv1d(256, 128, 1, stride=1, padding=0), 36 | nn.BatchNorm1d(128), 37 | nn.ReLU(inplace=True), 38 | 39 | nn.Conv1d(128, 3, 1, stride=1, padding=0), 40 | nn.Tanh()) 41 | 42 | #gd = int(np.sqrt(size)) 43 | #grid = np.indices((gd,gd)).T.reshape(-1, 2).T.astype('float32') 44 | #grid /= gd-1 45 | #self.z = Variable(torch.from_numpy(grid).unsqueeze(0).cuda()) 46 | self.z = Variable(torch.rand(1, 2, size).cuda()) 47 | self.global_feat = nn.Parameter(torch.rand(1, 1024, 1)) 48 | 49 | def forward(self): 50 | #self.z.data.uniform_(0, 1) 51 | inp = torch.cat((self.z, self.global_feat.expand(-1, -1, 1024)), 1) 52 | return self.fold(inp) 53 | 54 | 55 | class UNetMRTDecoder(Model): 56 | 57 | def __init__(self, size, dim, batch_size=64, enc_size=100, 58 | kernel_size=2, 59 | axis_file='rpt_axis.npy', 60 | name="PointSeg"): 61 | super(UNetMRTDecoder, self).__init__(name) 62 | 63 | self.init_channels = 128 64 | self.size = size 65 | self.dim = dim 66 | self.batch_size = batch_size 67 | self.kernel_size = kernel_size 68 | self.enc_size = enc_size 69 | self.enc_modules = nn.ModuleList() 70 | self.dec_modules = nn.ModuleList() 71 | self.upsample = Ops.NNUpsample1d() 72 | self.pool = nn.AvgPool1d(kernel_size=4, stride=4) 73 | #self.z = nn.Parameter(torch.randn(self.size*3).view(batch_size, 3, -1)) 74 | self.z = Variable(torch.randn(self.size*3).view(batch_size, 3, -1)).cuda() 75 | 76 | #custom_nfilters = [3, 128, 128, 128, 256, 265, 256, 77 | # 512, 512, 512, 1024, 1024, 2048] 78 | custom_nfilters = [3, 4, 8, 16, 32, 64, 128, 79 | 128, 128, 128, 256, 256, 256] 80 | custom_nfilters = np.array(custom_nfilters) 81 | #custom_nfilters[1:] /= 4 82 | 83 | current_size = self.size 84 | layer_num = 1 85 | padding = (self.kernel_size - 1)/2 86 | n_channels = [] 87 | while current_size > 64: 88 | in_channels = custom_nfilters[layer_num-1] 89 | out_channels = custom_nfilters[layer_num] 90 | conv_enc = MultiResConv1d('down{}'.format(layer_num), 91 | in_channels, out_channels) 92 | # conv_enc = nn.Sequential() 93 | # conv_enc.add_module('conv{}'.format(layer_num), 94 | # nn.Conv1d(in_channels, out_channels, self.kernel_size, 95 | # stride=2, 96 | # padding=padding)) 97 | # #axis=self.axis_on_level(layer_num-1))) 98 | # conv_enc.add_module('bn{}'.format(layer_num), 99 | # nn.BatchNorm1d(out_channels)) 100 | # conv_enc.add_module('lrelu{}'.format(layer_num), 101 | # nn.LeakyReLU(0.2, inplace=True)) 102 | current_size /= 2 103 | in_channels = out_channels 104 | n_channels.append(out_channels) 105 | if out_channels < 1024: 106 | out_channels *= 2 107 | layer_num += 1 108 | 109 | self.enc_modules.append(conv_enc) 110 | 111 | n_channels.reverse() 112 | current_size = 64 113 | layer_num = 1 114 | padding = (self.kernel_size - 1)/2 115 | while current_size < self.size//2: 116 | if layer_num == 1: 117 | in_channels = n_channels[layer_num-1] 118 | else: 119 | in_channels = n_channels[layer_num-1]*2 120 | out_channels = n_channels[layer_num] 121 | 122 | conv_dec = MultiResConvTranspose1d('up{}'.format(layer_num), 123 | in_channels, out_channels) 124 | # conv_dec = nn.Sequential() 125 | # conv_dec.add_module('conv{}'.format(layer_num), 126 | # nn.ConvTranspose1d(in_channels, 127 | # out_channels, 128 | # self.kernel_size, 129 | # stride=2, 130 | # padding=padding)) 131 | # conv_dec.add_module('bn{}'.format(layer_num), 132 | # nn.BatchNorm1d(out_channels)) 133 | # conv_dec.add_module('relu{}'.format(layer_num), 134 | # nn.ReLU(inplace=True)) 135 | 136 | current_size *= 2 137 | in_channels = out_channels 138 | layer_num += 1 139 | 140 | self.dec_modules.append(conv_dec) 141 | 142 | conv_dec = MultiResConvTranspose1d('up{}'.format(layer_num), 143 | in_channels, 256) 144 | self.dec_modules.append(conv_dec) 145 | 146 | # conv_dec = nn.Sequential() 147 | # conv_dec.add_module('conv{}'.format(layer_num), 148 | # nn.ConvTranspose1d(in_channels, 256, self.kernel_size, 149 | # stride=2, 150 | # padding=padding)) 151 | # conv_dec.add_module('bn{}'.format(layer_num), 152 | # nn.BatchNorm1d(256)) 153 | # conv_dec.add_module('relu{}'.format(layer_num), 154 | # nn.ReLU(inplace=True)) 155 | 156 | self.final_conv = nn.Sequential() 157 | self.final_conv.add_module('final_conv1', 158 | nn.ConvTranspose1d(256*3, 128, 1, stride=1, padding=0)) 159 | self.final_conv.add_module('bn_final', 160 | nn.BatchNorm1d(128)) 161 | self.final_conv.add_module('relu_final', 162 | nn.ReLU(inplace=True)) 163 | self.final_conv.add_module('final_conv2', 164 | nn.ConvTranspose1d(128, 3, 1, stride=1, padding=0)) 165 | self.final_conv.add_module('tanh_final', 166 | nn.Tanh()) 167 | 168 | def multires_cat(self, x, y): 169 | out0 = torch.cat((x[0], y[0]), 1) 170 | out1 = torch.cat((x[1], y[1]), 1) 171 | out2 = torch.cat((x[2], y[2]), 1) 172 | 173 | return [out0, out1, out2] 174 | 175 | def forward(self): 176 | x0 = self.z 177 | x1 = self.pool(x0) 178 | x2 = self.pool(x1) 179 | 180 | enc_tensors = [] 181 | enc_tensors.append([x0, x1, x2]) 182 | 183 | for enc_op in self.enc_modules: 184 | enc_tensors.append(enc_op(enc_tensors[-1])) 185 | # 186 | # for t in enc_tensors: 187 | # print t.size() 188 | # 189 | # print self.dec_modules 190 | # 191 | #t = enc_tensors[-1].view(self.batch_size, -1) 192 | #encoding = self.enc_fc(t) 193 | 194 | dec_tensors = [] 195 | #dec_tensors.append(self.dec_fc(encoding).view(self.batch_size, -1, 16)) 196 | dec_tensors.append(self.dec_modules[0](enc_tensors[-1])) 197 | 198 | for i in xrange(1, len(self.dec_modules)-1): 199 | in_tensor = enc_tensors[-(i+1)] 200 | #in_tensor = torch.cat((dec_tensors[-1], in_tensor), 1) 201 | in_tensor = self.multires_cat(in_tensor, dec_tensors[-1]) 202 | dec_tensors.append(self.dec_modules[i](in_tensor)) 203 | 204 | conv_out = self.dec_modules[-1](dec_tensors[-1]) 205 | 206 | out0 = conv_out[0] 207 | out1 = self.upsample(conv_out[1]) 208 | out2 = self.upsample(self.upsample(conv_out[2])) 209 | 210 | out = torch.cat((out0, out1, out2), 1) 211 | 212 | return self.final_conv(out) 213 | 214 | 215 | def axis_on_level(self, l): 216 | nlevels = np.log2(self.axis.shape[0]+1) 217 | level = nlevels - l - 1 218 | a = int(2**level - 1) 219 | b = int(2**(level+1) - 1) 220 | return self.axis[a:b, :] 221 | 222 | 223 | def save_results(self, path, data): 224 | results = data.cpu().data.numpy() 225 | results = results.transpose(0, 2, 1) 226 | save_segs(results, path) 227 | print "Segmentations saved." 228 | 229 | 230 | 231 | class MRTDecoder(Model): 232 | 233 | def __init__(self, size, dim, batch_size=64, kernel_size=2, name="MRTDecoder"): 234 | super(MRTDecoder, self).__init__(name) 235 | 236 | self.size = size 237 | self.dim = dim 238 | self.kernel_size = kernel_size 239 | self.batch_size = batch_size 240 | 241 | self.z = nn.Parameter(torch.randn(16*1024)) 242 | self.dec_modules = nn.ModuleList() 243 | self.base_size = 16 244 | 245 | self.upsample = Ops.NNUpsample1d() 246 | self.pool = nn.MaxPool1d(kernel_size=4, stride=4) 247 | 248 | custom_nfilters = [128, 128, 128, 256, 512, 512, 1024, 1024, 1024] 249 | custom_nfilters.reverse() 250 | custom_nfilters = np.array(custom_nfilters) 251 | custom_nfilters[1:] /= 2 252 | 253 | current_size = self.base_size 254 | layer_num = 1 255 | padding = (self.kernel_size - 1)/2 256 | while current_size < self.size: 257 | in_channels = custom_nfilters[layer_num-1] 258 | out_channels = custom_nfilters[layer_num] 259 | conv_dec = MultiResConvTranspose1d('up{}'.format(layer_num), 260 | in_channels, out_channels) 261 | current_size *= 2 262 | in_channels = out_channels 263 | layer_num += 1 264 | 265 | self.dec_modules.append(conv_dec) 266 | 267 | self.final_conv = nn.Sequential() 268 | self.final_conv.add_module('final_conv1', 269 | nn.ConvTranspose1d(custom_nfilters[-1]*3, 128, 1, stride=1, padding=0)) 270 | self.final_conv.add_module('bn_final', 271 | nn.BatchNorm1d(128)) 272 | self.final_conv.add_module('relu_final', 273 | nn.ReLU(inplace=True)) 274 | self.final_conv.add_module('final_conv2', 275 | nn.ConvTranspose1d(128, 3, 1, stride=1, padding=0)) 276 | self.final_conv.add_module('tanh_final', 277 | nn.Tanh()) 278 | 279 | 280 | def forward(self): 281 | mr_enc0 = self.z.view(self.batch_size, -1, self.base_size) 282 | mr_enc1 = self.pool(mr_enc0) 283 | mr_enc2 = self.pool(mr_enc1) 284 | mr_enc = [mr_enc0, mr_enc1, mr_enc2] 285 | 286 | dec_tensors = [] 287 | dec_tensors.append(mr_enc) 288 | 289 | for i in xrange(0, len(self.dec_modules)-1): 290 | dec_tensors.append(self.dec_modules[i](dec_tensors[-1])) 291 | 292 | conv_out = self.dec_modules[-1](dec_tensors[-1]) 293 | out0 = conv_out[0] 294 | out1 = self.upsample(conv_out[1]) 295 | out2 = self.upsample(self.upsample(conv_out[2])) 296 | out = torch.cat((out0, out1, out2), 1) 297 | return self.final_conv(out) 298 | 299 | 300 | def save_results(self, path, data): 301 | results = data.cpu().data.numpy() 302 | results = results.transpose(0, 2, 1) 303 | save_objs(results, path) 304 | print "Points saved." 305 | 306 | 307 | class SRTDecoder(Model): 308 | 309 | def __init__(self, size, dim, batch_size=64, kernel_size=2, name="SRTDecoder"): 310 | super(SRTDecoder, self).__init__(name) 311 | 312 | self.size = size 313 | self.dim = dim 314 | self.kernel_size = kernel_size 315 | self.batch_size = batch_size 316 | 317 | self.z = nn.Parameter(torch.randn(16*1024)) 318 | self.dec_modules = nn.ModuleList() 319 | self.base_size = 16 320 | 321 | self.upsample = Ops.NNUpsample1d() 322 | self.pool = nn.MaxPool1d(kernel_size=4, stride=4) 323 | 324 | custom_nfilters = [128, 128, 128, 256, 512, 512, 1024, 1024, 1024] 325 | custom_nfilters.reverse() 326 | custom_nfilters = np.array(custom_nfilters) 327 | custom_nfilters[1:] /= 2 328 | 329 | self.conv_dec = nn.Sequential() 330 | 331 | current_size = self.base_size 332 | layer_num = 1 333 | padding = (self.kernel_size - 1)/2 334 | while current_size < self.size: 335 | in_channels = custom_nfilters[layer_num-1] 336 | out_channels = custom_nfilters[layer_num] 337 | 338 | self.conv_dec.add_module('conv{}'.format(layer_num), 339 | nn.ConvTranspose1d(in_channels, out_channels, self.kernel_size, 340 | stride=2, 341 | padding=padding)) 342 | self.conv_dec.add_module('bn{}'.format(layer_num), 343 | nn.BatchNorm1d(out_channels)) 344 | self.conv_dec.add_module('lrelu{}'.format(layer_num), 345 | nn.LeakyReLU(0.2, inplace=True)) 346 | 347 | current_size *= 2 348 | in_channels = out_channels 349 | layer_num += 1 350 | 351 | self.final_conv = nn.Sequential() 352 | self.final_conv.add_module('final_conv1', 353 | nn.ConvTranspose1d(custom_nfilters[-1], 128, 1, stride=1, padding=0)) 354 | self.final_conv.add_module('bn_final', 355 | nn.BatchNorm1d(128)) 356 | self.final_conv.add_module('relu_final', 357 | nn.ReLU(inplace=True)) 358 | self.final_conv.add_module('final_conv2', 359 | nn.ConvTranspose1d(128, 3, 1, stride=1, padding=0)) 360 | self.final_conv.add_module('tanh_final', 361 | nn.Tanh()) 362 | 363 | 364 | def forward(self): 365 | feat = self.z.view(self.batch_size, -1, self.base_size) 366 | feat = self.conv_dec(feat) 367 | out = self.final_conv(feat) 368 | return out 369 | 370 | 371 | def save_results(self, path, data): 372 | results = data.cpu().data.numpy() 373 | results = results.transpose(0, 2, 1) 374 | save_objs(results, path) 375 | print "Points saved." 376 | 377 | 378 | -------------------------------------------------------------------------------- /models/Model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import os 4 | import glob 5 | 6 | 7 | class Model(nn.Module): 8 | 9 | def __init__(self, name): 10 | super(Model, self).__init__() 11 | self.name = name 12 | 13 | 14 | def save(self, path, epoch=0): 15 | complete_path = os.path.join(path, self.name) 16 | if not os.path.exists(complete_path): 17 | os.makedirs(complete_path) 18 | torch.save(self.state_dict(), 19 | os.path.join(complete_path, 20 | "model-{}.pth".format(str(epoch).zfill(5)))) 21 | 22 | 23 | def save_results(self, path, data): 24 | raise NotImplementedError("Model subclass must implement this method.") 25 | 26 | 27 | def load(self, path, modelfile=None): 28 | complete_path = os.path.join(path, self.name) 29 | if not os.path.exists(complete_path): 30 | raise IOError("{} directory does not exist in {}".format(self.name, path)) 31 | 32 | if modelfile is None: 33 | model_files = glob.glob(complete_path+"/*") 34 | mf = max(model_files) 35 | else: 36 | mf = os.path.join(complete_path, modelfile) 37 | 38 | self.load_state_dict(torch.load(mf)) 39 | 40 | 41 | -------------------------------------------------------------------------------- /models/VoxUNet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from torch.autograd import Variable 7 | import torchvision 8 | 9 | from Model import Model 10 | from tools.PointCloudDataset import save_objs 11 | from tools import Ops 12 | import tools.DataVis as DataVis 13 | 14 | from AutoEncoder import MultiResConvTranspose1d 15 | 16 | 17 | class UpConv(nn.Module): 18 | 19 | def __init__(self, in_channels, out_channels, kernel_size=3): 20 | super(UpConv, self).__init__() 21 | padding = (kernel_size - 1)/2 22 | self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, 23 | padding=padding) 24 | self.upsample = nn.Upsample(scale_factor=2, mode='trilinear') 25 | 26 | def forward(self, x): 27 | return self.conv(self.upsample(x)) 28 | 29 | 30 | class VoxUNet(Model): 31 | 32 | def __init__(self, size, batch_size=64, kernel_size=3, name="VoxUNet"): 33 | super(VoxUNet, self).__init__(name) 34 | 35 | self.size = size 36 | self.kernel_size = kernel_size 37 | self.batch_size = batch_size 38 | 39 | self.z = nn.Parameter(torch.randn(1,1,size,size,size)) 40 | self.enc_modules = nn.ModuleList() 41 | self.dec_modules = nn.ModuleList() 42 | 43 | custom_nfilters = [1, 16, 32, 64, 128, 256, 512, 1024, 1024, 1024] 44 | #custom_nfilters.reverse() 45 | #custom_nfilters = np.array(custom_nfilters) 46 | #custom_nfilters[1:] /= 2 47 | 48 | nfilters = [] 49 | 50 | #Encoder creation 51 | current_size = size 52 | layer_num = 1 53 | padding = (self.kernel_size - 1)/2 54 | while current_size > 2: 55 | in_channels = custom_nfilters[layer_num-1] 56 | out_channels = custom_nfilters[layer_num] 57 | conv_enc = nn.Sequential( 58 | nn.Conv3d(in_channels, out_channels, 59 | kernel_size=self.kernel_size, 60 | padding=padding, 61 | stride=2), 62 | nn.BatchNorm3d(out_channels), 63 | nn.ReLU(inplace=False)) 64 | current_size /= 2 65 | in_channels = out_channels 66 | layer_num += 1 67 | nfilters.append(out_channels) 68 | 69 | self.enc_modules.append(conv_enc) 70 | 71 | nfilters.reverse() 72 | 73 | #Decoder creation 74 | current_size = 2 75 | layer_num = 1 76 | padding = (self.kernel_size - 1)/2 77 | while current_size < self.size: 78 | if layer_num == 1: 79 | in_channels = nfilters[layer_num-1] 80 | else: 81 | in_channels = nfilters[layer_num-1]*2 82 | out_channels = nfilters[layer_num] 83 | conv_dec = nn.Sequential( 84 | UpConv(in_channels, out_channels, 85 | kernel_size=self.kernel_size), 86 | nn.BatchNorm3d(out_channels), 87 | nn.ReLU(inplace=False)) 88 | current_size *= 2 89 | in_channels = out_channels 90 | layer_num += 1 91 | nfilters.append(out_channels) 92 | 93 | self.dec_modules.append(conv_dec) 94 | 95 | self.final_conv = nn.Sequential( 96 | UpConv(nfilters[-2]*2, 1, 97 | kernel_size=self.kernel_size), 98 | nn.Sigmoid()) 99 | 100 | 101 | def forward(self): 102 | #Encoding 103 | enc_tensors = [] 104 | enc_tensors.append(self.z) 105 | 106 | for enc_op in self.enc_modules: 107 | enc_tensors.append(enc_op(enc_tensors[-1])) 108 | 109 | #Decoding 110 | dec_tensors = [] 111 | dec_tensors.append(self.dec_modules[0](enc_tensors[-1])) 112 | print dec_tensors[0].size() 113 | 114 | for i in xrange(1, len(self.dec_modules)-1): 115 | in_tensor = enc_tensors[-(i+1)] 116 | in_tensor = torch.cat((in_tensor, dec_tensors[-1]), 1) 117 | dec_tensors.append(self.dec_modules[i](in_tensor)) 118 | 119 | final_input = torch.cat((dec_tensors[-1], enc_tensors[1]), 1) 120 | out = self.final_conv(final_input) 121 | 122 | return out 123 | 124 | 125 | def save_results(self, path, data): 126 | results = data.cpu().data.numpy() 127 | results = results.transpose(0, 2, 1) 128 | save_objs(results, path) 129 | print "Points saved." 130 | 131 | if __name__ == '__main__': 132 | net = VoxUNet(256, batch_size=1).cuda() 133 | net() 134 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/models/__init__.py -------------------------------------------------------------------------------- /modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/modules/__init__.py -------------------------------------------------------------------------------- /modules/nnd.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | from functions.nnd import NNDFunction 3 | 4 | class NNDModule(Module): 5 | def forward(self, input1, input2): 6 | return NNDFunction()(input1, input2) 7 | -------------------------------------------------------------------------------- /mr_train_4k.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.optim as optim 4 | import torch.nn as nn 5 | 6 | import argparse 7 | 8 | from tools.Trainer import VAETrainer 9 | from tools.PointCloudDataset import PointCloudDataset 10 | from models.AutoEncoder import MultiResVAE 11 | from models.AutoEncoder import ChamferLoss 12 | from models.AutoEncoder import ChamferWithNormalLoss 13 | from models.AutoEncoder import L2WithNormalLoss 14 | 15 | parser = argparse.ArgumentParser(description='Point Cloud Generator.') 16 | parser.add_argument("-d", "--datapath", type=str, help="Dataset path.", default="") 17 | parser.add_argument("-n", "--name", type=str, help="Name of the experiment", default="PointGen") 18 | parser.add_argument("-bs", "--batchSize", type=int, help="Batch size", default=64) 19 | parser.add_argument("-e", "--encSize", type=int, help="Encoding size", default=128) 20 | parser.add_argument("-f", "--factorNoise", type=float, help="Noise factor", default=0.0) 21 | parser.add_argument("--train", dest='train', action='store_true') 22 | parser.set_defaults(train=False) 23 | 24 | 25 | if __name__ == '__main__': 26 | args = parser.parse_args() 27 | 28 | vae = MultiResVAE(4096, 3, name=args.name, enc_size=args.encSize, 29 | noise=args.factorNoise, 30 | batch_size=args.batchSize) 31 | #vae.load('checkpoint') 32 | optimizer = optim.Adam(vae.parameters(), lr=1e-5) 33 | 34 | dataset = PointCloudDataset(args.datapath) 35 | loader = torch.utils.data.DataLoader(dataset, batch_size=args.batchSize, 36 | shuffle=True, num_workers=2) 37 | trainer = VAETrainer(vae, loader, optimizer, ChamferLoss()) 38 | trainer.train(2000) 39 | 40 | -------------------------------------------------------------------------------- /nndistance/.gitignore: -------------------------------------------------------------------------------- 1 | *.so 2 | *.o 3 | -------------------------------------------------------------------------------- /nndistance/README.md: -------------------------------------------------------------------------------- 1 | # An example C extension for PyTorch 2 | 3 | This example showcases adding a neural network layer that adds two input Tensors 4 | 5 | - src: C source code 6 | - functions: the autograd functions 7 | - modules: code of the nn module 8 | - build.py: a small file that compiles your module to be ready to use 9 | - test.py: an example file that loads and uses the extension 10 | 11 | ```bash 12 | cd src 13 | nvcc -c -o nnd_cuda.cu.o nnd_cuda.cu -x cu -Xcompiler -fPIC -arch=sm_52 14 | cd .. 15 | python build.py 16 | python test.py 17 | ``` 18 | -------------------------------------------------------------------------------- /nndistance/_ext/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/nndistance/_ext/__init__.py -------------------------------------------------------------------------------- /nndistance/_ext/my_lib/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._my_lib import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /nndistance/build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.utils.ffi import create_extension 4 | 5 | this_file = os.path.dirname(__file__) 6 | 7 | sources = ['src/my_lib.c'] 8 | headers = ['src/my_lib.h'] 9 | defines = [] 10 | with_cuda = False 11 | 12 | if torch.cuda.is_available(): 13 | print('Including CUDA code.') 14 | sources += ['src/my_lib_cuda.c'] 15 | headers += ['src/my_lib_cuda.h'] 16 | defines += [('WITH_CUDA', None)] 17 | with_cuda = True 18 | 19 | this_file = os.path.dirname(os.path.realpath(__file__)) 20 | print(this_file) 21 | extra_objects = ['src/nnd_cuda.cu.o'] 22 | extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] 23 | 24 | ffi = create_extension( 25 | '_ext.my_lib', 26 | headers=headers, 27 | sources=sources, 28 | define_macros=defines, 29 | relative_to=__file__, 30 | with_cuda=with_cuda, 31 | extra_objects=extra_objects 32 | ) 33 | 34 | if __name__ == '__main__': 35 | ffi.build() 36 | -------------------------------------------------------------------------------- /nndistance/functions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/nndistance/functions/__init__.py -------------------------------------------------------------------------------- /nndistance/functions/nnd.py: -------------------------------------------------------------------------------- 1 | # functions/add.py 2 | import torch 3 | from torch.autograd import Function 4 | from _ext import my_lib 5 | 6 | 7 | class NNDFunction(Function): 8 | def forward(self, xyz1, xyz2): 9 | batchsize, n, _ = xyz1.size() 10 | _, m, _ = xyz2.size() 11 | self.xyz1 = xyz1 12 | self.xyz2 = xyz2 13 | dist1 = torch.zeros(batchsize, n) 14 | dist2 = torch.zeros(batchsize, m) 15 | 16 | self.idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) 17 | self.idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) 18 | 19 | if not xyz1.is_cuda: 20 | my_lib.nnd_forward(xyz1, xyz2, dist1, dist2, self.idx1, self.idx2) 21 | else: 22 | dist1 = dist1.cuda() 23 | dist2 = dist2.cuda() 24 | self.idx1 = self.idx1.cuda() 25 | self.idx2 = self.idx2.cuda() 26 | my_lib.nnd_forward_cuda(xyz1, xyz2, dist1, dist2, self.idx1, self.idx2) 27 | 28 | self.dist1 = dist1 29 | self.dist2 = dist2 30 | 31 | #print(batchsize, n, m) 32 | 33 | return dist1, dist2 34 | 35 | def backward(self, graddist1, graddist2): 36 | #print(self.idx1, self.idx2) 37 | 38 | 39 | graddist1 = graddist1.contiguous() 40 | graddist2 = graddist2.contiguous() 41 | 42 | gradxyz1 = torch.zeros(self.xyz1.size()) 43 | gradxyz2 = torch.zeros(self.xyz2.size()) 44 | 45 | if not graddist1.is_cuda: 46 | my_lib.nnd_backward(self.xyz1, self.xyz2, gradxyz1, gradxyz2, graddist1, graddist2, self.idx1, self.idx2) 47 | else: 48 | gradxyz1 = gradxyz1.cuda() 49 | gradxyz2 = gradxyz2.cuda() 50 | my_lib.nnd_backward_cuda(self.xyz1, self.xyz2, gradxyz1, gradxyz2, graddist1, graddist2, self.idx1, self.idx2) 51 | 52 | return gradxyz1, gradxyz2 -------------------------------------------------------------------------------- /nndistance/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/nndistance/modules/__init__.py -------------------------------------------------------------------------------- /nndistance/modules/nnd.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | from functions.nnd import NNDFunction 3 | 4 | class NNDModule(Module): 5 | def forward(self, input1, input2): 6 | return NNDFunction()(input1, input2) 7 | -------------------------------------------------------------------------------- /nndistance/src/make.sh: -------------------------------------------------------------------------------- 1 | set -e 2 | nvcc nnd_cuda.cu -o nnd_cuda.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | -------------------------------------------------------------------------------- /nndistance/src/my_lib.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | 4 | void nnsearch(int b,int n,int m,const float * xyz1,const float * xyz2,float * dist,int * idx){ 5 | for (int i=0;isize[0]; 30 | int n = xyz1->size[1]; 31 | int m = xyz2->size[1]; 32 | 33 | //printf("%d %d %d\n", batchsize, n, m); 34 | 35 | float *xyz1_data = THFloatTensor_data(xyz1); 36 | float *xyz2_data = THFloatTensor_data(xyz2); 37 | float *dist1_data = THFloatTensor_data(dist1); 38 | float *dist2_data = THFloatTensor_data(dist2); 39 | int *idx1_data = THIntTensor_data(idx1); 40 | int *idx2_data = THIntTensor_data(idx2); 41 | 42 | nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); 43 | nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); 44 | 45 | return 1; 46 | } 47 | 48 | 49 | int nnd_backward(THFloatTensor *xyz1, THFloatTensor *xyz2, THFloatTensor *gradxyz1, THFloatTensor *gradxyz2, THFloatTensor *graddist1, THFloatTensor *graddist2, THIntTensor *idx1, THIntTensor *idx2) { 50 | 51 | int b = xyz1->size[0]; 52 | int n = xyz1->size[1]; 53 | int m = xyz2->size[1]; 54 | 55 | //printf("%d %d %d\n", batchsize, n, m); 56 | 57 | float *xyz1_data = THFloatTensor_data(xyz1); 58 | float *xyz2_data = THFloatTensor_data(xyz2); 59 | float *gradxyz1_data = THFloatTensor_data(gradxyz1); 60 | float *gradxyz2_data = THFloatTensor_data(gradxyz2); 61 | float *graddist1_data = THFloatTensor_data(graddist1); 62 | float *graddist2_data = THFloatTensor_data(graddist2); 63 | int *idx1_data = THIntTensor_data(idx1); 64 | int *idx2_data = THIntTensor_data(idx2); 65 | 66 | 67 | for (int i=0;i 2 | #include "nnd_cuda.h" 3 | 4 | 5 | 6 | extern THCState *state; 7 | 8 | 9 | int nnd_forward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *dist1, THCudaTensor *dist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2) { 10 | int success = 0; 11 | success = NmDistanceKernelLauncher(xyz1->size[0], 12 | xyz1->size[1], 13 | THCudaTensor_data(state, xyz1), 14 | xyz2->size[1], 15 | THCudaTensor_data(state, xyz2), 16 | THCudaTensor_data(state, dist1), 17 | THCudaIntTensor_data(state, idx1), 18 | THCudaTensor_data(state, dist2), 19 | THCudaIntTensor_data(state, idx2), 20 | THCState_getCurrentStream(state) 21 | ); 22 | //int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream) 23 | 24 | 25 | if (!success) { 26 | THError("aborting"); 27 | } 28 | return 1; 29 | } 30 | 31 | 32 | int nnd_backward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *gradxyz1, THCudaTensor *gradxyz2, THCudaTensor *graddist1, 33 | THCudaTensor *graddist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2) { 34 | 35 | int success = 0; 36 | success = NmDistanceGradKernelLauncher(xyz1->size[0], 37 | xyz1->size[1], 38 | THCudaTensor_data(state, xyz1), 39 | xyz2->size[1], 40 | THCudaTensor_data(state, xyz2), 41 | THCudaTensor_data(state, graddist1), 42 | THCudaIntTensor_data(state, idx1), 43 | THCudaTensor_data(state, graddist2), 44 | THCudaIntTensor_data(state, idx2), 45 | THCudaTensor_data(state, gradxyz1), 46 | THCudaTensor_data(state, gradxyz2), 47 | THCState_getCurrentStream(state) 48 | ); 49 | //int NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream) 50 | 51 | if (!success) { 52 | THError("aborting"); 53 | } 54 | 55 | return 1; 56 | } 57 | 58 | 59 | 60 | -------------------------------------------------------------------------------- /nndistance/src/my_lib_cuda.h: -------------------------------------------------------------------------------- 1 | int nnd_forward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *dist1, THCudaTensor *dist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2); 2 | 3 | 4 | int nnd_backward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *gradxyz1, THCudaTensor *gradxyz2, THCudaTensor *graddist1, THCudaTensor *graddist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2); 5 | 6 | -------------------------------------------------------------------------------- /nndistance/src/nnd_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include "nnd_cuda.h" 3 | 4 | 5 | 6 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 7 | const int batch=512; 8 | __shared__ float buf[batch*3]; 9 | for (int i=blockIdx.x;ibest){ 121 | result[(i*n+j)]=best; 122 | result_i[(i*n+j)]=best_i; 123 | } 124 | } 125 | __syncthreads(); 126 | } 127 | } 128 | } 129 | int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ 130 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 131 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 132 | 133 | cudaError_t err = cudaGetLastError(); 134 | if (err != cudaSuccess) { 135 | printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); 136 | //THError("aborting"); 137 | return 0; 138 | } 139 | return 1; 140 | 141 | 142 | } 143 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 144 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 167 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 168 | 169 | cudaError_t err = cudaGetLastError(); 170 | if (err != cudaSuccess) { 171 | printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); 172 | //THError("aborting"); 173 | return 0; 174 | } 175 | return 1; 176 | 177 | } 178 | 179 | -------------------------------------------------------------------------------- /nndistance/src/nnd_cuda.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 6 | 7 | int NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 8 | 9 | #ifdef __cplusplus 10 | } 11 | #endif -------------------------------------------------------------------------------- /nndistance/test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | 5 | from modules.nnd import NNDModule 6 | 7 | dist = NNDModule() 8 | 9 | p1 = torch.rand(10,1000,3) 10 | p2 = torch.rand(10,1500,3) 11 | points1 = Variable(p1,requires_grad = True) 12 | points2 = Variable(p2) 13 | dist1, dist2 = dist(points1, points2) 14 | print(dist1, dist2) 15 | loss = torch.sum(dist1) 16 | print(loss) 17 | loss.backward() 18 | print(points1.grad, points2.grad) 19 | 20 | 21 | points1 = Variable(p1.cuda(), requires_grad = True) 22 | points2 = Variable(p2.cuda()) 23 | dist1, dist2 = dist(points1, points2) 24 | print(dist1, dist2) 25 | loss = torch.sum(dist1) 26 | print(loss) 27 | loss.backward() 28 | print(points1.grad, points2.grad) -------------------------------------------------------------------------------- /run_img2pc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.optim as optim 4 | import torch.nn as nn 5 | 6 | import argparse 7 | import os 8 | 9 | from tools.Trainer import ImageToPCTrainer 10 | from tools.PointCloudDataset import ImageDataset 11 | from models.AutoEncoder import PointCloudVAE 12 | from models.AutoEncoder import ChamferLoss 13 | from models.AutoEncoder import ChamferWithNormalLoss 14 | from models.AutoEncoder import L2WithNormalLoss 15 | from models.ImageToShape import MultiResImageToShape 16 | from models.ImageToShape import SingleResImageToShape 17 | from models.ImageToShape import FCImageToShape 18 | 19 | parser = argparse.ArgumentParser(description='MultiResolution image to shape model.') 20 | parser.add_argument("-n", "--name", type=str, help="Name of the experiment.", default="MRI2PC") 21 | parser.add_argument("-a", "--arch", type=str, help="Encoder architecture.", default="vgg") 22 | parser.add_argument("-pt", "--pretrained", type=str, help="Use pretrained net", default="True") 23 | parser.add_argument("-c", "--category", type=str, help="Category code (all is possible)", default="all") 24 | parser.add_argument("--train", dest='train', action='store_true') 25 | parser.set_defaults(train=False) 26 | 27 | 28 | #Change this for a path with images you want to test 29 | image_datapath = "notseen_real" 30 | 31 | if __name__ == '__main__': 32 | args = parser.parse_args() 33 | 34 | ptrain = None 35 | if args.pretrained == "False": 36 | ptrain = False 37 | elif args.pretrained == "True": 38 | ptrain = True 39 | 40 | full_name = "{}_{}_{}_{}".format(args.name, args.category, args.arch, ptrain) 41 | #full_name = args.name 42 | print full_name 43 | 44 | 45 | #mri2pc = FCImageToShape(size=4096, dim=3, batch_size=1, 46 | # name=full_name, pretrained=ptrain, arch=args.arch) 47 | #mri2pc = SingleResImageToShape(size=4096, dim=3, batch_size=1, 48 | # name=full_name, pretrained=ptrain, arch=args.arch) 49 | mri2pc = MultiResImageToShape(size=4096, dim=3, batch_size=1, 50 | name=full_name, pretrained=ptrain, arch=args.arch) 51 | mri2pc.load('checkpoint') 52 | 53 | optimizer = optim.Adam(mri2pc.parameters(), lr=0.001) 54 | 55 | test_dataset = ImageDataset(image_datapath) 56 | test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, 57 | shuffle=True, num_workers=2) 58 | 59 | log_dir = os.path.join("log", full_name) 60 | if not os.path.exists(log_dir): 61 | os.makedirs(log_dir) 62 | 63 | trainer = ImageToPCTrainer(mri2pc, None, test_loader, 64 | optimizer, ChamferLoss(), log_dir=log_dir) 65 | trainer.run() 66 | 67 | -------------------------------------------------------------------------------- /sampler/README: -------------------------------------------------------------------------------- 1 | To compile: 2 | 3 | javac -cp .:./vecmath.jar cotSpectral.java 4 | 5 | To run: 6 | 7 | java -cp .:./vecmath.jar cotSpectral in.obj out.obj nsample [sorting] [dmetric] 8 | 9 | nsample: number of samples (e.g. 1024) 10 | sorting: choices are kd or xyz (kd is default) 11 | dmetric: choices are geodesic or random (geodesic is default) 12 | 13 | [] options are optional, if you don't provide these, it will use default value 14 | 15 | -------------------------------------------------------------------------------- /sampler/cotSpectral.java: -------------------------------------------------------------------------------- 1 | import java.io.*; 2 | 3 | import java.util.ArrayList; 4 | import java.util.Arrays; 5 | import java.util.Comparator; 6 | import java.util.HashMap; 7 | import java.util.Random; 8 | 9 | import javax.vecmath.Point3f; 10 | import javax.vecmath.Point3i; 11 | import javax.vecmath.Vector3f; 12 | 13 | class cotSpectral { 14 | 15 | /* 'input_verts' and 'input_faces' store the vertex and face data of the input mesh respectively 16 | * each element in 'input_verts' is a 3D points defining the vertex 17 | * every three integers in 'input_faces' define the indexes of the three vertices that make a triangle 18 | * there are in total input_faces.size()/3 triangles 19 | */ 20 | private static ArrayList input_verts = new ArrayList (); 21 | private static ArrayList input_faces = new ArrayList (); 22 | 23 | /* 'curr_verts' and 'curr_faces' store the subdivided mesh that will be displayed on the screen 24 | * the elements stored in these arrays are the same with the input mesh 25 | * 'curr_normals' stores the normal of each face and is necessary for shading calculation 26 | * you don't have to compute the normals yourself: once you have curr_verts and curr_faces 27 | * ready, you just need to call estimateFaceNormal to generate normal data 28 | */ 29 | private static ArrayList curr_verts = new ArrayList (); 30 | private static ArrayList curr_faces = new ArrayList (); 31 | private static ArrayList curr_normals = new ArrayList (); 32 | 33 | private static String dmetric = "geodesic"; 34 | private static String sortmethod = "kd"; 35 | 36 | private static float xmin, ymin, zmin; 37 | private static float xmax, ymax, zmax; 38 | private static String output_filename; 39 | 40 | private static ArrayList samples = null; 41 | private static ArrayList aspSamples = null; 42 | private static float cum_area[]; 43 | private static Vector3f normals[]; 44 | 45 | //minimum radius points must be apart from one another. 46 | private static float radius = 0.f, radiusSq = radius*radius; 47 | private static float rho = 0.75f; 48 | private static int NSamples = 256; 49 | 50 | /* load a simple .obj mesh from disk file 51 | * note that each face *must* be a triangle and cannot be a quad or other types of polygon 52 | */ 53 | private static void loadMesh(String filename) { 54 | BufferedReader in = null; 55 | try { 56 | in = new BufferedReader(new FileReader(filename)); 57 | } catch (IOException e) { 58 | System.out.println("Error reading from file " + filename); 59 | System.exit(0); 60 | } 61 | 62 | float x, y, z; 63 | int v1, v2, v3; 64 | //int dummy; 65 | String line; 66 | String[] tokens; 67 | String[] subtokens; 68 | try { 69 | while ((line = in.readLine()) != null) { 70 | if (line.length() == 0) 71 | continue; 72 | switch(line.charAt(0)) { 73 | case 'v': 74 | if (line.charAt(1) == 't') { 75 | continue; // do not process texture coords 76 | } 77 | if (line.charAt(1) == 'n') { 78 | continue; // do not process normal 79 | } 80 | tokens = line.split("[ ]+"); 81 | x = Float.valueOf(tokens[1]); 82 | y = Float.valueOf(tokens[2]); 83 | z = Float.valueOf(tokens[3]); 84 | input_verts.add(new Point3f(x, y, z)); 85 | break; 86 | case 'f': 87 | tokens = line.split("[ ]+"); 88 | /* when defining faces, .obj assumes the vertex index starts from 1 89 | * so we should subtract 1 from each index 90 | */ 91 | subtokens = tokens[1].split("/+"); 92 | v1 = Integer.valueOf(subtokens[0])-1; 93 | 94 | subtokens = tokens[2].split("/+"); 95 | v2 = Integer.valueOf(subtokens[0])-1; 96 | 97 | subtokens = tokens[3].split("/+"); 98 | v3 = Integer.valueOf(subtokens[0])-1; 99 | 100 | input_faces.add(v1); 101 | input_faces.add(v2); 102 | input_faces.add(v3); 103 | break; 104 | default: 105 | continue; 106 | } 107 | } 108 | in.close(); 109 | } catch(IOException e) { 110 | // error reading file 111 | } 112 | 113 | //System.out.println("Read " + input_verts.size() + " vertices and " + input_faces.size() + " faces."); 114 | System.out.print(filename+" "); 115 | 116 | int tris = input_faces.size()/3; 117 | cum_area = new float[tris]; 118 | normals = new Vector3f[tris]; 119 | float cur_area = 0; 120 | for(int i=0; i=4) sortmethod = args[3]; 175 | if(args.length>=5) dmetric = args[4]; 176 | 177 | loadMesh(inputFilename); 178 | computeBoundingBox(); 179 | 180 | curr_verts = input_verts; 181 | curr_faces = input_faces; 182 | estimateFaceNormal(); 183 | 184 | for (int i=0; i<1; ++i) 185 | { 186 | while(generateSamples()==false) { 187 | rho-=0.05; 188 | float total_area = cum_area[cum_area.length-1]; 189 | radius = (float)(rho*Math.pow(3.464101615f*NSamples/total_area, -0.5f)*2); 190 | radiusSq = radius*radius; 191 | } 192 | sortASPSamples(); 193 | saveSamples(true, i); 194 | } 195 | } 196 | 197 | private static String addId(String filepath, int id){ 198 | return filepath.replaceAll(".obj", "_"+String.valueOf(id)+".obj"); 199 | } 200 | 201 | private static void saveSamples(boolean write_normal, int id){ 202 | String samples_filename = addId(output_filename, id); 203 | File f = new File(samples_filename); 204 | f.getParentFile().mkdirs(); 205 | try{ 206 | BufferedWriter w = new BufferedWriter(new FileWriter(samples_filename)); 207 | for(ASPPoint s : aspSamples) { 208 | w.write("v "+s.point.x+" "+s.point.y+" "+s.point.z+"\n"); 209 | if(write_normal) w.write("vn "+s.normal.x+" "+s.normal.y+" "+s.normal.z+"\n"); 210 | } 211 | w.flush(); 212 | w.close(); 213 | } catch(Exception e){e.printStackTrace();} 214 | 215 | System.out.println(samples_filename+" saved"); 216 | } 217 | 218 | private static class AuxComparatorX implements Comparator { 219 | public int compare(ASPPoint o1, ASPPoint o2) { 220 | if (o1.point.x > o2.point.x) return 1; if (o1.point.x < o2.point.x) return -1; return 0; 221 | } 222 | } 223 | private static class AuxComparatorY implements Comparator { 224 | public int compare(ASPPoint o1, ASPPoint o2) { 225 | if (o1.point.y > o2.point.y) return 1; if (o1.point.y < o2.point.y) return -1; return 0; 226 | } 227 | } 228 | private static class AuxComparatorZ implements Comparator { 229 | public int compare(ASPPoint o1, ASPPoint o2) { 230 | if (o1.point.z > o2.point.z) return 1; if (o1.point.z < o2.point.z) return -1; return 0; 231 | } 232 | } 233 | private static class AuxComparatorXYZ implements Comparator { 234 | public int compare(ASPPoint o1, ASPPoint o2) { 235 | float xyz1 = o1.point.x+o1.point.y+o1.point.z; 236 | float xyz2 = o2.point.x+o2.point.y+o2.point.z; 237 | if (xyz1 > xyz2) return 1; if (xyz1 < xyz2) return -1; return 0; 238 | } 239 | } 240 | private static void sortASPSamples() { 241 | ASPPoint[] pts = new ASPPoint[aspSamples.size()]; 242 | for(int i=0;i(Arrays.asList(pts)); 250 | } 251 | private static void recursiveSortAspSamples(int start, int end, ASPPoint[] pts, int axis) { 252 | if(end-start<=1) return; 253 | if(axis==0) 254 | Arrays.sort(pts, start, end, new AuxComparatorX()); 255 | else if(axis==1) 256 | Arrays.sort(pts, start, end, new AuxComparatorY()); 257 | else 258 | Arrays.sort(pts, start, end, new AuxComparatorZ()); 259 | int mid=(start+end)/2; 260 | recursiveSortAspSamples(start, mid, pts, (axis+1)%3); 261 | recursiveSortAspSamples(mid, end, pts, (axis+1)%3); 262 | } 263 | 264 | private static class Grid { 265 | private static final int GRID_SIZE = 32; 266 | private final HashMap> grid = new HashMap>(); 267 | public Grid(){} 268 | public void insert(ASPPoint p){ 269 | Point3i g = new Point3i(); 270 | g.x = (int)(((p.point.x-xmin)/(xmax-xmin))*GRID_SIZE); 271 | g.y = (int)(((p.point.y-ymin)/(ymax-ymin))*GRID_SIZE); 272 | g.z = (int)(((p.point.z-zmin)/(zmax-zmin))*GRID_SIZE); 273 | ArrayList pts = grid.get(g); 274 | if(pts == null){pts = new ArrayList(); grid.put(g, pts);} 275 | pts.add(p); 276 | } 277 | public boolean legalPoint(ASPPoint p){ 278 | Point3i min = new Point3i(); 279 | Point3i max = new Point3i(); 280 | 281 | min.x = (int)(((p.point.x-radius-xmin)/(xmax-xmin))*GRID_SIZE); 282 | min.y = (int)(((p.point.y-radius-ymin)/(ymax-ymin))*GRID_SIZE); 283 | min.z = (int)(((p.point.z-radius-zmin)/(zmax-zmin))*GRID_SIZE); 284 | 285 | max.x = (int)(((p.point.x+radius-xmin)/(xmax-xmin))*GRID_SIZE); 286 | max.y = (int)(((p.point.y+radius-ymin)/(ymax-ymin))*GRID_SIZE); 287 | max.z = (int)(((p.point.z+radius-zmin)/(zmax-zmin))*GRID_SIZE); 288 | 289 | if (xmax == xmin) { 290 | min.x = 0; 291 | max.x = 0; 292 | } 293 | if (ymax == ymin) { 294 | min.y = 0; 295 | max.y = 0; 296 | } 297 | if (zmax == zmin) { 298 | min.z = 0; 299 | max.z = 0; 300 | } 301 | 302 | Point3i g = new Point3i(); 303 | for(g.x = min.x; g.x <= max.x; g.x++) 304 | for(g.y = min.y; g.y <= max.y; g.y++) 305 | for(g.z = min.z; g.z <= max.z; g.z++){ 306 | ArrayList pts = grid.get(g); 307 | if(pts != null) 308 | for(ASPPoint s : pts){ 309 | float distSq = s.point.distanceSquared(p.point); 310 | if(distSq > radiusSq) continue; 311 | float scaling_factor = 0.f; 312 | if(dmetric.equals("random")) scaling_factor = 10000.f; 313 | else scaling_factor = getScalingFactor(s, p); 314 | if(distSq*scaling_factor < radiusSq) return false; 315 | } 316 | } 317 | return true; 318 | } 319 | private void clearGrid(){ 320 | grid.clear(); 321 | } 322 | } 323 | 324 | private static float getScalingFactor(ASPPoint s, ASPPoint p){ 325 | float scaling_factor = 0; 326 | Vector3f D = new Vector3f(s.point.x - p.point.x, s.point.y - p.point.y, s.point.z - p.point.z); 327 | 328 | D.normalize(); 329 | float a = D.dot(s.normal), b = D.dot(p.normal); 330 | 331 | 332 | if (a == b) scaling_factor = (float)(1.0 / Math.sqrt(1.0 - a * a)); 333 | else scaling_factor = (float)((Math.asin(a) - Math.asin(b)) / (a - b)); 334 | return (float)scaling_factor * scaling_factor; 335 | 336 | /* 337 | if(a*b >= 0) return 1; 338 | a = Math.abs(a); b = Math.abs(b); 339 | float total = a+b; 340 | scaling_factor = a/total*1.f/(float)Math.sqrt(1 - a*a)+b/total*1.f/(float)Math.sqrt(1 - b*b); 341 | System.out.println(scaling_factor); 342 | return scaling_factor * scaling_factor;*/ 343 | 344 | 345 | /* 346 | for (int i = 0; i < 21; i++) { 347 | float t = (i / 20.f); 348 | float tPow = (float)Math.pow(t, 20); 349 | float tm1Pow = (float)Math.pow(1-t, 20); 350 | //sample the function at t and add f(t)*0.05; 351 | Vector3f N = new Vector3f( 352 | tPow*s.normal.x + tm1Pow*p.normal.x, 353 | tPow*s.normal.y + tm1Pow*p.normal.y, 354 | tPow*s.normal.z + tm1Pow*p.normal.z 355 | ); 356 | N.normalize(); 357 | float ND = N.dot(D); 358 | float f = 1.f / (float)Math.sqrt(1 - ND*ND); 359 | scaling_factor += (f*0.0476190476f); 360 | 361 | 362 | } 363 | System.out.println(scaling_factor); 364 | return scaling_factor * scaling_factor; 365 | */ 366 | 367 | } 368 | 369 | private static final Grid grid = new Grid(); 370 | 371 | private static class ASPPoint { 372 | public final Point3f point; 373 | public final Vector3f normal; 374 | public final int v1, v2, v3; 375 | public final float l1, l2; 376 | public ASPPoint(Vector3f n, int v1, int v2, int v3, float l1, float l2){ 377 | normal = n; this.v1 = v1; this.v2 = v2; this.v3 = v3; this.l1 = l1; this.l2 = l2; 378 | Point3f p1 = curr_verts.get(v1); 379 | Point3f p2 = curr_verts.get(v2); 380 | Point3f p3 = curr_verts.get(v3); 381 | float l3 = 1.f-l1-l2; 382 | point = new Point3f(l1*p1.x+l2*p2.x+l3*p3.x, l1*p1.y+l2*p2.y+l3*p3.y, l1*p1.z+l2*p2.z+l3*p3.z); 383 | } 384 | /*public void write(Writer w) throws Exception {w.write(v1+" "+v2+" "+v3+" "+l1+" "+l2+"\n");} 385 | public void writeobj(Writer w) throws Exception { 386 | w.write("v "+point.x+" "+point.y+" "+point.z+"\n"); 387 | }*/ 388 | 389 | } 390 | 391 | private static int getIdx(int i){return (i<0)?-i-1:i;} 392 | 393 | private static boolean generateSamples(){ 394 | float cur_area = cum_area[cum_area.length-1]; 395 | samples = new ArrayList (); 396 | aspSamples = new ArrayList(); 397 | grid.clearGrid(); System.gc(); 398 | boolean foundPoint = true; 399 | int minTries = 1000; 400 | while(foundPoint){ 401 | foundPoint = false; 402 | if(samples.size() >= NSamples) break; 403 | 404 | int actualNumTries = Math.max(minTries, samples.size() * 15); 405 | for(int i=0; i 2 | #include "emd_cuda.h" 3 | 4 | 5 | __global__ void approxmatch(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,float * __restrict__ match,float * temp){ 6 | float * remainL=temp+blockIdx.x*(n+m)*2, * remainR=temp+blockIdx.x*(n+m)*2+n,*ratioL=temp+blockIdx.x*(n+m)*2+n+m,*ratioR=temp+blockIdx.x*(n+m)*2+n+m+n; 7 | float multiL,multiR; 8 | if (n>=m){ 9 | multiL=1; 10 | multiR=n/m; 11 | }else{ 12 | multiL=m/n; 13 | multiR=1; 14 | } 15 | const int Block=1024; 16 | __shared__ float buf[Block*7]; 17 | for (int i=blockIdx.x;i=-2;j--){ 26 | float level=-powf(4.0f,j); 27 | if (j==-2){ 28 | level=0; 29 | } 30 | for (int k0=0;k0>>(b,n,m,xyz1,xyz2,match,temp); 172 | cudaError_t err = cudaGetLastError(); 173 | if (err != cudaSuccess) { 174 | printf("error in matching: %s\n", cudaGetErrorString(err)); 175 | //THError("aborting"); 176 | return 0; 177 | } 178 | return 1; 179 | } 180 | __global__ void matchcost(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ out){ 181 | __shared__ float allsum[512]; 182 | const int Block=1024; 183 | __shared__ float buf[Block*6]; 184 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,out); 231 | 232 | cudaError_t err = cudaGetLastError(); 233 | if (err != cudaSuccess) { 234 | printf("error in emd updateOutput: %s\n", cudaGetErrorString(err)); 235 | //THError("aborting"); 236 | return 0; 237 | } 238 | return 1; 239 | 240 | 241 | } 242 | __global__ void matchcostgrad2(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ grad2){ 243 | __shared__ float sum_grad[256*6]; 244 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,grad1); 337 | matchcostgrad2<<>>(b,n,m,xyz1,xyz2,match,grad2); 338 | 339 | cudaError_t err = cudaGetLastError(); 340 | if (err != cudaSuccess) { 341 | printf("error in emd backward: %s\n", cudaGetErrorString(err)); 342 | //THError("aborting"); 343 | return 0; 344 | } 345 | return 1; 346 | 347 | } 348 | 349 | -------------------------------------------------------------------------------- /src/emd_cuda.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/src/emd_cuda.cu.o -------------------------------------------------------------------------------- /src/emd_cuda.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | //int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 6 | //int NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 7 | 8 | int matchcostgradLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * grad1,float * grad2); 9 | int matchcostLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * out); 10 | int approxmatchLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,float * match, float * temp); 11 | 12 | 13 | #ifdef __cplusplus 14 | } 15 | #endif 16 | -------------------------------------------------------------------------------- /src/make.sh: -------------------------------------------------------------------------------- 1 | set -e 2 | nvcc nnd_cuda.cu -o nnd_cuda.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | -------------------------------------------------------------------------------- /src/my_lib.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | 4 | void nnsearch(int b,int n,int m,const float * xyz1,const float * xyz2,float * dist,int * idx){ 5 | for (int i=0;isize[0]; 30 | int n = xyz1->size[1]; 31 | int m = xyz2->size[1]; 32 | 33 | //printf("%d %d %d\n", batchsize, n, m); 34 | 35 | float *xyz1_data = THFloatTensor_data(xyz1); 36 | float *xyz2_data = THFloatTensor_data(xyz2); 37 | float *dist1_data = THFloatTensor_data(dist1); 38 | float *dist2_data = THFloatTensor_data(dist2); 39 | int *idx1_data = THIntTensor_data(idx1); 40 | int *idx2_data = THIntTensor_data(idx2); 41 | 42 | nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); 43 | nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); 44 | 45 | return 1; 46 | } 47 | 48 | 49 | int nnd_backward(THFloatTensor *xyz1, THFloatTensor *xyz2, THFloatTensor *gradxyz1, THFloatTensor *gradxyz2, THFloatTensor *graddist1, THFloatTensor *graddist2, THIntTensor *idx1, THIntTensor *idx2) { 50 | 51 | int b = xyz1->size[0]; 52 | int n = xyz1->size[1]; 53 | int m = xyz2->size[1]; 54 | 55 | //printf("%d %d %d\n", batchsize, n, m); 56 | 57 | float *xyz1_data = THFloatTensor_data(xyz1); 58 | float *xyz2_data = THFloatTensor_data(xyz2); 59 | float *gradxyz1_data = THFloatTensor_data(gradxyz1); 60 | float *gradxyz2_data = THFloatTensor_data(gradxyz2); 61 | float *graddist1_data = THFloatTensor_data(graddist1); 62 | float *graddist2_data = THFloatTensor_data(graddist2); 63 | int *idx1_data = THIntTensor_data(idx1); 64 | int *idx2_data = THIntTensor_data(idx2); 65 | 66 | 67 | for (int i=0;i 2 | #include "nnd_cuda.h" 3 | 4 | 5 | 6 | extern THCState *state; 7 | 8 | 9 | int nnd_forward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *dist1, THCudaTensor *dist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2) { 10 | int success = 0; 11 | success = NmDistanceKernelLauncher(xyz1->size[0], 12 | xyz1->size[1], 13 | THCudaTensor_data(state, xyz1), 14 | xyz2->size[1], 15 | THCudaTensor_data(state, xyz2), 16 | THCudaTensor_data(state, dist1), 17 | THCudaIntTensor_data(state, idx1), 18 | THCudaTensor_data(state, dist2), 19 | THCudaIntTensor_data(state, idx2), 20 | THCState_getCurrentStream(state) 21 | ); 22 | //int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream) 23 | 24 | 25 | if (!success) { 26 | THError("aborting"); 27 | } 28 | return 1; 29 | } 30 | 31 | 32 | int nnd_backward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *gradxyz1, THCudaTensor *gradxyz2, THCudaTensor *graddist1, 33 | THCudaTensor *graddist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2) { 34 | 35 | int success = 0; 36 | success = NmDistanceGradKernelLauncher(xyz1->size[0], 37 | xyz1->size[1], 38 | THCudaTensor_data(state, xyz1), 39 | xyz2->size[1], 40 | THCudaTensor_data(state, xyz2), 41 | THCudaTensor_data(state, graddist1), 42 | THCudaIntTensor_data(state, idx1), 43 | THCudaTensor_data(state, graddist2), 44 | THCudaIntTensor_data(state, idx2), 45 | THCudaTensor_data(state, gradxyz1), 46 | THCudaTensor_data(state, gradxyz2), 47 | THCState_getCurrentStream(state) 48 | ); 49 | //int NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream) 50 | 51 | if (!success) { 52 | THError("aborting"); 53 | } 54 | 55 | return 1; 56 | } 57 | 58 | 59 | 60 | -------------------------------------------------------------------------------- /src/my_lib_cuda.h: -------------------------------------------------------------------------------- 1 | int nnd_forward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *dist1, THCudaTensor *dist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2); 2 | 3 | 4 | int nnd_backward_cuda(THCudaTensor *xyz1, THCudaTensor *xyz2, THCudaTensor *gradxyz1, THCudaTensor *gradxyz2, THCudaTensor *graddist1, THCudaTensor *graddist2, THCudaIntTensor *idx1, THCudaIntTensor *idx2); 5 | 6 | -------------------------------------------------------------------------------- /src/nnd_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include "nnd_cuda.h" 3 | 4 | 5 | 6 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 7 | const int batch=512; 8 | __shared__ float buf[batch*3]; 9 | for (int i=blockIdx.x;ibest){ 121 | result[(i*n+j)]=best; 122 | result_i[(i*n+j)]=best_i; 123 | } 124 | } 125 | __syncthreads(); 126 | } 127 | } 128 | } 129 | int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ 130 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 131 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 132 | 133 | cudaError_t err = cudaGetLastError(); 134 | if (err != cudaSuccess) { 135 | printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); 136 | //THError("aborting"); 137 | return 0; 138 | } 139 | return 1; 140 | 141 | 142 | } 143 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 144 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 167 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 168 | 169 | cudaError_t err = cudaGetLastError(); 170 | if (err != cudaSuccess) { 171 | printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); 172 | //THError("aborting"); 173 | return 0; 174 | } 175 | return 1; 176 | 177 | } 178 | 179 | -------------------------------------------------------------------------------- /src/nnd_cuda.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/src/nnd_cuda.cu.o -------------------------------------------------------------------------------- /src/nnd_cuda.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | int NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 6 | 7 | int NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 8 | 9 | #ifdef __cplusplus 10 | } 11 | #endif -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | 5 | from modules.nnd import NNDModule 6 | 7 | dist = NNDModule() 8 | 9 | p1 = torch.rand(10,1000,3) 10 | p2 = torch.rand(10,1500,3) 11 | points1 = Variable(p1,requires_grad = True) 12 | points2 = Variable(p2) 13 | dist1, dist2 = dist(points1, points2) 14 | print(dist1, dist2) 15 | loss = torch.sum(dist1) 16 | print(loss) 17 | loss.backward() 18 | #print(points1.grad, points2.grad) 19 | 20 | 21 | points1 = Variable(p1.cuda(), requires_grad = True) 22 | points2 = Variable(p2.cuda()) 23 | dist1, dist2 = dist(points1, points2) 24 | print(dist1, dist2) 25 | loss = torch.sum(dist1) 26 | print(loss) 27 | loss.backward() 28 | #print(points1.grad, points2.grad) 29 | -------------------------------------------------------------------------------- /testchamfer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import time 4 | import numpy as np 5 | from torch.autograd import Variable 6 | 7 | from modules.nnd import NNDModule 8 | from tools.Ops import batch_pairwise_dist 9 | 10 | dist = NNDModule() 11 | 12 | p1 = torch.rand(64, 3, 4096).cuda() 13 | p2 = torch.rand(64, 3, 4096).cuda() 14 | points1 = Variable(p1.cuda(), requires_grad=True) 15 | points2 = Variable(p2.cuda(), requires_grad=True) 16 | 17 | indices = np.arange(4096).astype(int) 18 | np.random.shuffle(indices) 19 | indices = torch.from_numpy(indices[:1024]).cuda() 20 | spoints1 = points1[:, :, :].contiguous() 21 | spoints2 = points2[:, :, indices].contiguous() 22 | 23 | #points1 = points1.transpose(1,2).contiguous() 24 | #points2 = points2.transpose(1,2).contiguous() 25 | 26 | #points1 = torch.transpose(points1, 1, 2) 27 | #points2 = torch.transpose(points2, 1, 2) 28 | 29 | #start = time.time() 30 | #pd = batch_pairwise_dist(spoints1.transpose(1,2).contiguous(), spoints2.transpose(1,2).contiguous()) 31 | #loss = pd.min(dim=1)[0].sum() + pd.min(dim=2)[0].sum() 32 | #loss.backward() 33 | #end = time.time() 34 | #print(points1.grad[0, :, :], points2.grad[0, :, :]) 35 | #print loss 36 | #print "Runtime: {}s".format(end-start) 37 | 38 | start = time.time() 39 | dist1, dist2 = dist(spoints1.transpose(1,2).contiguous(), spoints2.transpose(1,2).contiguous()) 40 | #print(dist1, dist2) 41 | loss = dist1.sum() + dist2.sum() 42 | #print(loss) 43 | loss.backward() 44 | end = time.time() 45 | print(points1.grad[0, :, :], points2.grad[0, :, :]) 46 | print loss 47 | print "Runtime: {}s".format(end-start) 48 | 49 | -------------------------------------------------------------------------------- /tools/.ImageToPCDataset.py.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/tools/.ImageToPCDataset.py.swp -------------------------------------------------------------------------------- /tools/.PointCloudDataset.py.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/tools/.PointCloudDataset.py.swp -------------------------------------------------------------------------------- /tools/DataVis.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import os 4 | 5 | try: 6 | import visdom 7 | vis = visdom.Visdom() 8 | except ImportError: 9 | print "Visdom not available. No loss plotted." 10 | vis = None 11 | 12 | 13 | class LossLogger(object): 14 | 15 | def __init__(self, name, caption=None): 16 | self.name = name 17 | self.caption = caption 18 | self.window = None 19 | self.history = [] 20 | 21 | 22 | def update(self, value, step=10): 23 | val = value.cpu().data.numpy() 24 | self.history.append(val) 25 | 26 | if len(self.history) % step == 0: 27 | xs = np.arange(len(self.history)-step, len(self.history)) 28 | ys = np.array(self.history[-step:]).reshape((step,)) 29 | 30 | if vis is not None: 31 | if self.window is None: 32 | self.window = vis.line(X=xs, Y=ys, 33 | opts=dict(title=self.name, caption=self.caption)) 34 | else: 35 | vis.line(X=xs, Y=ys, win=self.window, update='append') 36 | 37 | else: 38 | if not os.path.exists("log"): 39 | os.makedirs("log") 40 | np.save(os.path.join("log", self.name+".npy"), self.history) 41 | 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /tools/ImageToPCDataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import torch.utils.data 4 | import os 5 | 6 | from tools.PointCloudDataset import * 7 | 8 | 9 | class ImageToPC(torch.utils.data.Dataset): 10 | 11 | def __init__(self, root_dir): 12 | self.filepaths = glob.glob(root_dir+"/*.obj") 13 | self.root_dir = root_dir 14 | 15 | 16 | def __len__(self): 17 | return len(self.filepaths) 18 | 19 | 20 | def __getitem__(self, idx): 21 | return torch.FloatTensor(read_obj(self.filepaths[idx])[:, 0:3].transpose()) 22 | 23 | -------------------------------------------------------------------------------- /tools/Ops.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | import numpy as np 6 | 7 | def cov(data): 8 | dmean = torch.mean(data, dim=0) 9 | centered_data = data - dmean.expand_as(data) 10 | return torch.mm(centered_data.transpose(0,1), centered_data) 11 | 12 | 13 | def cosine_similarity(x1, x2, dim=1, eps=1e-8): 14 | w12 = torch.sum(x1 * x2, dim) 15 | w1 = torch.norm(x1, 2, dim) 16 | w2 = torch.norm(x2, 2, dim) 17 | return (w12 / (w1 * w2).clamp(min=eps)).squeeze() 18 | 19 | 20 | def batch_pairwise_dist(a,b): 21 | x,y = a,b 22 | bs, num_points, points_dim = x.size() 23 | xx = torch.bmm(x, x.transpose(2,1)) 24 | yy = torch.bmm(y, y.transpose(2,1)) 25 | zz = torch.bmm(x, y.transpose(2,1)) 26 | diag_ind = torch.arange(0, num_points).type(torch.cuda.LongTensor) 27 | rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx) 28 | ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy) 29 | P = (rx.transpose(2,1) + ry - 2*zz) 30 | return P 31 | 32 | 33 | class NNUpsample1d(nn.Module): 34 | 35 | def __init__(self, scale=4): 36 | self.scale = scale 37 | super(NNUpsample1d, self).__init__() 38 | 39 | def forward(self, x): 40 | out = F.upsample(x.unsqueeze(3), scale_factor=self.scale, 41 | mode='nearest')[:, :, :, 0] 42 | return out 43 | 44 | 45 | def nplerp(x0, xn, n): 46 | interps = [] 47 | interps.append(x0) 48 | for i in xrange(1, n): 49 | alpha = i * 1.0/(n-1) 50 | interps.append(x0 + alpha*(xn-x0)) 51 | interps = np.array(interps) 52 | return interps 53 | 54 | 55 | def rotmat_2d(theta): 56 | mat = torch.zeros(2,2) 57 | mat[0,0] = torch.cos(theta) 58 | mat[0,1] = -torch.sin(theta) 59 | mat[1,0] = torch.sin(theta) 60 | mat[1,1] = torch.cos(theta) 61 | 62 | return mat 63 | 64 | 65 | def gridcoord_2d(w, h): 66 | max_dim = max(w, h) 67 | xs = torch.linspace(-max_dim/w, -max_dim/w, steps=w) 68 | ys = torch.linspace(-max_dim/h, -max_dim/h, steps=h) 69 | 70 | xc = xs.repeat(h) 71 | yc = ys.repeat(w,1).t().contiguous().view(-1) 72 | 73 | out = torch.cat((xc.unsqueeze(1), yc.unsqueeze(1)), 1) 74 | return out 75 | 76 | 77 | def resample_img(img, gridcoords, method='nearest'): 78 | if method=="nearest": 79 | round_coords = torch.round(gridcoords) 80 | out = img[round_coords] 81 | out.reshape(*(img.size())) 82 | return out 83 | 84 | 85 | def transform_image(img, t): 86 | height = img.size()[0] 87 | width = img.size()[1] 88 | grid2d = gridcoord_2d(width, height) 89 | 90 | transf_grid = torch.mm(grid2d, t) 91 | resample_img(img, transf_grid) 92 | 93 | 94 | if __name__ == '__main__': 95 | g = gridcoord_2d(2, 2) 96 | from IPython import embed; embed(); exit(-1) 97 | -------------------------------------------------------------------------------- /tools/PointCloudDataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import torch.utils.data 4 | import os 5 | from skimage import io, transform 6 | 7 | def save_obj_file(pts, filepath): 8 | f = open(filepath, 'w') 9 | for v in xrange(pts.shape[0]): 10 | line = "v {} {} {}\n".format(pts[v, 0], pts[v, 1], pts[v, 2]) 11 | f.write(line) 12 | if pts[v, :].shape[0] > 3: 13 | length = np.sqrt(np.sum(np.power(pts[v, 3:6], 2))) 14 | pts[v, 3:6] /= length 15 | line = "vn {} {} {}\n".format(pts[v, 3], pts[v, 4], pts[v, 5]) 16 | f.write(line) 17 | f.close() 18 | 19 | 20 | def save_torch_pc(path, pc): 21 | results = pc.cpu().data.numpy() 22 | results = results.transpose(0, 2, 1)[0, :, :] 23 | save_obj_file(results, path) 24 | 25 | def save_objs(pts, path, start_idx=0): 26 | filename = "pc_{}.obj" 27 | npts = pts.shape[1] 28 | ndims = pts.shape[2] 29 | 30 | for i in xrange(pts.shape[0]): 31 | save_obj_file(pts[i, :, :].reshape(npts, ndims), 32 | os.path.join(path, filename.format(str(i+start_idx).zfill(4)))) 33 | 34 | 35 | def save_points(pts, path): 36 | filename = "pc_{}.npy" 37 | npts = pts.shape[1] 38 | ndims = pts.shape[3] 39 | 40 | for i in xrange(pts.shape[0]): 41 | np.save(path+"/"+filename.format(i), 42 | pts[i, :, :, :].reshape(npts, ndims)) 43 | 44 | 45 | def point_from_vline(vline): 46 | sp = vline.split(' ')[1:4] 47 | pt = [] 48 | for v in sp: 49 | if np.isnan(float(v)): 50 | raise ValueError 51 | pt.append(float(v)) 52 | return np.array(pt) 53 | 54 | 55 | def read_obj(opath): 56 | points = [] 57 | with open(opath) as f: 58 | for line in f: 59 | if line.startswith('vn'): 60 | points[-1] = np.append(points[-1], point_from_vline(line)) 61 | else: 62 | points.append(point_from_vline(line)) 63 | return np.array(points) 64 | 65 | 66 | def write_image_pc(path, pair): 67 | #mean = np.array([0.485, 0.456, 0.406]).astype(float) 68 | #std = np.array([0.229, 0.224, 0.225]).astype(float) 69 | 70 | #img = pair[0].numpy().transpose((1, 2, 0)) 71 | img = pair[0] 72 | #img = np.clip(img*std + mean, 0.0, 1.0) 73 | #io.imsave(path+'.png', (img*255).astype('uint8')) 74 | io.imsave(path+'.png', img) 75 | save_obj_file(pair[1].transpose(0,1).numpy(), path+'.obj') 76 | 77 | 78 | class ImageDataset(torch.utils.data.Dataset): 79 | 80 | def __init__(self, image_dir): 81 | self.img_paths = sorted(glob.glob(image_dir+"/*.png")) 82 | 83 | self.mean = np.array([0.485, 0.456, 0.406]).astype(float) 84 | self.std = np.array([0.229, 0.224, 0.225]).astype(float) 85 | 86 | def __len__(self): 87 | return len(self.img_paths) 88 | 89 | def __getitem__(self, idx): 90 | img = io.imread(self.img_paths[idx])[:, :, 0:3].astype('float32') 91 | img /= 255.0 92 | img = transform.resize(img, (224, 224)) 93 | img = (img - self.mean)/self.std 94 | img = img.transpose((2, 0, 1)) 95 | img = torch.from_numpy(img).float() 96 | 97 | name = self.img_paths[idx].split("/")[-1].split(".")[-2] 98 | 99 | return (img, io.imread(self.img_paths[idx]), name) 100 | 101 | 102 | class ImageToPointCloudDataset(torch.utils.data.Dataset): 103 | 104 | def __init__(self, image_dir, pc_dir, category="all", train_mode=True): 105 | self.train_mode = train_mode 106 | self.category = category 107 | 108 | self.classids= ['02691156', '02828884', '02933112', '02958343', '03001627', 109 | '03211117', '03636649', '03691459', '04090263', '04256520', '04379243', 110 | '04401088', '04530566'] 111 | 112 | self.img_paths = [] 113 | if self.category == "all": 114 | self.img_paths = sorted(glob.glob(image_dir+"/*/*/rendering/*.png"), 115 | key=lambda p: p.split('/')[-3]) 116 | else: 117 | self.img_paths = sorted(glob.glob( 118 | image_dir+"/{}/*/rendering/*.png".format(self.category)), 119 | key=lambda p: p.split('/')[-3]) 120 | 121 | self.mean = np.array([0.485, 0.456, 0.406]).astype(float) 122 | self.std = np.array([0.229, 0.224, 0.225]).astype(float) 123 | 124 | ntrain_id = int(len(self.img_paths) * 0.8) 125 | if train_mode: 126 | self.img_paths = self.img_paths[0:ntrain_id] 127 | else: 128 | self.img_paths = self.img_paths[ntrain_id:] 129 | self.pc_paths = [] 130 | 131 | for img in self.img_paths: 132 | model_signature = img.split('/')[-3] 133 | pc_path = "{}/{}_0.npy".format(pc_dir, model_signature) 134 | self.pc_paths.append(pc_path) 135 | 136 | def __len__(self): 137 | return len(self.img_paths) 138 | 139 | def __getitem__(self, idx): 140 | img = io.imread(self.img_paths[idx])[:, :, 0:3].astype('float32') 141 | img /= 255.0 142 | img = transform.resize(img, (224, 224)) 143 | img = (img - self.mean)/self.std 144 | img = img.transpose((2, 0, 1)) 145 | img = torch.from_numpy(img).float() 146 | 147 | pc = torch.FloatTensor(np.load(self.pc_paths[idx])[:, 0:3].transpose()) 148 | 149 | #if self.train_mode: 150 | # out = (img, pc) 151 | #else: 152 | # out = (img, pc, self.classids.index(self.img_paths[idx].split('/')[-4])) 153 | out = (img, pc, self.classids.index(self.img_paths[idx].split('/')[-4]), 154 | io.imread(self.img_paths[idx])) 155 | 156 | return out 157 | 158 | 159 | class PointCloudDataset(torch.utils.data.Dataset): 160 | 161 | def __init__(self, root_dir): 162 | self.filepaths = glob.glob(root_dir+"/*.obj") 163 | self.root_dir = root_dir 164 | 165 | def __len__(self): 166 | return len(self.filepaths) 167 | 168 | 169 | def __getitem__(self, idx): 170 | return torch.FloatTensor(read_obj(self.filepaths[idx])[:, 0:3].transpose()) 171 | 172 | 173 | class PointModelNetDataset(torch.utils.data.Dataset): 174 | 175 | def __init__(self, root_dir, scale_aug=True, rot_aug=False, test_mode=False): 176 | self.filepaths = sorted(glob.glob(root_dir+"/*.obj_0.kdt.npy")) 177 | #self.filepaths = sorted(glob.glob(root_dir+"/*.rpt.npy")) 178 | self.root_dir = root_dir 179 | self.scale_aug = scale_aug 180 | self.rot_aug = rot_aug 181 | self.test_mode = test_mode 182 | 183 | self.class_files = {} 184 | 185 | self.classnames=[ 186 | 'airplane', 187 | 'bathtub', 188 | 'bed', 189 | 'bench', 190 | 'bookshelf', 191 | 'bottle', 192 | 'bowl', 193 | 'car', 194 | 'chair', 195 | 'cone', 196 | 'cup', 197 | 'curtain', 198 | 'desk', 199 | 'door', 200 | 'dresser', 201 | 'flower_pot', 202 | 'glass_box', 203 | 'guitar', 204 | 'keyboard', 205 | 'lamp', 206 | 'laptop', 207 | 'mantel', 208 | 'monitor', 209 | 'night_stand', 210 | 'person', 211 | 'piano', 212 | 'plant', 213 | 'radio', 214 | 'range_hood', 215 | 'sink', 216 | 'sofa', 217 | 'stairs', 218 | 'stool', 219 | 'table', 220 | 'tent', 221 | 'toilet', 222 | 'tv_stand', 223 | 'vase', 224 | 'wardrobe', 225 | 'xbox'] 226 | 227 | self.class_splits = [] 228 | 229 | current_class = '###' 230 | for f in self.filepaths: 231 | if current_class in f: 232 | self.class_splits[-1] += 1 233 | else: 234 | for c in self.classnames: 235 | if c in f and c != current_class: 236 | self.class_splits.append(1) 237 | current_class = c 238 | 239 | self.class_splits = np.cumsum(np.array(self.class_splits)) 240 | 241 | 242 | def __len__(self): 243 | return len(self.filepaths) 244 | 245 | 246 | def __getitem__(self, idx): 247 | path = self.filepaths[idx] 248 | class_idx = None 249 | class_name = path.split('/')[5] 250 | i = self.classnames.index(class_name) 251 | 252 | pc = np.load(self.filepaths[idx]) 253 | if self.scale_aug: 254 | scale_factor = np.random.uniform(0.75, 1.5, 3) 255 | pc[:, 0:3] *= scale_factor 256 | 257 | if self.rot_aug: 258 | rot_matrix = rotation_matrix(np.array([0, 0, 1]), 259 | np.random.choice(np.arange(0, 2*np.pi, np.pi/4)) + 260 | (np.random.random_sample()-0.5) + 1e-2) 261 | pc[:, 0:3] = pc[:, 0:3].dot(rot_matrix) 262 | 263 | dims = pc[:, 0:3].max(axis=0) - pc[:, 0:3].min(axis=0) 264 | pc[:, 0:3] /= dims.max() 265 | 266 | return (i, torch.FloatTensor(pc.transpose())) 267 | 268 | 269 | -------------------------------------------------------------------------------- /tools/Trainer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Variable 5 | 6 | import DataVis 7 | import numpy as np 8 | import os 9 | 10 | from tensorboardX import SummaryWriter 11 | 12 | from tools import Ops 13 | from tools.PointCloudDataset import save_torch_pc 14 | from tools.PointCloudDataset import write_image_pc 15 | 16 | class AutoEncoderTrainer(object): 17 | 18 | def __init__(self, model, loader, optimizer, loss_fn): 19 | 20 | self.optimizer = optimizer 21 | self.model = model 22 | self.loader = loader 23 | self.loss_fn = loss_fn 24 | self.logger = DataVis.LossLogger(model.name) 25 | 26 | self.model.cuda() 27 | 28 | 29 | def train(self, n_epochs, save_step=5): 30 | 31 | for epoch in range(n_epochs): 32 | out_data = None 33 | 34 | for i, data in enumerate(self.loader, 0): 35 | in_data = Variable(data.cuda()) 36 | 37 | 38 | self.optimizer.zero_grad() 39 | 40 | out_data = self.model(in_data) 41 | loss = self.loss_fn(out_data, in_data) 42 | 43 | self.logger.update(loss) 44 | 45 | loss.backward() 46 | self.optimizer.step() 47 | 48 | if epoch % save_step == 0: 49 | self.model.save("checkpoint", epoch) 50 | 51 | results_dir = os.path.join("results", self.model.name, 52 | "epoch_{}".format(str(epoch).zfill(4))) 53 | if not os.path.exists(results_dir): 54 | os.makedirs(results_dir) 55 | 56 | try: 57 | self.model.save_results(results_dir, out_data) 58 | except NotImplementedError: 59 | print "Intermediate results not saved." 60 | 61 | 62 | class ImageToPCTrainer(object): 63 | 64 | def __init__(self, model, train_loader, val_loader, optimizer, loss_fn, 65 | log_dir="log"): 66 | 67 | self.optimizer = optimizer 68 | self.model = model 69 | self.train_loader = train_loader 70 | self.val_loader = val_loader 71 | self.loss_fn = loss_fn 72 | self.log_dir = log_dir 73 | 74 | self.writer = SummaryWriter(self.log_dir) 75 | 76 | self.model.cuda() 77 | 78 | 79 | def train(self, n_epochs, save_step=1): 80 | 81 | best_acc = 0 82 | it_num = 0 83 | 84 | for epoch in range(n_epochs): 85 | out_data = None 86 | in_data = None 87 | 88 | for i, data in enumerate(self.train_loader, 0): 89 | if data[1].size()[0] != self.model.batch_size: 90 | continue 91 | 92 | in_data = Variable(data[0].cuda()) 93 | target = Variable(data[1]).cuda() 94 | 95 | self.optimizer.zero_grad() 96 | 97 | out_data = self.model(in_data) 98 | loss = self.loss_fn(out_data, target) 99 | print loss 100 | 101 | self.writer.add_scalar('train_loss', loss, it_num) 102 | 103 | loss.backward() 104 | self.optimizer.step() 105 | it_num += 1 106 | 107 | if epoch % save_step == 0: 108 | self.model.save("checkpoint", epoch) 109 | 110 | results_dir = os.path.join("results", self.model.name, 111 | "epoch_{}".format(str(epoch).zfill(4))) 112 | if not os.path.exists(results_dir): 113 | os.makedirs(results_dir) 114 | try: 115 | self.model.save_results(results_dir, out_data) 116 | except NotImplementedError: 117 | print "Results not saved." 118 | 119 | if epoch > 0 and epoch % 5 == 0: 120 | for param_group in self.optimizer.param_groups: 121 | param_group['lr'] /= 2 122 | 123 | self.writer.export_scalars_to_json(self.log_dir+"/all_scalars.json") 124 | self.writer.close() 125 | 126 | 127 | def run(self): 128 | 129 | self.model.eval() 130 | 131 | for i, data in enumerate(self.val_loader, 0): 132 | if data[1].size()[0] != self.model.batch_size: 133 | continue 134 | 135 | in_data = Variable(data[0].cuda()) 136 | 137 | out_data = self.model(in_data) 138 | 139 | results_dir = os.path.join("run", self.model.name) 140 | print data[2] 141 | if not os.path.exists(results_dir): 142 | os.makedirs(results_dir) 143 | write_image_pc(os.path.join(results_dir, data[2][0]+"_mrt"), 144 | (data[1][0, :, :, :], out_data[0, :, :].data.cpu())) 145 | print "Test PC saved." 146 | 147 | 148 | def evaluate(self): 149 | n_batches = 0 150 | all_correct_points = 0 151 | miou = 0 152 | 153 | total_error_class = np.zeros((13, 2)) 154 | total_count_class = np.zeros(13) 155 | 156 | in_data = None 157 | out_data = None 158 | target = None 159 | 160 | n_iter = 0.0 161 | total_d = 0.0 162 | 163 | self.model.eval() 164 | 165 | for i, data in enumerate(self.val_loader, 0): 166 | if data[1].size()[0] != self.model.batch_size: 167 | continue 168 | 169 | in_data = Variable(data[0].cuda()) 170 | target = Variable(data[1]).cuda() 171 | class_id = data[2][0] 172 | 173 | out_data = self.model(in_data) 174 | 175 | pd = Ops.batch_pairwise_dist(out_data.transpose(1,2), 176 | target.transpose(1,2)) 177 | pd = torch.sqrt(pd) 178 | 179 | total_error_class[class_id, 0] += torch.min(pd, dim=2)[0].data.cpu().numpy().mean() 180 | total_error_class[class_id, 1] += torch.min(pd, dim=1)[0].data.cpu().numpy().mean() 181 | total_count_class[class_id] += 1.0 182 | 183 | scalar_group = {} 184 | 185 | #Iterate over classes 186 | for c in xrange(13): 187 | if total_count_class[c] > 0.0: 188 | scalar_group['class{}_error_pred'.format(c)] = total_error_class[c, 0]/total_count_class[c] 189 | scalar_group['class{}_error_gt'.format(c)] = total_error_class[c, 1]/total_count_class[c] 190 | 191 | np.save('total_error_class.npy', total_error_class) 192 | np.save('total_count_class.npy', total_count_class) 193 | self.writer.add_scalars('class_errors', scalar_group, i) 194 | 195 | n_iter += self.model.batch_size 196 | 197 | #Save some point clouds for visualization 198 | if i < 50: 199 | results_dir = os.path.join("eval", self.model.name) 200 | if not os.path.exists(results_dir): 201 | os.makedirs(results_dir) 202 | write_image_pc(os.path.join(results_dir, "out_{}".format(str(2*i).zfill(4))), 203 | (data[3][0, :, :, :], out_data[0, :, :].data.cpu())) 204 | save_torch_pc(os.path.join(results_dir, "out_{}.obj".format(str(2*i+1).zfill(4))), target) 205 | print "Test PC saved." 206 | 207 | #Saves results 208 | np.save('total_error_class.npy', total_error_class) 209 | np.save('total_count_class.npy', total_count_class) 210 | print total_d/n_iter 211 | 212 | 213 | class ModelNetTrainer(object): 214 | 215 | def __init__(self, model, train_loader, val_loader, optimizer, loss_fn): 216 | 217 | self.optimizer = optimizer 218 | self.model = model 219 | self.train_loader = train_loader 220 | self.val_loader = val_loader 221 | self.loss_fn = loss_fn 222 | self.logger = DataVis.LossLogger(model.name) 223 | self.val_logger = DataVis.LossLogger("Validation-"+model.name) 224 | self.train_logger = DataVis.LossLogger("Training acc. "+model.name) 225 | 226 | self.model.cuda() 227 | 228 | 229 | def train(self, n_epochs, save_step=5): 230 | 231 | best_acc = 0 232 | 233 | for epoch in range(n_epochs): 234 | out_data = None 235 | in_data = None 236 | 237 | for i, data in enumerate(self.train_loader, 0): 238 | if data[1].size()[0] != self.model.batch_size: 239 | continue 240 | 241 | in_data = Variable(data[1].cuda()) 242 | target = Variable(data[0]).cuda().long() 243 | 244 | self.optimizer.zero_grad() 245 | 246 | out_data = self.model(in_data) 247 | loss = self.loss_fn(out_data, target) 248 | 249 | self.logger.update(loss) 250 | 251 | pred = torch.max(out_data, 1)[1] 252 | results = pred == target 253 | correct_points = torch.sum(results.long()) 254 | 255 | self.train_logger.update(correct_points.float()/self.model.batch_size) 256 | 257 | loss.backward() 258 | self.optimizer.step() 259 | 260 | 261 | epoch_acc = self.update_validation_accuracy(epoch) 262 | if epoch_acc > best_acc: 263 | best_acc = epoch_acc 264 | self.model.save("checkpoint", epoch) 265 | #self.update_validation_accuracy(epoch) 266 | 267 | if epoch > 0 and epoch % 5 == 0: 268 | for param_group in self.optimizer.param_groups: 269 | param_group['lr'] /= 2 270 | 271 | def update_validation_accuracy(self, epoch): 272 | n_batches = 0 273 | all_correct_points = 0 274 | miou = 0 275 | 276 | in_data = None 277 | out_data = None 278 | target = None 279 | 280 | results_dir = os.path.join("bad_results", self.model.name, 281 | "epoch_{}".format(str(epoch).zfill(4))) 282 | if not os.path.exists(results_dir): 283 | os.makedirs(results_dir) 284 | 285 | wrong_results = [] 286 | wrong_label = [] 287 | wrong_class = np.zeros(40) 288 | samples_class = np.zeros(40) 289 | 290 | self.model.eval() 291 | 292 | for i, data in enumerate(self.val_loader, 0): 293 | if data[1].size()[0] != self.val_loader.batch_size: 294 | continue 295 | 296 | in_data = Variable(data[1], volatile=True).cuda() 297 | target = Variable(data[0], volatile=True).cuda().long() 298 | 299 | out_data = self.model(in_data) 300 | 301 | pred = torch.max(out_data, 1)[1] 302 | 303 | results = pred == target 304 | for i in xrange(results.size()[0]): 305 | if not bool(results[i].cpu().data.numpy()): 306 | wrong_results.append(in_data[i, :, :].cpu().data.numpy()) 307 | wrong_label.append(pred.cpu().data.numpy().astype('int')[i]) 308 | wrong_class[target.cpu().data.numpy().astype('int')[i]] += 1 309 | samples_class[target.cpu().data.numpy().astype('int')[i]] += 1 310 | correct_points = torch.sum(results.long()) 311 | 312 | all_correct_points += correct_points 313 | n_batches += 1 314 | 315 | #self.model.save_bad_results(wrong_results, wrong_label, results_dir) 316 | acc = all_correct_points.float() / (n_batches*self.val_loader.batch_size) 317 | self.val_logger.update(acc, step=1) 318 | 319 | print np.mean((samples_class-wrong_class)/samples_class) 320 | 321 | self.model.train() 322 | 323 | return acc.cpu().data.numpy()[0] 324 | 325 | 326 | class VAETrainer(object): 327 | 328 | def __init__(self, model, loader, optimizer, loss_fn): 329 | 330 | self.optimizer = optimizer 331 | self.model = model 332 | self.loader = loader 333 | self.loss_fn = loss_fn 334 | self.logger = DataVis.LossLogger(model.name) 335 | self.reg_logger = DataVis.LossLogger("encreg-{}".format(model.name)) 336 | 337 | self.model.cuda() 338 | 339 | 340 | def train(self, n_epochs, save_step=50): 341 | 342 | for epoch in range(n_epochs): 343 | out_data = None 344 | 345 | for i, data in enumerate(self.loader, 0): 346 | if data.size()[0] != self.model.batch_size: 347 | continue 348 | 349 | in_data = Variable(data.cuda()) 350 | 351 | self.optimizer.zero_grad() 352 | 353 | out_data = self.model(in_data) 354 | loss = self.loss_fn(out_data, in_data) 355 | reg = self.model.encoding_regularizer(in_data) 356 | print "Total loss: {} | Rec Loss: {} | KLD: {}".format( 357 | loss.cpu().data.numpy()[0] + reg.cpu().data.numpy()[0], 358 | loss.cpu().data.numpy()[0], 359 | reg.cpu().data.numpy()[0]) 360 | 361 | loss += reg 362 | 363 | self.logger.update(loss) 364 | self.reg_logger.update(reg) 365 | 366 | loss.backward() 367 | self.optimizer.step() 368 | 369 | if epoch % save_step == 0: 370 | self.model.save("checkpoint", epoch) 371 | 372 | results_dir = os.path.join("results", self.model.name, 373 | "epoch_{}".format(str(epoch).zfill(4))) 374 | if not os.path.exists(results_dir): 375 | os.makedirs(results_dir) 376 | 377 | samples_dir = os.path.join("samples", self.model.name, 378 | "epoch_{}".format(str(epoch).zfill(4))) 379 | if not os.path.exists(samples_dir): 380 | os.makedirs(samples_dir) 381 | 382 | for param_group in self.optimizer.param_groups: 383 | param_group['lr'] /= 2 384 | 385 | try: 386 | self.model.save_results(results_dir, out_data) 387 | self.model.save_results(samples_dir, self.model.sample()) 388 | except NotImplementedError: 389 | print "Intermediate results not saved." 390 | 391 | 392 | class AdversarialTrainer(object): 393 | 394 | def __init__(self, 395 | d_model, 396 | g_model, 397 | d_opt, 398 | g_opt, 399 | d_loss_fn, 400 | g_loss_fn, 401 | loader, 402 | encoder, 403 | decoder, 404 | batch_size=64): 405 | 406 | self.discriminator = d_model 407 | self.generator = g_model 408 | self.d_opt = d_opt 409 | self.g_opt = g_opt 410 | self.d_loss_fn = d_loss_fn 411 | self.g_loss_fn = g_loss_fn 412 | self.loader = loader 413 | 414 | self.d_logger = DataVis.LossLogger(self.discriminator.name) 415 | self.dacc_real_logger = DataVis.LossLogger("ACCReal"+self.discriminator.name) 416 | self.dacc_fake_logger = DataVis.LossLogger("ACCFake_"+self.discriminator.name) 417 | self.g_logger = DataVis.LossLogger(self.generator.name) 418 | 419 | self.encoder = encoder 420 | self.decoder = decoder 421 | 422 | self.batch_size = batch_size 423 | 424 | self.noise = torch.FloatTensor(self.batch_size, self.generator.enc_size) 425 | 426 | self.discriminator.cuda() 427 | self.generator.cuda() 428 | self.encoder.cuda() 429 | self.decoder.cuda() 430 | 431 | self.encoder.eval() 432 | self.decoder.eval() 433 | 434 | 435 | def reset_grad(self): 436 | self.discriminator.zero_grad() 437 | self.generator.zero_grad() 438 | 439 | 440 | def train(self, n_epochs, save_step=5): 441 | 442 | for epoch in range(n_epochs): 443 | fake_data = None 444 | real_data = None 445 | 446 | for i, data in enumerate(self.loader, 0): 447 | if data.size()[0] != self.batch_size: 448 | continue 449 | 450 | self.noise.normal_() 451 | real_data = self.encoder(Variable(data.cuda())) 452 | fake_data = self.generator(Variable(self.noise.cuda())) 453 | 454 | d_out_real = self.discriminator(real_data) 455 | d_out_fake = self.discriminator(fake_data) 456 | 457 | #Discriminator update 458 | self.d_opt.zero_grad() 459 | dloss_real = self.d_loss_fn((d_out_real, 460 | Variable(torch.ones(self.batch_size).cuda()))) 461 | dloss_fake = self.d_loss_fn((d_out_fake, 462 | Variable(torch.zeros(self.batch_size).cuda()))) 463 | 464 | dacc_real = torch.mean(F.sigmoid(d_out_real[0])) 465 | dacc_fake = 1 - torch.mean(F.sigmoid(d_out_fake[0])) 466 | self.dacc_real_logger.update(dacc_real) 467 | self.dacc_fake_logger.update(dacc_fake) 468 | dacc = 0.5*(dacc_real + dacc_fake) 469 | 470 | dloss = dloss_real + dloss_fake 471 | self.d_logger.update(dloss) 472 | 473 | self.reset_grad() 474 | dloss.backward(retain_variables=True) 475 | if dacc < 0.8: 476 | self.d_opt.step() 477 | 478 | #Generator update 479 | self.g_opt.zero_grad() 480 | 481 | gloss = self.g_loss_fn((d_out_real, d_out_fake)) 482 | #gloss = self.d_loss_fn((d_out_fake, 483 | # Variable(torch.ones(self.batch_size).cuda()))) 484 | self.g_logger.update(gloss) 485 | 486 | self.reset_grad() 487 | gloss.backward() 488 | self.g_opt.step() 489 | 490 | if epoch % save_step == 0: 491 | self.generator.save("checkpoint", epoch) 492 | self.discriminator.save("checkpoint", epoch) 493 | 494 | results_dir = os.path.join("results", self.generator.name, 495 | "epoch_{}".format(str(epoch).zfill(4))) 496 | if not os.path.exists(results_dir): 497 | os.makedirs(results_dir) 498 | 499 | try: 500 | self.generator.save_results(results_dir, self.decoder(fake_data)) 501 | except NotImplementedError: 502 | print "Intermediate results not saved." 503 | 504 | 505 | -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/matheusgadelha/MRTNet/3b540189c56aec201c3423469cb497730f7797e2/tools/__init__.py -------------------------------------------------------------------------------- /train_img2pc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.optim as optim 4 | import torch.nn as nn 5 | 6 | import argparse 7 | import os 8 | 9 | from tools.Trainer import ImageToPCTrainer 10 | from tools.PointCloudDataset import ImageToPointCloudDataset 11 | from models.AutoEncoder import PointCloudVAE 12 | from models.AutoEncoder import ChamferLoss 13 | from models.AutoEncoder import ChamferWithNormalLoss 14 | from models.AutoEncoder import L2WithNormalLoss 15 | from models.ImageToShape import MultiResImageToShape 16 | 17 | parser = argparse.ArgumentParser(description='MultiResolution image to shape model.') 18 | parser.add_argument("-n", "--name", type=str, help="Name of the experiment.", default="MRI2PC") 19 | parser.add_argument("-a", "--arch", type=str, help="Encoder architecture.", default="vgg") 20 | parser.add_argument("-bs", "--batchSize", type=int, help="Batch size", default=64) 21 | parser.add_argument("-pt", "--pretrained", type=str, help="Use pretrained net", default=True) 22 | parser.add_argument("-c", "--category", type=str, help="Category code (all is possible)", default="all") 23 | parser.add_argument("--train", dest='train', action='store_true') 24 | parser.set_defaults(train=False) 25 | 26 | 27 | #CHANGE THESE!!! 28 | image_datapath = "/media/mgadelha/hd2/ShapenetRenderings" 29 | pc_datapath = "/media/mgadelha/hd2/shapenet_4k" 30 | 31 | if __name__ == '__main__': 32 | args = parser.parse_args() 33 | 34 | ptrain = None 35 | if args.pretrained == "False": 36 | ptrain = False 37 | elif args.pretrained == "True": 38 | ptrain = True 39 | 40 | full_name = "{}_{}_{}_{}".format(args.name, args.category, args.arch, ptrain) 41 | print full_name 42 | 43 | mri2pc = MultiResImageToShape(size=4096, dim=3, batch_size=args.batchSize, 44 | name=full_name, pretrained=ptrain, arch=args.arch) 45 | #mri2pc.load('checkpoint') 46 | optimizer = optim.Adam(mri2pc.parameters(), lr=0.001) 47 | 48 | train_dataset = ImageToPointCloudDataset(image_datapath, pc_datapath, category=args.category, train_mode=True) 49 | test_dataset = ImageToPointCloudDataset(image_datapath, pc_datapath, category=args.category, train_mode=False) 50 | train_loader = torch.utils.data.DataLoader(train_dataset, 51 | batch_size=args.batchSize, shuffle=True, num_workers=2) 52 | test_loader = torch.utils.data.DataLoader(test_dataset, 53 | batch_size=args.batchSize, 54 | shuffle=True, num_workers=2) 55 | 56 | log_dir = os.path.join("log", full_name) 57 | if not os.path.exists(log_dir): 58 | os.makedirs(log_dir) 59 | 60 | trainer = ImageToPCTrainer(mri2pc, train_loader, test_loader, 61 | optimizer, ChamferLoss(cuda_opt=True), log_dir=log_dir) 62 | trainer.train(2000) 63 | 64 | --------------------------------------------------------------------------------