├── .gitignore ├── LICENSE ├── README.md ├── config ├── all.csv └── chair_T26AK2FES_model_699 ├── demo └── 03001627 │ └── dac4af24e2facd7d3000ca4b04fcd6ac │ └── models │ ├── model_normalized.json │ ├── model_normalized.mtl │ ├── model_normalized.obj │ ├── model_normalized.solid.binvox │ └── model_normalized.surface.binvox ├── img ├── 00.png ├── 01.png └── 02.png ├── learnable_primitives ├── __init__.py ├── common │ ├── __init__.py │ ├── batch_provider.py │ ├── dataset.py │ └── model_factory.py ├── equal_distance_sampler_sq.py ├── fast_sampler │ ├── __init__.py │ ├── _sampler.c │ ├── _sampler.pyx │ ├── sampling.cpp │ └── sampling.hpp ├── loss_functions.py ├── mesh.py ├── models.py ├── pointcloud.py ├── primitives.py ├── regularizers.py ├── utils │ ├── __init__.py │ ├── pcl_voxelization.py │ └── progbar.py ├── volumetric_iou.py └── voxelizers.py ├── misc ├── __init__.py ├── chamfer_vs_inside_outside_local_minima.py ├── create_simple_shapes_dataset.py ├── create_spheres_dataset.py ├── shapes.py └── visualize_distance_function.py ├── requirements.txt ├── scripts ├── arguments.py ├── benchmark_sampler.py ├── compute_chamfer_loss.py ├── forward_pass.py ├── output_logger.py ├── train_network.py ├── utils.py ├── visualization_utils.py └── visualize_sq.py └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.bbl 2 | *.blg 3 | *.log 4 | *.out 5 | *.aux 6 | *.toc 7 | *.brf 8 | *.nav 9 | *.snm 10 | *.bz2 11 | *pyc 12 | *.swp 13 | top.pdf 14 | supp.pdf 15 | 16 | *.so 17 | build/ 18 | *.egg-info 19 | 20 | scripts/learnable_primitives 21 | misc/learnable_primitives 22 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Max Planck Institute for Intelligent Systems, http://www.is.mpg.de 4 | Written by Despoina Paschalidou 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy of 7 | this software and associated documentation files (the "Software"), to deal in 8 | the Software without restriction, including without limitation the rights to 9 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 10 | of the Software, and to permit persons to whom the Software is furnished to do 11 | so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Superquadrics Revisited: Learning 3D Shape Parsing beyond Cuboids 2 | 3 | This repository contains the code that accompanies our CVPR 2019 paper 4 | [Superquadrics Revisited: Learning 3D Shape Parsing beyond Cuboids](https://arxiv.org/pdf/1904.09970.pdf) 5 | 6 | ![Teaser Image](https://avg.is.tuebingen.mpg.de/uploads/publication/image/22555/superquadrics_parsing.png) 7 | 8 | You can find detailed instructions for both training your own models and using 9 | pretrained models in the examples below. 10 | 11 | Dependencies & Installation 12 | ---------------------------- 13 | 14 | Our library has the following dependencies: 15 | 16 | - Python 2.7 17 | - [PyTorch](https://pytorch.org/get-started/locally/) 18 | - torchvision 19 | - numpy 20 | - [scikit-learn](https://scikit-learn.org/stable/install.html) 21 | - [progress](https://pypi.org/project/progress/) 22 | - [pyquaternion](http://kieranwynn.github.io/pyquaternion/) 23 | - [trimesh](https://github.com/mikedh/trimesh) 24 | - [backports.functools_lru_cache](https://pypi.org/project/backports.functools_lru_cache/) 25 | - [Cython](https://cython.readthedocs.io/en/latest/src/quickstart/install.html) 26 | - [sympy](https://pypi.org/project/sympy/) 27 | - seaborn 28 | - matplotlib 29 | - [PIL](https://pillow.readthedocs.io/en/stable/installation.html#basic-installation) 30 | - [mayavi](https://docs.enthought.com/mayavi/mayavi/installation.html#installing-with-pip) 31 | 32 | They should be automatically installed by running 33 | ``` 34 | pip install --user -e . 35 | ``` 36 | 37 | In case you this doesn't work automatically try this instead, 38 | ``` 39 | pip install -r requirements.txt 40 | pip install --user -e . 41 | ``` 42 | 43 | Please note that you might need to install `python-qt4` in order to be able to 44 | use mayavi. You can do that by simply typing 45 | ``` 46 | sudo apt install python-qt4 47 | ``` 48 | Now you are ready to start playing with the code! 49 | 50 | Evaluation 51 | ---------- 52 | 53 | For evaluating a previously trained model, we provide the `forward_pass.py` 54 | script. This script performs a forward pass and predicts the parameters of the 55 | superquadric surfaces used to represent the 3D object. With this script you can 56 | visualize the predicted superquadrics using `mayavi` as well as save them as a 57 | mesh file. 58 | 59 | You can run it using 60 | ``` 61 | $ ./forward_pass.py ../demo/03001627/ /tmp/ --model_tag "dac4af24e2facd7d3000ca4b04fcd6ac" --n_primitives 18 --weight_file ../config/chair_T26AK2FES_model_699 --train_with_bernoulli --use_deformations --use_sq --dataset_type shapenet_v2 62 | ``` 63 | 64 | The script requires two mandatory arguments, the path to the directory that 65 | contains the dataset, in this case it is `../demo/03001627` and the path to a 66 | directory that will be used for saving the generated files, here `/tmp`. You 67 | should also provide (even if it is not mandatory) the path to the previously 68 | trained model via the `--weight_file` argument, as well as the tag of the model 69 | you want to reconstruct (`--model_tag`) and the type of the dataset you are 70 | using (`--dataset_type`). Note that you should provide the same arguments that 71 | you used when training the model, regarding the configuration of the geometric 72 | primitives (e.g number of primitives, whether or not to use superquadrics 73 | etc.). This script automatically visualizes the predicted superquadrics using 74 | `mayavi`. To save these predictions as a mesh file, simply add the 75 | `save_prediction_as_mesh` argument. 76 | 77 | Running the above command, will result in something like the following: 78 | ``` 79 | $ ./forward_pass.py ~/data/03001627/ /tmp/ --model_tag "dac4af24e2facd7d3000ca4b04fcd6ac" --n_primitives 18 --weight_file ../config/chair_T26AK2FES_model_699 --train_with_bernoulli --use_deformations --use_sq --dataset_type shapenet_v2 80 | No handlers could be found for logger "trimesh" 81 | Running code on cpu 82 | Found 6778 'ShapeNetV2' models 83 | 1 models in total ... 84 | R: [[-0.38369974 0.57926947 0.7191812 ] 85 | [-0.5103006 0.5160787 -0.6879362 ] 86 | [-0.7696545 -0.6309594 0.09758216]] 87 | t: [[-0.01014123] 88 | [-0.04051362] 89 | [ 0.00080032]] 90 | e: [1.3734075, 1.3150274] 91 | K: [-0.33036062, 0.23313229] 92 | R: [[-0.5362834 0.04193148 -0.8429957 ] 93 | [ 0.78693724 0.3859552 -0.4814232 ] 94 | [ 0.3051718 -0.92156404 -0.23997879]] 95 | t: [[-0.18105912] 96 | [ 0.11244812] 97 | [ 0.09697253]] 98 | e: [0.7390462, 1.157957] 99 | ... 100 | R: [[ 0.9701394 -0.24186404 -0.01820319] 101 | [-0.24252652 -0.96831805 -0.059508 ] 102 | [-0.00323363 0.06214581 -0.99806195]] 103 | t: [[-0.02585344] 104 | [-0.15909581] 105 | [-0.091534 ]] 106 | e: [0.4002924, 0.40005013] 107 | K: [0.41454253, 0.3229351] 108 | 0 3.581116e-08 109 | 1 0.99999857 110 | 2 0.99999917 111 | 3 3.7113843e-08 112 | 4 0.9999982 113 | 5 0.9999975 114 | 6 0.9999988 115 | 7 0.9999999 116 | 8 3.5458616e-08 117 | 9 3.721507e-08 118 | 10 3.711448e-08 119 | 11 3.9621053e-08 120 | 12 3.7611613e-08 121 | 13 0.9999976 122 | 14 0.99999905 123 | 15 0.9999981 124 | 16 0.9999982 125 | 17 0.99999785 126 | Using 11 primitives out of 18 127 | ``` 128 | and a chair should appear. 129 | 130 | Training 131 | -------- 132 | 133 | To train a new network from scratch we provide the `train_network.py` script. 134 | 135 | You can simply execute it by typing 136 | ``` 137 | $ ./train_network.py ~/data/03001627/ /tmp/ --use_sq --lr 1e-4 --n_primitives 20 --train_with_bernoulli --dataset_type shapenet_v2 --use_chamfer 138 | Running code on cpu 139 | Save experiment statistics in 26EKQBNTG 140 | Found 6778 'ShapeNetV2' models 141 | 6778 models in total ... 142 | 1000/1000 [==============================] - 6s 6ms/step 143 | Epoch 1/150 | | 15/500 - loss: 0.0110078 - pcl_to_prim: 0.0024909 - prim_to_pcl: 0.0085169 - exp_n_prims: 9.8473 144 | 145 | ``` 146 | 147 | You need to specify the path to the directory containing the dataset directory 148 | as well as the path to save the generated files such as the trained models. 149 | Note that the script automatically generates a subfolder inside the specified 150 | output directory (in this case `26EKQBNTG`), where it saves the trained models, 151 | three .txt files with the loss evolution and a .json file with the parameters 152 | used for the current experiment. 153 | 154 | Visualizing Superquadrics 155 | ------------------------- 156 | 157 | We also provide the `visualize_sq.py` script which allows you to quickly 158 | visualize superquadrics given a set of parameters as a set of points sampled on 159 | the surface of the superquadric surface. Please note that you need to install 160 | python-tk to be able to use this script. You can do this my simply writing 161 | `sudo apt install python-tk`. 162 | 163 | You can simply execute it by providing your preferred shape and size parametrization as follows: 164 | ``` 165 | $ ./visualize_sq.py --shape 1.0,1.0 --size 0.25,0.25,0.25 166 | ``` 167 | Below are some example images of various superquadrics using different shape 168 | and size parameters 169 | ![Example 1](img/00.png) 170 | ![Example 2](img/01.png) 171 | ![Example 3](img/02.png) 172 | 173 | 174 | Contribution 175 | ------------ 176 | Contributions such as bug fixes, bug reports, suggestions etc. are more than 177 | welcome and should be submitted in the form of new issues and/or pull requests 178 | on Github. 179 | 180 | License 181 | ------- 182 | Our code is released under the MIT license which practically allows anyone to do anything with it. 183 | MIT license found in the LICENSE file. 184 | 185 | 186 | Relevant Research 187 | ------------------ 188 | Below we list some papers that are relevant to the provided code. 189 | 190 | **Ours:** 191 | - Superquadrics Revisited: Learning 3D Shape Parsing beyond Cuboids [pdf](https://arxiv.org/pdf/1904.09970.pdf) [blog](https://autonomousvision.github.io/superquadrics-revisited/) 192 | 193 | **By Others:** 194 | - Learning Shape Abstractions by Assembling Volumetric Primitives [pdf](https://arxiv.org/pdf/1612.00404.pdf) 195 | - 3D-PRNN: Generating Shape Primitives with Recurrent Neural Networks [pdf](https://arxiv.org/abs/1708.01648.pdf) 196 | - Im2Struct: Recovering 3D Shape Structure From a Single RGB Image [pdf](http://openaccess.thecvf.com/content_cvpr_2018/html/Niu_Im2Struct_Recovering_3D_CVPR_2018_paper.pdf) 197 | 198 | Below we also list some more papers that are more closely related to superquadrics 199 | - Equal-Distance Sampling of Supercllipse Models [pdf](https://pdfs.semanticscholar.org/3e6f/f812b392f9eb70915b3c16e7bfbd57df379d.pdf) 200 | - Revisiting Superquadric Fitting: A Numerically Stable Formulation [link](https://ieeexplore.ieee.org/document/8128485) 201 | 202 | 203 | Citation 204 | -------- 205 | If you found this work influential or helpful for your research, please consider citing 206 | 207 | ``` 208 | @Inproceedings{Paschalidou2019CVPR, 209 | title = {Superquadrics Revisited: Learning 3D Shape Parsing beyond Cuboids}, 210 | author = {Paschalidou, Despoina and Ulusoy, Ali Osman and Geiger, Andreas}, 211 | booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)}, 212 | year = {2019} 213 | } 214 | ``` 215 | -------------------------------------------------------------------------------- /config/chair_T26AK2FES_model_699: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/config/chair_T26AK2FES_model_699 -------------------------------------------------------------------------------- /demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.json: -------------------------------------------------------------------------------- 1 | {"max": [0.320516, 0.734928, 0.328415], "centroid": [2.4996886303102894e-08, 0.34210536311030604, 0.030300756370850807], "id": "dac4af24e2facd7d3000ca4b04fcd6ac", "numVertices": 1106, "min": [-0.320516, 0.0, -0.257756]} 2 | -------------------------------------------------------------------------------- /demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.mtl: -------------------------------------------------------------------------------- 1 | # File produced by Open Asset Import Library (http://www.assimp.sf.net) 2 | # (assimp v3.2.202087883) 3 | 4 | newmtl ForegroundColor 5 | Kd 0 0 0 6 | Ka 0 0 0 7 | Ks 0.33 0.33 0.33 8 | Ke 0 0 0 9 | d 1 10 | Ns 20 11 | illum 2 12 | 13 | newmtl FrontColorNoCulling 14 | Kd 0.882353 0.882353 0.784314 15 | Ka 0 0 0 16 | Ks 0.33 0.33 0.33 17 | Ke 0 0 0 18 | d 1 19 | Ns 20 20 | illum 2 21 | 22 | newmtl __Charcoal_noCulling 23 | Kd 0.137255 0.137255 0.137255 24 | Ka 0 0 0 25 | Ks 0.33 0.33 0.33 26 | Ke 0 0 0 27 | d 1 28 | Ns 20 29 | illum 2 30 | 31 | newmtl __GhostWhite_noCulling 32 | Kd 0.972549 0.972549 1 33 | Ka 0 0 0 34 | Ks 0.33 0.33 0.33 35 | Ke 0 0 0 36 | d 1 37 | Ns 20 38 | illum 2 39 | 40 | newmtl __GreenYellow_1noCulling 41 | Kd 0.72549 0.94902 0.411765 42 | Ka 0 0 0 43 | Ks 0.33 0.33 0.33 44 | Ke 0 0 0 45 | d 1 46 | Ns 20 47 | illum 2 48 | 49 | newmtl __Silver_ 50 | Kd 0.752941 0.752941 0.752941 51 | Ka 0 0 0 52 | Ks 0.33 0.33 0.33 53 | Ke 0 0 0 54 | d 1 55 | Ns 20 56 | illum 2 57 | 58 | -------------------------------------------------------------------------------- /demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.solid.binvox: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.solid.binvox -------------------------------------------------------------------------------- /demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.surface.binvox: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/demo/03001627/dac4af24e2facd7d3000ca4b04fcd6ac/models/model_normalized.surface.binvox -------------------------------------------------------------------------------- /img/00.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/img/00.png -------------------------------------------------------------------------------- /img/01.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/img/01.png -------------------------------------------------------------------------------- /img/02.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/img/02.png -------------------------------------------------------------------------------- /learnable_primitives/__init__.py: -------------------------------------------------------------------------------- 1 | """Recover the geometry of a 3D mesh as a set of learnable superquadric 2 | surfaces. 3 | """ 4 | 5 | __author__ = "Despoina Paschalidou" 6 | __copyright__ = "Copyright (c) 2018 Max Planck Institute for Intelligent Systems" 7 | __license__ = "MIT" 8 | __maintainer__ = "Despoina Paschalidou" 9 | __email__ = "despoina.paschalidou@tue.mpg.de" 10 | __url__ = "http://superquadrics.com/" 11 | __version__ = "0.1" 12 | -------------------------------------------------------------------------------- /learnable_primitives/common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/learnable_primitives/common/__init__.py -------------------------------------------------------------------------------- /learnable_primitives/common/batch_provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import threading 4 | import time 5 | 6 | import numpy as np 7 | import torch 8 | 9 | from ..utils.progbar import Progbar 10 | 11 | 12 | class BatchProvider(object): 13 | """BatchProvider class is a wrapper class to generate batches to train the 14 | network. 15 | """ 16 | def __init__( 17 | self, 18 | dataset, 19 | batch_size, 20 | cache_size=500, 21 | verbose=1 22 | ): 23 | self._dataset = dataset 24 | self.batch_size = batch_size 25 | self.cache_size = cache_size 26 | self.verbose = verbose 27 | 28 | # Shape of the input tensor to the network 29 | self.input_shape = self._dataset.input_dim 30 | # Shape of the output tensor from the network 31 | self.output_shape = self._dataset.output_dim 32 | 33 | # Member variable to stop the thread (to be set via a call to stop) 34 | self._stop = False 35 | self._ready = False 36 | 37 | # Start a thread to fill the cache 38 | self._start_producer_thread() 39 | 40 | def ready(self, blocking=True): 41 | while blocking and not self._ready: 42 | time.sleep(0.1) 43 | return self._ready 44 | 45 | def stop(self): 46 | self._stop = True 47 | self._producer_thread.join() 48 | 49 | def __iter__(self): 50 | return self 51 | 52 | def __next__(self): 53 | return next() 54 | 55 | def next(self): 56 | idxs = np.random.randint(0, self.cache_size, size=self.batch_size) 57 | with self.cache_lock: 58 | return self.X[idxs], self.y[idxs] 59 | 60 | def _start_producer_thread(self): 61 | # This is going to be the amount of cached elements 62 | N = self.cache_size 63 | self.X = torch.empty((N,) + self.input_shape, dtype=torch.float32) 64 | self.y = torch.empty((N,) + self.output_shape, dtype=torch.float32) 65 | 66 | self.cache_lock = threading.RLock() 67 | self._producer_thread = threading.Thread(target=self._producer) 68 | self._producer_thread.daemon = True 69 | self._producer_thread.start() 70 | 71 | def _producer(self): 72 | N = self.cache_size 73 | passes = 0 74 | if self.verbose > 0: 75 | prog = Progbar(N) 76 | 77 | while True: 78 | # Acquire the lock for the whole first pass 79 | if passes == 0: 80 | self.cache_lock.acquire() 81 | 82 | for idx in range(N): 83 | # We 're done stop now 84 | if self._stop and passes > 0: 85 | return 86 | 87 | while True: 88 | sample = self._dataset.get_random_datapoint() 89 | X = sample[0] 90 | y = sample[1] 91 | if X is None or y is None: 92 | continue 93 | break 94 | 95 | # Do the copy to the cache but make sure you lock first and 96 | # unlock afterwards 97 | with self.cache_lock: 98 | try: 99 | self.X[idx] = X 100 | self.y[idx] = y 101 | except Exception as e: 102 | sys.stderr.write( 103 | "Exception caught in producer thread " + str(e) 104 | ) 105 | 106 | # Show progress if it is the first pass 107 | if passes == 0 and self.verbose > 0: 108 | prog.update(idx + 1) 109 | 110 | # Release the lock if it was the first pass 111 | if passes == 0: 112 | self._ready = True 113 | self.cache_lock.release() 114 | 115 | # Count the passes 116 | passes += 1 117 | -------------------------------------------------------------------------------- /learnable_primitives/common/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import torch 5 | from torch.utils.data import Dataset 6 | from torchvision import transforms 7 | from torchvision.transforms import Normalize as TorchNormalize 8 | 9 | 10 | class BaseDataset(Dataset): 11 | """Dataset is a wrapper for all datasets we have 12 | """ 13 | def __init__(self, dataset_object, voxelizer_factory, 14 | n_points_from_mesh=1000, transform=None): 15 | """ 16 | Arguments: 17 | --------- 18 | dataset_object: a dataset object that can be either ShapeNetObject 19 | or SurrealBodiesObject 20 | voxelizer_factory: a factory that creates a voxelizer object to 21 | voxelizes the ground-truth mesh of a 22 | ShapeNetModel instance 23 | n_points_from_mesh: int, the number of points to be sampled from 24 | the groundtruth mesh 25 | transform: Callable that applies a transform to a sample 26 | """ 27 | self._dataset_object = dataset_object 28 | print "%d models in total ..." % (len(self._dataset_object)) 29 | 30 | # Number of samples to use for supervision 31 | self._n_points_from_mesh = n_points_from_mesh 32 | # Get the voxelizer to be used 33 | self._voxelizer = voxelizer_factory.voxelizer 34 | # Operations on the datapoints 35 | self.transform = transform 36 | 37 | self._input_dim = self._voxelizer.output_shape 38 | # 3 for the points and 3 for the normals of each face 39 | self._output_dim = (n_points_from_mesh, 6) 40 | 41 | def __len__(self): 42 | return len(self._dataset_object) 43 | 44 | def __getitem__(self, idx): 45 | m = self._dataset_object[idx] 46 | # print m.path_to_mesh_file 47 | # print m.path_to_tsdf_file 48 | X = self._voxelizer.get_X(m) 49 | mesh = m.groundtruth_mesh 50 | y_target = mesh.sample_faces(self._n_points_from_mesh) 51 | 52 | datapoint = ( 53 | X, 54 | y_target.astype(np.float32) 55 | ) 56 | 57 | # Store the dimentionality of the input tensor and the y_target tensor 58 | self._input_dim = datapoint[0].shape 59 | self._output_dim = datapoint[1].shape 60 | 61 | if self.transform: 62 | datapoint = self.transform(datapoint) 63 | 64 | return datapoint 65 | 66 | def get_random_datapoint(self, idx=None): 67 | if idx is None: 68 | idx = np.random.choice(np.arange(self.__len__())) 69 | return self.__getitem__(idx) 70 | 71 | @property 72 | def input_dim(self): 73 | return self._input_dim 74 | 75 | @property 76 | def output_dim(self): 77 | return self._output_dim 78 | 79 | 80 | class MeshParserDataset(BaseDataset): 81 | """MeshDataset is a class that simply parses meshes 82 | """ 83 | def __init__(self, dataset_object): 84 | self._dataset_object = dataset_object 85 | print "%d models in total ..." % (len(self._dataset_object)) 86 | 87 | def __getitem__(self, idx): 88 | m = self._dataset_object[idx] 89 | mesh = m.groundtruth_mesh 90 | return (mesh, m.tag) 91 | 92 | @property 93 | def input_dim(self): 94 | raise NotImplementedError() 95 | 96 | @property 97 | def output_dim(self): 98 | raise NotImplementedError() 99 | 100 | 101 | class DatasetWithTags(BaseDataset): 102 | def __getitem__(self, idx): 103 | m = self._dataset_object[idx] 104 | X = self._voxelizer.get_X(m) 105 | mesh = m.groundtruth_mesh 106 | y_target = mesh.sample_faces(self._n_points_from_mesh).T 107 | 108 | datapoint = ( 109 | X, 110 | y_target.astype(np.float32).T 111 | ) 112 | 113 | # Store the dimentionality of the input tensor and the y_target tensor 114 | self._input_dim = datapoint[0].shape 115 | self._output_dim = datapoint[1].shape 116 | 117 | if self.transform: 118 | datapoint = self.transform(datapoint) 119 | 120 | return datapoint + (m.tag.split("/")[-1],) 121 | 122 | 123 | class DatasetWithTagsAndFaces(BaseDataset): 124 | def __getitem__(self, idx): 125 | m = self._dataset_object[idx] 126 | X = self._voxelizer.get_X(m) 127 | mesh = m.groundtruth_mesh 128 | y_target = mesh.sample_faces(self._n_points_from_mesh).T 129 | 130 | datapoint = ( 131 | X, 132 | y_target.astype(np.float32).T 133 | ) 134 | 135 | # Store the dimentionality of the input tensor and the y_target tensor 136 | self._input_dim = datapoint[0].shape 137 | self._output_dim = datapoint[1].shape 138 | 139 | if self.transform: 140 | datapoint = self.transform(datapoint) 141 | 142 | tag = m.tag.split("/")[-1] 143 | return datapoint + (tag, m.path_to_mesh_file, mesh.points) 144 | 145 | 146 | class ToTensor(object): 147 | """Convert ndarrays in sample to Tensors.""" 148 | def __call__(self, sample): 149 | X, y_target = sample 150 | 151 | # Do some sanity checks to ensure that the inputs have the appropriate 152 | # dimensionality 153 | # assert len(X.shape) == 4 154 | return (torch.from_numpy(X), torch.from_numpy(y_target).float()) 155 | 156 | 157 | class Normalize(object): 158 | """Normalize image based based on ImageNet.""" 159 | def __call__(self, sample): 160 | X, y_target = sample 161 | X = X.float() 162 | 163 | # The normalization will only affect X 164 | normalize = TorchNormalize( 165 | mean=[0.485, 0.456, 0.406], 166 | std=[0.229, 0.224, 0.225] 167 | ) 168 | X = X.float() / 255.0 169 | return (normalize(X), y_target) 170 | 171 | 172 | def compose_transformations(voxelizer_factory): 173 | transformations = [ToTensor()] 174 | if voxelizer_factory == "image": 175 | transformations.append(Normalize()) 176 | 177 | return transforms.Compose(transformations) 178 | 179 | 180 | def get_dataset_type(loss_type): 181 | return { 182 | "euclidean_dual_loss": BaseDataset 183 | }[loss_type] 184 | -------------------------------------------------------------------------------- /learnable_primitives/common/model_factory.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import re 3 | try: 4 | from functools import lru_cache 5 | except ImportError: 6 | from backports.functools_lru_cache import lru_cache 7 | import os 8 | import pickle 9 | from PIL import Image 10 | 11 | from ..mesh import MeshFactory 12 | 13 | 14 | class BaseModel(object): 15 | """BaseModel class is wrapper for all models, independent of dataset. Every 16 | model has a unique model_tag, mesh_file and Mesh object. Optionally, it 17 | can also have a tsdf file. 18 | """ 19 | def __init__(self, tag, path_to_mesh_file, path_to_tsdf_file, 20 | mesh=None, images_dir=None): 21 | self._path_to_tsdf_file = path_to_tsdf_file 22 | self._path_to_mesh_file = path_to_mesh_file 23 | self._tag = tag 24 | self.images_dir = images_dir 25 | 26 | # A Mesh object 27 | self._gt_mesh = mesh 28 | # A numpy array containing the TSDF 29 | self._tsdf = None 30 | # A list containing the images 31 | self._images = [] 32 | self._points_on_surface = None 33 | # Variable to cache mesh internal points 34 | self._internal_points = None 35 | 36 | @property 37 | def tag(self): 38 | return self._tag 39 | 40 | @property 41 | def path_to_mesh_file(self): 42 | return self._path_to_mesh_file 43 | 44 | @property 45 | def path_to_tsdf_file(self): 46 | return self._path_to_tsdf_file 47 | 48 | @property 49 | def groundtruth_mesh(self): 50 | if self._gt_mesh is None: 51 | self._gt_mesh = MeshFactory.from_file(self._path_to_mesh_file) 52 | return self._gt_mesh 53 | 54 | @groundtruth_mesh.setter 55 | def groundtruth_mesh(self, mesh): 56 | if self._gt_mesh is not None: 57 | raise RuntimeError("Trying to overwrite a mesh") 58 | self._gt_mesh = mesh 59 | 60 | @property 61 | def tsdf(self): 62 | if self._tsdf is None: 63 | # Check that the file exists 64 | if not os.path.isfile(self._path_to_tsdf_file): 65 | raise Exception( 66 | "The mesh path %s does not exist" % ( 67 | self._path_to_tsdf_file 68 | ) 69 | ) 70 | self._tsdf = np.load(self._path_to_tsdf_file) 71 | return self._tsdf 72 | 73 | def _n_images(self): 74 | return len(os.listdir(self.images_dir)) 75 | 76 | @property 77 | def images(self): 78 | if not self._images: 79 | for path_to_img in sorted(os.listdir(self.images_dir)): 80 | if path_to_img.endswith((".jpg", ".png")): 81 | image_path = os.path.join(self.images_dir, path_to_img) 82 | self._images.append( 83 | np.array(Image.open(image_path).convert("RGB")) 84 | ) 85 | return self._images 86 | 87 | @property 88 | def random_image(self): 89 | ri = np.random.choice(len(self.images)) 90 | return self.images[ri] 91 | 92 | 93 | class ModelsCollection(object): 94 | def __len__(self): 95 | raise NotImplementedError() 96 | 97 | def _get_model(self, i): 98 | raise NotImplementedError() 99 | 100 | def __getitem__(self, i): 101 | if i >= len(self): 102 | raise IndexError() 103 | return self._get_model(i) 104 | 105 | 106 | class ModelsSubset(ModelsCollection): 107 | def __init__(self, collection, subset): 108 | self._collection = collection 109 | self._subset = subset 110 | 111 | def __len__(self): 112 | return len(self._subset) 113 | 114 | def _get_sample(self, i): 115 | return self._collection[self._subset[i]] 116 | 117 | def __getitem__(self, i): 118 | if i >= len(self): 119 | raise IndexError() 120 | return self._get_sample(i) 121 | 122 | 123 | class TagSubset(ModelsSubset): 124 | def __init__(self, collection, tags): 125 | subset = [i for (i, m) in enumerate(collection) if m.tag in tags] 126 | super(TagSubset, self).__init__(collection, subset) 127 | 128 | 129 | class ShapeNetQuad(ModelsCollection): 130 | def __init__(self, base_dir): 131 | self._tags = sorted(os.listdir(base_dir)) 132 | self._paths = [os.path.join(base_dir, x) for x in self._tags] 133 | 134 | print "Found {} 'ShapeNetQuad' models".format(len(self)) 135 | 136 | def __len__(self): 137 | return len(self._paths) 138 | 139 | def _get_model(self, i): 140 | return BaseModel( 141 | self._tags[i], 142 | os.path.join(self._paths[i], "model.obj"), 143 | None 144 | ) 145 | 146 | 147 | class ShapeNetV1(ModelsCollection): 148 | def __init__(self, base_dir): 149 | self._tags = sorted( 150 | x for x in os.listdir(base_dir) 151 | if os.path.isdir(os.path.join(base_dir, x)) 152 | ) 153 | self._paths = [os.path.join(base_dir, x) for x in self._tags] 154 | 155 | print "Found {} 'ShapeNetV1' models".format(len(self)) 156 | 157 | def __len__(self): 158 | return len(self._paths) 159 | 160 | def _get_model(self, i): 161 | return BaseModel( 162 | self._tags[i], 163 | os.path.join(self._paths[i], "model_watertight.off"), 164 | None, 165 | images_dir=os.path.join(self._paths[i], "img_choy2016") 166 | ) 167 | 168 | 169 | class ShapeNetV2(ModelsCollection): 170 | def __init__(self, base_dir): 171 | self._tags = sorted( 172 | x for x in os.listdir(base_dir) 173 | if os.path.isdir(os.path.join(base_dir, x)) 174 | ) 175 | self._paths = [os.path.join(base_dir, x) for x in self._tags] 176 | 177 | print "Found {} 'ShapeNetV2' models".format(len(self)) 178 | 179 | def __len__(self): 180 | return len(self._paths) 181 | 182 | def _get_model(self, i): 183 | return BaseModel( 184 | self._tags[i], 185 | os.path.join(self._paths[i], "models", "model_normalized.obj"), 186 | None 187 | ) 188 | 189 | 190 | class SurrealHumanBodies(ModelsCollection): 191 | def __init__(self, base_dir): 192 | self._base_dir = base_dir 193 | self._tags = sorted({x[:6] for x in os.listdir(base_dir)}) 194 | 195 | print "Found {} 'Surreal' models".format(len(self)) 196 | 197 | def __len__(self): 198 | return len(self._tags) 199 | 200 | def _get_model(self, i): 201 | return BaseModel( 202 | self._tags[i], 203 | os.path.join(self._base_dir, "{}.obj".format(self._tags[i])), 204 | None 205 | ) 206 | 207 | 208 | class DynamicFaust(ModelsCollection): 209 | def __init__(self, base_dir): 210 | self._base_dir = base_dir 211 | self._paths = sorted([ 212 | d 213 | for d in os.listdir(self._base_dir) 214 | if os.path.isdir(os.path.join(self._base_dir, d)) 215 | ]) 216 | 217 | self._tags = sorted([ 218 | "{}:{}".format(d, l[:-4]) for d in self._paths 219 | for l in os.listdir(os.path.join(self._base_dir, d, "mesh_seq")) 220 | ]) 221 | 222 | print "Found {} 'Surreal' models".format(len(self)) 223 | 224 | def __len__(self): 225 | return len(self._tags) 226 | 227 | def _get_model(self, i): 228 | tag_parts = self._tags[i].split(":") 229 | model_dir = os.path.join(self._base_dir, tag_parts[0]) 230 | return BaseModel( 231 | self._tags[i], 232 | os.path.join(model_dir, "mesh_seq", "{}.obj".format(tag_parts[1])), 233 | None 234 | ) 235 | 236 | 237 | class MeshCache(ModelsCollection): 238 | """Cache the meshes from a collection and give them to the model before 239 | returning it.""" 240 | def __init__(self, collection): 241 | self._collection = collection 242 | self._meshes = [None]*len(collection) 243 | 244 | def __len__(self): 245 | return len(self._collection) 246 | 247 | def _get_model(self, i): 248 | model = self._collection._get_model(i) 249 | if self._meshes[i] is not None: 250 | model.groundtruth_mesh = self._meshes[i] 251 | else: 252 | self._meshes[i] = model.groundtruth_mesh 253 | 254 | return model 255 | 256 | 257 | class LRUCache(ModelsCollection): 258 | def __init__(self, collection, n=2000): 259 | self._collection = collection 260 | self._model_getter = lru_cache(n)(self._inner_get_model) 261 | 262 | def __len__(self): 263 | return len(self._collection) 264 | 265 | def _inner_get_model(self, i): 266 | return self._collection._get_model(i) 267 | 268 | def _get_model(self, i): 269 | return self._model_getter(i) 270 | 271 | 272 | def model_factory(dataset_type): 273 | return { 274 | "shapenet_quad": ShapeNetQuad, 275 | "shapenet_v1": ShapeNetV1, 276 | "shapenet_v2": ShapeNetV2, 277 | "surreal_bodies": SurrealHumanBodies, 278 | "dynamic_faust": DynamicFaust 279 | }[dataset_type] 280 | 281 | 282 | class DatasetBuilder(object): 283 | def __init__(self): 284 | self._dataset_class = None 285 | self._cache_meshes = False 286 | self._lru_cache = 0 287 | self._tags = [] 288 | 289 | def with_dataset(self, dataset): 290 | self._dataset_class = model_factory(dataset) 291 | return self 292 | 293 | def with_cache_meshes(self): 294 | self._cache_meshes = True 295 | return self 296 | 297 | def without_cache_meshes(self): 298 | self._cache_meshes = False 299 | return self 300 | 301 | def lru_cache(self, n=2000): 302 | self._lru_cache = n 303 | return self 304 | 305 | def filter_tags(self, tags): 306 | self._tags = tags 307 | return self 308 | 309 | def build(self, base_dir): 310 | dataset = self._dataset_class(base_dir) 311 | if self._cache_meshes: 312 | dataset = MeshCache(dataset) 313 | if self._lru_cache > 0: 314 | dataset = LRUCache(dataset, self._lru_cache) 315 | if self._tags: 316 | dataset = TagSubset(dataset, self._tags) 317 | 318 | return dataset 319 | -------------------------------------------------------------------------------- /learnable_primitives/fast_sampler/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from ._sampler import step_eta, step_omega, collect_etas, collect_omegas, \ 3 | fast_sample, fast_sample_on_batch 4 | -------------------------------------------------------------------------------- /learnable_primitives/fast_sampler/_sampler.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cython 3 | from libc.math cimport sqrt, sin, cos, floor, M_PI_2, M_PI 4 | from libc.float cimport FLT_MIN 5 | cimport numpy as cnp 6 | 7 | 8 | cdef inline float sign(float x): 9 | if x > 0: 10 | return 1.0 11 | if x < 0: 12 | return -1.0 13 | if x == 0: 14 | return 0.0 15 | 16 | 17 | @cython.wraparound(False) 18 | cdef inline float _step_eta( 19 | float D_eta, 20 | float a1, 21 | float a2, 22 | float a3, 23 | float e1, 24 | float e2, 25 | float eta, 26 | float omega 27 | ): 28 | cdef float t1, t2, t3, t 29 | t1 = (a1**2) * sin(eta)**4 30 | t1 *= abs(cos(eta))**(2*e1) 31 | t1 *= abs(cos(omega))**(2*e2) 32 | t1 *= sign(cos(eta))**2 33 | t1 *= sign(cos(omega))**2 34 | t1 *= sign(cos(eta))**2 35 | 36 | t2 = (a2**2)*(sin(eta)**4) 37 | t2 *= abs(sin(omega))**(2*e2) 38 | t2 *= abs(cos(eta))**(2*e1) 39 | t2 *= sign(sin(omega))**2 40 | t2 *= sign(cos(eta))**2 41 | t2 *= sign(cos(eta))**2 42 | 43 | t3 = (a3**2)*(cos(eta)**4) 44 | t3 *= abs(sin(eta))**(2*e1) 45 | t3 *= sign(sin(eta))**2 46 | t3 *= sign(sin(eta))**2 47 | 48 | t = sqrt(1.0 / (t1 + t2 + t3))*sin(eta)*cos(eta) 49 | return abs(D_eta * (t / e1)) 50 | 51 | 52 | @cython.wraparound(False) 53 | cpdef float step_eta( 54 | float D_eta, 55 | float a1, 56 | float a2, 57 | float a3, 58 | float e1, 59 | float e2, 60 | float eta, 61 | float omega 62 | ): 63 | cdef float step 64 | step = _step_eta( 65 | D_eta, 66 | a1, 67 | a2, 68 | a3, 69 | e1, 70 | e2, 71 | eta, 72 | omega 73 | ) 74 | return step 75 | 76 | 77 | @cython.wraparound(False) 78 | cdef inline float _step_omega( 79 | float D_omega, 80 | float a1, 81 | float a2, 82 | float a3, 83 | float e1, 84 | float e2, 85 | float eta, 86 | float omega 87 | ): 88 | cdef float t1, t2, t3, t 89 | t1 = (a1**2) * sin(omega)**4 90 | t1 *= abs(cos(omega))**(2*e2) 91 | t1 *= sign(cos(omega))**2 92 | t1 *= sign(cos(omega))**2 93 | 94 | t2 = (a2**2) * cos(omega)**4 95 | t2 *= abs(sin(omega))**(2*e2) 96 | t2 *= sign(sin(omega))**2 97 | t2 *= sign(sin(omega))**2 98 | 99 | t = sqrt(abs(cos(eta))**(-2*e1) / (t1 + t2)) 100 | t *= D_omega * sin(omega) * cos(omega) 101 | return abs(t / (e2 * sign(cos(eta)))) 102 | 103 | 104 | @cython.wraparound(False) 105 | cpdef float step_omega( 106 | float D_omega, 107 | float a1, 108 | float a2, 109 | float a3, 110 | float e1, 111 | float e2, 112 | float eta, 113 | float omega 114 | ): 115 | cdef float step 116 | step = _step_omega( 117 | D_omega, 118 | a1, 119 | a2, 120 | a3, 121 | e1, 122 | e2, 123 | eta, 124 | omega 125 | ) 126 | return step 127 | 128 | 129 | @cython.boundscheck(False) 130 | @cython.wraparound(False) 131 | cdef inline int update_etas( 132 | float eta_initial, 133 | float D_eta, 134 | float a1, 135 | float a2, 136 | float a3, 137 | float e1, 138 | float e2, 139 | float eta, 140 | float omega, 141 | float[:] etas 142 | ): 143 | # counter to keep track of the number of etas added in the list 144 | cdef int i = 0 145 | cdef float eta_value = eta_initial 146 | while eta_value < M_PI/2: 147 | # Update etas with the value 148 | etas[i] = eta_value 149 | if eta_value == 0.0: 150 | eta_value += 0.01 151 | continue 152 | # increment counter by 1 153 | eta_update = _step_eta(D_eta, a1, a2, a3, e1, e2, eta_value, omega) 154 | eta_value += eta_update 155 | i += 1 156 | return i 157 | 158 | 159 | @cython.boundscheck(False) 160 | @cython.wraparound(False) 161 | cdef inline int update_omegas( 162 | float omega_initial, 163 | float D_omega, 164 | float a1, 165 | float a2, 166 | float a3, 167 | float e1, 168 | float e2, 169 | float eta, 170 | float omega, 171 | float[:] omegas 172 | ): 173 | # counter to keep track of the number of omegas added in the list 174 | cdef int i = 0 175 | cdef float omega_value = omega_initial 176 | while omega_value <= M_PI: 177 | # Update etas with the value 178 | omegas[i] = omega_value 179 | if omega_value == 0.0: 180 | omega_value += 0.01 181 | continue 182 | # increment counter by 1 183 | omega_update = _step_omega(D_omega, a1, a2, a3, e1, e2, 184 | eta, omega_value) 185 | omega_value += omega_update 186 | i += 1 187 | return i 188 | 189 | 190 | @cython.boundscheck(False) 191 | @cython.wraparound(False) 192 | cdef inline int etas_is_not_acceptable(float[:] etas, int N): 193 | cdef float C = 0.1 194 | cdef int c1, c2, c3, i 195 | cdef float max_d = 0.0 196 | cdef float current_val = 0.0 197 | for i in range(1, N-1): 198 | current_val = etas[i] - etas[i-1] 199 | if current_val > max_d: 200 | max_d = current_val 201 | 202 | c1 = max_d > C 203 | c2 = floor(abs(etas[0] + M_PI_2)) > C 204 | c3 = floor(abs(M_PI_2 - etas[N-1])) > 0.01 205 | return c1 or c2 or c3 206 | 207 | 208 | @cython.boundscheck(False) 209 | @cython.wraparound(False) 210 | cdef inline int omegas_is_not_acceptable(float[:] omegas, int N): 211 | cdef float C = 0.1 212 | cdef int c1, c2, c3, i 213 | cdef float max_d = 0.0 214 | cdef float current_val = 0.0 215 | for i in range(1, N-1): 216 | current_val = omegas[i] - omegas[i-1] 217 | if current_val > max_d: 218 | max_d = current_val 219 | 220 | c1 = max_d > C 221 | c2 = floor(abs(omegas[0] + M_PI)) > C 222 | c3 = floor(abs(M_PI - omegas[N-1])) > 0.01 223 | return c1 or c2 or c3 224 | 225 | 226 | @cython.boundscheck(False) 227 | @cython.wraparound(False) 228 | cpdef cnp.ndarray[cnp.float32_t, ndim=1] collect_etas( 229 | float eta_initial, 230 | float D_eta, 231 | float a1, 232 | float a2, 233 | float a3, 234 | float e1, 235 | float e2, 236 | float eta, 237 | float omega 238 | ): 239 | cdef float D_eta_value = D_eta 240 | cdef cnp.ndarray[cnp.float32_t, ndim=1] etas = \ 241 | np.empty(shape=(100000,), dtype=np.float32) 242 | # Do the first update 243 | cdef int N 244 | N = update_etas(eta_initial, D_eta_value, a1, a2, a3, e1, e2, 245 | eta, omega, etas) 246 | 247 | # while (etas_is_not_acceptable(etas, N) == 1) and (N>500 and N<200): 248 | while (etas_is_not_acceptable(etas, N) == 1): 249 | D_eta_value = D_eta_value / 2.0 250 | # Update etas with the new D_eta_value 251 | N = update_etas(eta_initial, D_eta_value, a1, a2, a3, e1, e2, 252 | eta, omega, etas) 253 | return etas[:N] 254 | 255 | 256 | @cython.boundscheck(False) 257 | @cython.wraparound(False) 258 | cpdef cnp.ndarray[cnp.float32_t, ndim=1] collect_omegas( 259 | float omega_initial, 260 | float D_omega, 261 | float a1, 262 | float a2, 263 | float a3, 264 | float e1, 265 | float e2, 266 | float eta, 267 | float omega 268 | ): 269 | cdef float D_omega_value = D_omega 270 | cdef cnp.ndarray[cnp.float32_t, ndim=1] omegas = \ 271 | np.empty(shape=(100000,), dtype=np.float32) 272 | # Do the first update 273 | cdef int N 274 | N = update_omegas(omega_initial, D_omega, a1, a2, a3, e1, e2, 275 | eta, omega, omegas) 276 | while (omegas_is_not_acceptable(omegas, N) == 1): 277 | D_omega_value = D_omega_value / 2.0 278 | # Update omegas with the new D_omega_value 279 | N = update_omegas(omega_initial, D_omega_value, a1, a2, a3, e1, e2, 280 | eta, omega, omegas) 281 | 282 | return omegas[:N] 283 | 284 | 285 | cdef inline float fexp(float x, float p): 286 | return sign(x) * (abs(x)**p) 287 | 288 | 289 | @cython.boundscheck(False) 290 | @cython.wraparound(False) 291 | cdef inline void xy(float theta, float a1, float a2, float e, float[2] C): 292 | C[0] = a1 * fexp(cos(theta), e) 293 | C[1] = a2 * fexp(sin(theta), e) 294 | 295 | 296 | @cython.boundscheck(False) 297 | @cython.wraparound(False) 298 | cdef inline float distance(float[2] a, float[2] b): 299 | return sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2) 300 | 301 | 302 | @cython.boundscheck(False) 303 | @cython.wraparound(False) 304 | cdef void _sample_superellipse_divide_conquer_inner( 305 | float a1, 306 | float a2, 307 | float e, 308 | float theta_A, 309 | float theta_B, 310 | float[2] A, 311 | float[2] B, 312 | int N, 313 | float[::1] thetas, 314 | int * idx 315 | ): 316 | if N <= 0: 317 | return 318 | 319 | cdef float theta 320 | cdef float dA 321 | cdef float dB 322 | cdef float[2] C 323 | cdef int nA 324 | cdef int nB 325 | 326 | theta = (theta_A + theta_B)/2 327 | xy(theta, a1, a2, e, C) 328 | dA = distance(A, C) 329 | dB = distance(C, B) 330 | nA = int((dA/(dA+dB))*(N-1)) 331 | nB = N - nA - 1 332 | 333 | _sample_superellipse_divide_conquer_inner( 334 | a1, a2, e, 335 | theta_A, theta, 336 | A, C, 337 | nA, 338 | thetas, 339 | idx 340 | ) 341 | thetas[idx[0]] = theta 342 | idx[0] += 1 343 | _sample_superellipse_divide_conquer_inner( 344 | a1, a2, e, 345 | theta, theta_B, 346 | C, B, 347 | nB, 348 | thetas, 349 | idx 350 | ) 351 | 352 | 353 | @cython.boundscheck(False) 354 | @cython.wraparound(False) 355 | cdef void sample_superellipse_divide_conquer( 356 | float a1, 357 | float a2, 358 | float e, 359 | float theta_A, 360 | float theta_B, 361 | float[::1] thetas 362 | ): 363 | # Decleare some useful variables 364 | cdef int N = thetas.shape[0] 365 | cdef float[2] A 366 | cdef float[2] B 367 | cdef int i = 0 # counter for the added elements in the list 368 | 369 | # Compute the two points on the superellipse 370 | xy(theta_A, a1, a2, e, A) 371 | xy(theta_B, a1, a2, e, B) 372 | 373 | # Add etas_init in the beginnint 374 | thetas[i] = theta_A 375 | i += 1 376 | _sample_superellipse_divide_conquer_inner( 377 | a1, a2, e, theta_A, theta_B, A, B, N-2, thetas, &i 378 | ) 379 | # Add etas_end in the end 380 | thetas[i] = theta_B 381 | 382 | 383 | def fast_sample( 384 | float a1, 385 | float a2, 386 | float a3, 387 | float e1, 388 | float e2, 389 | int N 390 | ): 391 | # Allocate memory for etas and omegas 392 | etas = np.empty((201,), dtype=np.float32) 393 | omegas = np.empty((201,), dtype=np.float32) 394 | sample_superellipse_divide_conquer( 395 | a1, a3, e1, M_PI_2, -M_PI_2, etas 396 | ) 397 | sample_superellipse_divide_conquer( 398 | a1, a2, e2, M_PI, -M_PI, omegas 399 | ) 400 | 401 | # Do the random sampling 402 | idxs = np.random.choice( 403 | etas.size*omegas.size, N, replace=False 404 | ) 405 | idxs_unraveled = np.unravel_index(idxs, (etas.size, omegas.size)) 406 | 407 | etas = etas[idxs_unraveled[0]] 408 | omegas = omegas[idxs_unraveled[1]] 409 | 410 | return etas, omegas 411 | 412 | 413 | @cython.boundscheck(False) 414 | @cython.wraparound(False) 415 | def fast_sample_on_batch( 416 | float[:, :, ::1] shapes, 417 | float[:, :, ::1] epsilons, 418 | int N 419 | ): 420 | # Declare some variables 421 | cdef int B = shapes.shape[0] 422 | cdef int M = shapes.shape[1] 423 | cdef int buffer_size = 201 424 | 425 | # Allocate memory for the etas and omegas 426 | cdef cnp.ndarray[cnp.float32_t, ndim=3] etas = \ 427 | np.zeros((B, M, N), dtype=np.float32) 428 | cdef cnp.ndarray[cnp.float32_t, ndim=3] omegas = \ 429 | np.zeros((B, M, N), dtype=np.float32) 430 | with nogil: 431 | sample_on_batch( 432 | &shapes[0, 0, 0], 433 | &epsilons[0, 0, 0], 434 | etas.data, 435 | omegas.data, 436 | B, M, N, 437 | buffer_size, 438 | 0 439 | ) 440 | 441 | return etas, omegas 442 | 443 | cdef extern from "sampling.hpp" nogil: 444 | void sample_on_batch( 445 | float *shapes, 446 | float *epsilons, 447 | float *etas, 448 | float *omegas, 449 | int B, 450 | int M, 451 | int N, 452 | int buffer_size, 453 | int seed 454 | ) 455 | -------------------------------------------------------------------------------- /learnable_primitives/fast_sampler/sampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | extern "C" { 10 | #include "sampling.hpp" 11 | } 12 | 13 | 14 | const float pi = std::acos(-1); 15 | const float pi_2 = pi/2; 16 | 17 | 18 | class prng { 19 | public: 20 | prng(int seed) : gen(seed), dis(0, 1) {} 21 | float operator()() { 22 | return dis(gen); 23 | } 24 | 25 | private: 26 | std::mt19937 gen; 27 | std::uniform_real_distribution dis; 28 | }; 29 | 30 | 31 | struct recursion_params { 32 | float A[2]; 33 | float B[2]; 34 | float theta_a; 35 | float theta_b; 36 | int N; 37 | int offset; 38 | 39 | recursion_params( 40 | float a[2], 41 | float b[2], 42 | float t_a, 43 | float t_b, 44 | int n, 45 | int o 46 | ) { 47 | A[0] = a[0]; 48 | A[1] = a[1]; 49 | B[0] = b[0]; 50 | B[1] = b[1]; 51 | theta_a = t_a; 52 | theta_b = t_b; 53 | N = n; 54 | offset = o; 55 | } 56 | }; 57 | 58 | 59 | inline float fexp(float x, float p) { 60 | return std::copysign(std::pow(std::abs(x), p), x); 61 | } 62 | 63 | 64 | inline void xy(float theta, float a1, float a2, float e, float C[2]) { 65 | C[0] = a1 * fexp(std::cos(theta), e); 66 | C[1] = a2 * fexp(std::sin(theta), e); 67 | } 68 | 69 | inline float distance(float A[2], float B[2]) { 70 | float d1 = A[0]-B[0]; 71 | float d2 = A[1]-B[1]; 72 | return std::sqrt(d1*d1 + d2*d2); 73 | } 74 | 75 | 76 | void sample_superellipse_divide_conquer( 77 | float a1, 78 | float a2, 79 | float e, 80 | float theta_a, 81 | float theta_b, 82 | std::vector &buffer, 83 | std::vector &stack 84 | ) { 85 | float A[2], B[2], C[2], theta, dA, dB; 86 | int nA, nB; 87 | 88 | xy(theta_a, a1, a2, e, A); 89 | xy(theta_b, a1, a2, e, B); 90 | buffer[0] = theta_a; 91 | stack.emplace_back(A, B, theta_a, theta_b, buffer.size()-2, 1); 92 | 93 | while (stack.size() > 0) { 94 | recursion_params params = stack.back(); 95 | stack.pop_back(); 96 | 97 | if (params.N <= 0) { 98 | continue; 99 | } 100 | 101 | theta = (params.theta_a + params.theta_b)/2; 102 | xy(theta, a1, a2, e, C); 103 | dA = distance(params.A, C); 104 | dB = distance(C, params.B); 105 | nA = static_cast(std::round((dA/(dA+dB))*(params.N-1))); 106 | nB = params.N - nA - 1; 107 | 108 | buffer[nA+params.offset] = theta; 109 | 110 | stack.emplace_back( 111 | params.A, C, 112 | params.theta_a, theta, 113 | nA, 114 | params.offset 115 | ); 116 | stack.emplace_back( 117 | C, params.B, 118 | theta, params.theta_b, 119 | nB, 120 | params.offset + nA + 1 121 | ); 122 | } 123 | 124 | buffer[buffer.size()-1] = theta_b; 125 | } 126 | 127 | 128 | void sample_etas( 129 | std::function rand, 130 | float a1a2, 131 | float e1, 132 | std::vector &buffer, 133 | std::vector &cdf, 134 | float *etas, 135 | int N 136 | ) { 137 | const float smoothing = 0.001; 138 | float s; 139 | 140 | // Make the sampling distribution's CDF 141 | cdf[0] = smoothing; 142 | for (unsigned int i=1; i buffer(buffer_size); 172 | std::vector eta_cdf(buffer_size); 173 | std::vector stack; 174 | 175 | for (int b=0; b(rand()*buffer_size)]; 212 | } 213 | } 214 | } 215 | } 216 | -------------------------------------------------------------------------------- /learnable_primitives/fast_sampler/sampling.hpp: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_HPP_ 2 | #define _SAMPLING_HPP_ 3 | 4 | 5 | void sample_on_batch( 6 | float *shapes, 7 | float *epsilons, 8 | float *etas, 9 | float *omegas, 10 | int B, 11 | int M, 12 | int N, 13 | int buffer_size, 14 | int seed 15 | ); 16 | 17 | 18 | #endif // _SAMPLING_HPP_ 19 | -------------------------------------------------------------------------------- /learnable_primitives/mesh.py: -------------------------------------------------------------------------------- 1 | """Create a Mesh object by parsing a file either in .ply format or .obj format 2 | """ 3 | import os 4 | import sys 5 | 6 | import numpy as np 7 | import trimesh 8 | import re 9 | 10 | from .pointcloud import PointcloudFromOBJ, Pointcloud, PointcloudFromOFF 11 | 12 | 13 | class Mesh(object): 14 | """A collection of vertices, normals and faces that construct a 3D mesh of 15 | an object. 16 | """ 17 | def __init__(self, points, normals, faces_idxs): 18 | self._normals = normals 19 | self._faces_idxs = faces_idxs 20 | self._face_normals = None 21 | 22 | # Pointcloud object that will hold the points 23 | self._pointcloud = Pointcloud(points) 24 | self._faces = None 25 | 26 | self._bbox = None 27 | 28 | @staticmethod 29 | def get_face_normals(faces): 30 | assert faces.shape[1] == 9 31 | A = faces[:, :3] 32 | B = faces[:, 3:6] 33 | C = faces[:, 6:] 34 | AB = B - A 35 | AC = C - A 36 | n = np.cross(AB, AC) 37 | n /= np.linalg.norm(n, axis=-1)[:, np.newaxis] 38 | # Make sure that n are unit vectors 39 | assert np.linalg.norm(n, axis=-1).sum() == n.shape[0] 40 | return n 41 | 42 | @staticmethod 43 | def get_face_normals_from_file(mesh_file): 44 | m = trimesh.load(mesh_file) 45 | # Make sure that the face orinetations are ok 46 | trimesh.repair.fix_normals(m, multibody=True) 47 | trimesh.repair.fix_winding(m) 48 | assert m.is_winding_consistent == True 49 | return m.face_normals 50 | 51 | @property 52 | def pointcloud(self): 53 | return self._pointcloud 54 | 55 | @property 56 | def points(self): 57 | return self.pointcloud.points 58 | 59 | @property 60 | def normals(self): 61 | return self._normals 62 | 63 | @property 64 | def faces(self): 65 | if self._faces is None: 66 | self._faces = np.array([np.hstack(( 67 | self.points.T[idx[0]], 68 | np.hstack((self.points.T[idx[1]], self.points.T[idx[2]])) 69 | )) for idx in self.faces_idxs.T 70 | ]) 71 | return self._faces 72 | 73 | @property 74 | def faces_idxs(self): 75 | return self._faces_idxs 76 | 77 | @property 78 | def face_normals(self): 79 | raise NotImplementedError() 80 | 81 | @property 82 | def bounding_box(self): 83 | if self._bbox is None: 84 | self._bbox = np.vstack([ 85 | self.points.min(-1), 86 | self.points.max(-1) 87 | ]) 88 | assert self._bbox.shape == (2, 3) 89 | return self._bbox 90 | 91 | def sample_points(self, N): 92 | return self.pointcloud.sample(N) 93 | 94 | def save_obj(self, file): 95 | with open(file, "w") as f: 96 | f.write("# OBJ file\n") 97 | for v in self.points.T: 98 | f.write("v %.4f %.4f %.4f\n" % tuple(v.tolist())) 99 | for fidx in self.faces_idxs.T: 100 | f.write("f") 101 | for i in fidx: 102 | f.write(" %d" % (i+1)) 103 | f.write("\n") 104 | 105 | def save_ply_as_ascii(self, file): 106 | # Number of points (vertices) 107 | N = self.points.shape[1] 108 | # Number of faces 109 | F = self.faces_idxs.shape[1] 110 | 111 | with open(file, "w") as f: 112 | f.write(("ply\nformat ascii 1.0\ncomment Raynet" 113 | " pointcloud!\nelement vertex %d\nproperty float x\n" 114 | "property float y\nproperty float z\nelement face %d\n" 115 | "property list uchar int vertex_indices\nend_header\n") 116 | % (N, F)) 117 | for p in self.points.T: 118 | f.write(" ".join(map(str, [p[0], p[1], p[2]])) + "\n") 119 | 120 | idxs = np.hstack([ 121 | 3*np.ones(self.faces_idxs.shape[1]).reshape(-1, 1), 122 | self.faces_idxs.T 123 | ]).astype(int) 124 | for p in idxs: 125 | f.write(" ".join(map(str, [p[0], p[1], p[2], p[3]])) + "\n") 126 | 127 | def save_ply(self, file): 128 | # Number of points (vertices) 129 | N = self.points.shape[1] 130 | # Number of faces 131 | F = self.faces_idxs.shape[1] 132 | 133 | with open(file, "w") as f: 134 | f.write(("ply\nformat binary_%s_endian 1.0\ncomment Raynet" 135 | " pointcloud!\nelement vertex %d\nproperty float x\n" 136 | "property float y\nproperty float z\nelement face %d\n" 137 | "property list uchar int vertex_index\nend_header\n") % ( 138 | sys.byteorder, N, F)) 139 | self.points.T.astype(np.float32).tofile(f) 140 | 141 | # TODO: Make this properly :-) 142 | for p in self.faces_idxs.T: 143 | np.array([3]).astype(np.uint8).tofile(f) 144 | np.array(p).astype(np.int32).tofile(f) 145 | 146 | def _sample_faces(self, N): 147 | # Change the shape of self.points and self.faces_idxs to match the 148 | # implementation 149 | vertices = self.points.T 150 | faces = self.faces_idxs.T 151 | 152 | n_faces = faces.shape[0] 153 | # Calculate all face areas, for total sum and cumsum sample uniformly 154 | # into the fractional size array, and then sample uniformly in that 155 | # triangle 156 | v0s = vertices[faces[:, 0], :] 157 | v1s = vertices[faces[:, 1], :] - v0s 158 | v2s = vertices[faces[:, 2], :] - v0s 159 | 160 | areas = np.power(np.sum(np.power(np.cross(v2s, v1s), 2), 1), 0.5) 161 | triangle_idxs = np.random.choice( 162 | len(areas), 163 | N, 164 | p=areas/np.sum(areas) 165 | ) 166 | 167 | v1_fracs = np.random.rand(N, 1) 168 | v2_fracs = np.random.rand(N, 1) 169 | frac_out = (v1_fracs + v2_fracs > 1) 170 | v1_fracs[frac_out] = 1 - v1_fracs[frac_out] 171 | v2_fracs[frac_out] = 1 - v2_fracs[frac_out] 172 | 173 | P = v0s[triangle_idxs, :] 174 | P += v1_fracs * v1s[triangle_idxs, :] 175 | P += v2_fracs * v2s[triangle_idxs, :] 176 | 177 | return P, triangle_idxs 178 | 179 | def sample_faces(self, N=10000): 180 | P, t = self._sample_faces(N) 181 | return np.hstack([P, self.face_normals[t, :]]) 182 | 183 | 184 | class MeshFromOBJ(Mesh): 185 | """Construct a Mesh Object from an OBJ file 186 | """ 187 | def __init__(self, obj_file): 188 | self.obj_file = obj_file 189 | # Raise Exception in case the given file does not exist 190 | if not os.path.exists(obj_file): 191 | raise IOException("File does not exist : %s" % (obj_file,)) 192 | 193 | self._normals = None 194 | self._faces = None 195 | self._faces_idxs = None 196 | self._face_normals = None 197 | 198 | # Pointcloud object that will hold the points 199 | self._pointcloud = None 200 | 201 | self._bbox = None 202 | 203 | @property 204 | def pointcloud(self): 205 | if self._pointcloud is None: 206 | self._pointcloud = PointcloudFromOBJ(self.obj_file) 207 | return self._pointcloud 208 | 209 | @staticmethod 210 | def parse_data(obj_file): 211 | # List to keep the unprocessed lines parsed from the file 212 | lines = [] 213 | with open(obj_file, "r") as f: 214 | lines = f.readlines() 215 | # Remove lines containing change of line 216 | lines = filter(None, [x.strip("\r\n") for x in lines]) 217 | 218 | # Keep only the lines that start with the letter v, that 219 | # correspond to vertices 220 | vertices = filter(lambda k: k.startswith("v "), lines) 221 | # Remove the unwanted "v" in front of every row and transform 222 | # it to float 223 | points = np.array([ 224 | map(float, k.strip().split(" ")[1:]) for k in vertices 225 | ]) 226 | 227 | # Keep only the lines that start with the bigramm vn, that 228 | # correspond to normals 229 | normals = filter(lambda k: k.startswith("vn"), lines) 230 | # Remove the unwanted "v" in front of every row and transform 231 | # it to float 232 | normals = np.array([ 233 | map(float, k.strip().split(" ")[1:]) for k in normals 234 | ]) 235 | 236 | # Keep only the lines that start with the letter f, that 237 | # correspond to faces 238 | f = filter(lambda k: k.startswith("f"), lines) 239 | # Remove all empty strings from the list and split it 240 | t = map( 241 | lambda x: filter(None, x), 242 | [k.strip().split(" ") for k in f] 243 | ) 244 | # Remove preficx "f" 245 | f_clean = map(lambda k: filter(lambda x: "f" not in x, k), t) 246 | 247 | faces_idxs = [] 248 | normals_idxs = [] 249 | for ff in f_clean: 250 | # Every row in f has the following format v1//vn1 v2//vn2 251 | # v3//vn3, where v* corresponds the vertex index while vn* 252 | # corresponds to the normal index. 253 | faces_idxs.append([re.split("/+", i)[0] for i in ff]) 254 | normals_idxs.append([re.split("/+", i)[-1] for i in ff]) 255 | faces_idxs = np.array([map(int, x) for x in faces_idxs]) 256 | normals_idxs = np.array([map(int, x) for x in normals_idxs]) 257 | # Remove 1 to make it compatible with the zero notation 258 | faces_idxs = faces_idxs - 1 259 | normals_idxs = normals_idxs - 1 260 | 261 | return points.T, normals.T, faces_idxs.T, normals_idxs.T 262 | 263 | @property 264 | def normals(self): 265 | if self._normals is None: 266 | _, self._normals, _, _ = MeshFromOBJ.parse_data(self.obj_file) 267 | return self._normals 268 | 269 | @property 270 | def faces_idxs(self): 271 | if self._faces_idxs is None: 272 | _, _, self._faces_idxs, _ = MeshFromOBJ.parse_data(self.obj_file) 273 | return self._faces_idxs 274 | 275 | @property 276 | def face_normals(self): 277 | if self._face_normals is None: 278 | self._face_normals = Mesh.get_face_normals_from_file(self.obj_file) 279 | return self._face_normals 280 | 281 | 282 | class MeshFromOFF(Mesh): 283 | """Construct a Mesh Object from an OFF file 284 | """ 285 | def __init__(self, off_file): 286 | self.off_file = off_file 287 | 288 | # Raise Exception in case the given file does not exist 289 | if not os.path.exists(off_file): 290 | raise IOException("File does not exist : %s" % (off_file,)) 291 | self._normals = None 292 | self._faces = None 293 | self._faces_idxs = None 294 | 295 | # Pointcloud object that will hold the points 296 | self._pointcloud = None 297 | 298 | self._bbox = None 299 | 300 | @property 301 | def pointcloud(self): 302 | if self._pointcloud is None: 303 | self._pointcloud = PointcloudFromOFF(self.off_file) 304 | return self._pointcloud 305 | 306 | @staticmethod 307 | def parse_data(off_file): 308 | with open(off_file, "r") as f: 309 | # Read the lines and drop the first to remove the "OFF" 310 | lines = f.readlines()[1:] 311 | # Parse the number of vertices and the number of faces from the 312 | # first line 313 | n_vertices, n_faces, _ = map(int, lines[0].strip().split()) 314 | vertices = lines[1:n_vertices+1] 315 | assert len(vertices) == n_vertices 316 | points = np.array([ 317 | map(float, vi) 318 | for vi in 319 | [v.strip().split() for v in vertices] 320 | ]) 321 | 322 | faces_idxs = lines[n_vertices+1:] 323 | assert len(faces_idxs) == n_faces 324 | faces_idxs = np.array([ 325 | map(int, vi) 326 | for vi in 327 | [v.strip().split()[1:] for v in faces_idxs] 328 | ]) 329 | 330 | return points.T, None, faces_idxs.T 331 | 332 | @property 333 | def normals(self): 334 | if self._normals is None: 335 | _, self._normals, _ = MeshFromOFF.parse_data(self.off_file) 336 | return self._normals 337 | 338 | @property 339 | def faces_idxs(self): 340 | if self._faces_idxs is None: 341 | _, _, self._faces_idxs = MeshFromOFF.parse_data(self.off_file) 342 | return self._faces_idxs 343 | 344 | 345 | class Trimesh(Mesh): 346 | "Wrapper when using the trimesh library" 347 | def __init__(self, mesh_file): 348 | self.mesh_file = mesh_file 349 | # Raise Exception in case the given file does not exist 350 | if not os.path.exists(mesh_file): 351 | raise ValueError("File does not exist : %s" % (mesh_file,)) 352 | self._mesh = None 353 | 354 | self._normals = None 355 | self._faces = None 356 | self._faces_idxs = None 357 | self._face_normals = None 358 | 359 | # Pointcloud object that will hold the points 360 | self._pointcloud = None 361 | self._points = None 362 | 363 | self._bbox = None 364 | 365 | def _normalize_points(self): 366 | """Make sure that points lie in the unit cube.""" 367 | points = self.mesh.vertices.T 368 | mins = np.min(points, axis=1, keepdims=True) 369 | steps = np.max(points, axis=1, keepdims=True) - mins 370 | points -= mins + steps/2 371 | if steps.max() > 1: 372 | points /= steps.max() 373 | 374 | @property 375 | def mesh(self): 376 | if self._mesh is None: 377 | self._mesh = trimesh.load(self.mesh_file) 378 | # Make sure that the face orinetations are ok 379 | # trimesh.repair.fix_normals(self.mesh, multibody=True) 380 | # trimesh.repair.fix_winding(self.mesh) 381 | # assert self.mesh.is_winding_consistent == True 382 | # Normalize the points to be in the unit cube 383 | self._normalize_points() 384 | 385 | return self._mesh 386 | 387 | def contains(self, points): 388 | return self.mesh.contains(points) 389 | 390 | @property 391 | def faces_idxs(self): 392 | if self._faces_idxs is None: 393 | self._faces_idxs = self.mesh.faces.T 394 | return self._faces_idxs 395 | 396 | @property 397 | def face_normals(self): 398 | if self._face_normals is None: 399 | self._face_normals = self.mesh.face_normals 400 | return self._face_normals 401 | 402 | @property 403 | def points(self): 404 | if self._points is None: 405 | self._points = self.mesh.vertices.T 406 | 407 | return self._points 408 | 409 | @property 410 | def normals(self): 411 | if self._normals is None: 412 | self._normals = self.mesh.vertex_normals.T 413 | return self._normals 414 | 415 | @property 416 | def pointcloud(self): 417 | if self._pointcloud is None: 418 | self._pointcloud = Pointcloud(self.points) 419 | return self._pointcloud 420 | 421 | @property 422 | def bounding_box(self): 423 | if self._bbox is None: 424 | # A numpy array of size 2x3 containing the bbox that contains the 425 | # mesh 426 | self._bbox = self.mesh.bounds 427 | return self._bbox 428 | 429 | def sample_faces(self, N=10000): 430 | P, t = trimesh.sample.sample_surface(self.mesh, N) 431 | return np.hstack([ 432 | P, self.face_normals[t, :] 433 | ]) 434 | 435 | 436 | class MeshFactory(object): 437 | """Static factory methods collected under the MeshFactory namespace.""" 438 | @staticmethod 439 | def from_file(filepath): 440 | return Trimesh(filepath) 441 | -------------------------------------------------------------------------------- /learnable_primitives/models.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | from torchvision import models 7 | 8 | 9 | class NetworkParameters(object): 10 | def __init__(self, architecture, n_primitives=32, 11 | mu=0.0, sigma=0.001, add_gaussian_noise=False, 12 | use_sq=False, make_dense=False, 13 | use_deformations=False, 14 | train_with_bernoulli=False): 15 | self.architecture = architecture 16 | self.n_primitives = n_primitives 17 | self.train_with_bernoulli = train_with_bernoulli 18 | self.add_gaussian_noise = add_gaussian_noise 19 | self.gaussian_noise_layer = get_gaussian_noise_layer( 20 | self.add_gaussian_noise, 21 | mu, 22 | sigma 23 | ) 24 | self.use_sq = use_sq 25 | self.use_deformations = use_deformations 26 | self.make_dense = make_dense 27 | 28 | @classmethod 29 | def from_options(cls, argument_parser): 30 | # Make Namespace to dictionary to be able to use it 31 | args = vars(argument_parser) 32 | 33 | architecture = args["architecture"] 34 | n_primitives = args.get("n_primitives", 32) 35 | 36 | add_gaussian_noise = args.get("add_gaussian_noise", False) 37 | mu = args.get("mu", 0.0) 38 | sigma = args.get("sigma", 0.001) 39 | 40 | # By default train without learning Bernoulli priors 41 | train_with_bernoulli = args.get("train_with_bernoulli", False) 42 | use_sq = args.get("use_sq", False) 43 | use_deformations = args.get("use_deformations", False) 44 | make_dense = args.get("make_dense", False) 45 | 46 | return cls( 47 | architecture, 48 | n_primitives=n_primitives, 49 | mu=mu, 50 | sigma=sigma, 51 | add_gaussian_noise=add_gaussian_noise, 52 | use_sq=use_sq, 53 | use_deformations=use_deformations, 54 | train_with_bernoulli=train_with_bernoulli, 55 | make_dense=make_dense 56 | ) 57 | 58 | @property 59 | def network(self): 60 | networks = dict( 61 | tulsiani=TulsianiNetwork, 62 | octnet=OctnetNetwork, 63 | resnet18=ResNet18 64 | ) 65 | 66 | return networks[self.architecture.lower()] 67 | 68 | def primitive_layer(self, n_primitives, input_channels): 69 | modules = self._build_modules(n_primitives, input_channels) 70 | module = GeometricPrimitive(n_primitives, modules) 71 | return module 72 | 73 | def _build_modules(self, n_primitives, input_channels): 74 | modules = { 75 | "translations": Translation(n_primitives, input_channels, self.make_dense), 76 | "rotations": Rotation(n_primitives, input_channels, self.make_dense), 77 | "sizes": Size(n_primitives, input_channels, self.make_dense) 78 | } 79 | if self.train_with_bernoulli: 80 | modules["probs"] = Probability(n_primitives, input_channels, self.make_dense) 81 | if self.use_sq and not self.use_deformations: 82 | modules["shapes"] = Shape(n_primitives, input_channels, self.make_dense) 83 | if self.use_sq and self.use_deformations: 84 | modules["shapes"] = Shape(n_primitives, input_channels, self.make_dense) 85 | modules["deformations"] = Deformation( 86 | n_primitives, input_channels, self.make_dense) 87 | 88 | return modules 89 | 90 | 91 | class TulsianiNetwork(nn.Module): 92 | def __init__(self, network_params): 93 | super(TulsianiNetwork, self).__init__() 94 | self._network_params = network_params 95 | 96 | # Initialize some useful variables 97 | n_filters = 4 98 | input_channels = 1 99 | 100 | encoder_layers = [] 101 | # Create an encoder using a stack of convolutions 102 | for i in range(5): 103 | encoder_layers.append( 104 | nn.Conv3d(input_channels, n_filters, kernel_size=3, padding=1) 105 | ) 106 | encoder_layers.append(nn.BatchNorm3d(n_filters)) 107 | encoder_layers.append(nn.LeakyReLU(0.2, True)) 108 | encoder_layers.append(nn.MaxPool3d(kernel_size=2, stride=2)) 109 | 110 | input_channels = n_filters 111 | # Double the number of filters after every layer 112 | n_filters *= 2 113 | 114 | # Add the two fully connected layers 115 | input_channels = n_filters / 2 116 | n_filters = 100 117 | for i in range(2): 118 | encoder_layers.append(nn.Conv3d(input_channels, n_filters, 1)) 119 | # encoder_layers.append(nn.BatchNorm3d(n_filters)) 120 | encoder_layers.append(nn.LeakyReLU(0.2, True)) 121 | 122 | input_channels = n_filters 123 | 124 | self._features_extractor = nn.Sequential(*encoder_layers) 125 | self._primitive_layer = self._network_params.primitive_layer( 126 | self._network_params.n_primitives, 127 | n_filters 128 | ) 129 | 130 | def forward(self, X): 131 | x = self._features_extractor(X) 132 | return self._primitive_layer(x) 133 | 134 | 135 | class OctnetNetwork(nn.Module): 136 | def __init__(self, network_params): 137 | super(OctnetNetwork, self).__init__() 138 | 139 | self.encoder_conv = nn.Sequential( 140 | nn.Conv3d(1, 8, kernel_size=3, padding=1), nn.ReLU(), 141 | nn.Conv3d(8, 8, kernel_size=3, padding=1), nn.ReLU(), 142 | nn.Conv3d(8, 8, kernel_size=3, padding=1), nn.ReLU(), 143 | nn.MaxPool3d(kernel_size=2, stride=2), 144 | 145 | nn.Conv3d(8, 16, kernel_size=3, padding=1), nn.ReLU(), 146 | nn.Conv3d(16, 16, kernel_size=3, padding=1), nn.ReLU(), 147 | nn.Conv3d(16, 16, kernel_size=3, padding=1), nn.ReLU(), 148 | nn.MaxPool3d(kernel_size=2, stride=2), 149 | 150 | nn.Conv3d(16, 32, kernel_size=3, padding=1), nn.ReLU(), 151 | nn.Conv3d(32, 32, kernel_size=3, padding=1), nn.ReLU(), 152 | nn.Conv3d(32, 32, kernel_size=3, padding=1), nn.ReLU(), 153 | nn.MaxPool3d(kernel_size=2, stride=2), 154 | 155 | nn.Conv3d(32, 64, kernel_size=3, padding=1), nn.ReLU(), 156 | nn.Conv3d(64, 64, kernel_size=3, padding=1), nn.ReLU(), 157 | nn.Conv3d(64, 64, kernel_size=3, padding=1), nn.ReLU(), 158 | nn.MaxPool3d(kernel_size=2, stride=2), 159 | ) 160 | self.encoder_fc = nn.Sequential( 161 | nn.Linear(2*2*2*64, 1024), nn.ReLU(), 162 | # nn.Dropout(0.5), 163 | nn.Linear(1024, 1024), nn.ReLU() 164 | ) 165 | self.primitive_layer = network_params.primitive_layer( 166 | network_params.n_primitives, 167 | 1024 168 | ) 169 | 170 | def forward(self, X): 171 | X = self.encoder_conv(X) 172 | X = self.encoder_fc(X.view(-1, 2*2*2*64)) 173 | return self.primitive_layer(X.view(-1, 1024, 1, 1, 1)) 174 | 175 | 176 | class ResNet18(nn.Module): 177 | def __init__(self, network_params): 178 | super(ResNet18, self).__init__() 179 | self._network_params = network_params 180 | 181 | self._features_extractor = models.resnet18(pretrained=True) 182 | self._features_extractor.fc = nn.Sequential( 183 | nn.Linear(512, 512), nn.ReLU(), 184 | nn.Linear(512, 512), nn.ReLU() 185 | ) 186 | self._features_extractor.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 187 | 188 | self._primitive_layer = self._network_params.primitive_layer( 189 | self._network_params.n_primitives, 190 | 512 191 | ) 192 | 193 | def forward(self, X): 194 | X = X.float() / 255.0 195 | x = self._features_extractor(X) 196 | return self._primitive_layer(x.view(-1, 512, 1, 1, 1)) 197 | 198 | 199 | class Translation(nn.Module): 200 | """A layer that predicts the translation vector 201 | """ 202 | def __init__(self, n_primitives, input_channels, make_dense=False): 203 | super(Translation, self).__init__() 204 | self._n_primitives = n_primitives 205 | 206 | self._make_dense = make_dense 207 | if self._make_dense: 208 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 209 | self._nonlin = nn.LeakyReLU(0.2, True) 210 | 211 | # Layer used to infer the translation vector of each primitive, namely 212 | # BxMx3 213 | self._translation_layer = nn.Conv3d( 214 | input_channels, self._n_primitives*3, 1 215 | ) 216 | 217 | def forward(self, X): 218 | if self._make_dense: 219 | X = self._nonlin(self._fc(X)) 220 | 221 | # Compute the BxM*3 translation vectors for every primitive and ensure 222 | # that they lie inside the unit cube 223 | translations = torch.tanh(self._translation_layer(X)) * 0.51 224 | 225 | return translations[:, :, 0, 0, 0] 226 | 227 | 228 | class Rotation(nn.Module): 229 | """A layer that predicts the rotation vector 230 | """ 231 | def __init__(self, n_primitives, input_channels, make_dense=False): 232 | super(Rotation, self).__init__() 233 | self._n_primitives = n_primitives 234 | 235 | self._make_dense = make_dense 236 | if self._make_dense: 237 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 238 | self._nonlin = nn.LeakyReLU(0.2, True) 239 | 240 | # Layer used to infer the 4 quaternions of each primitive, namely 241 | # BxMx4 242 | self._rotation_layer = nn.Conv3d( 243 | input_channels, self._n_primitives*4, 1 244 | ) 245 | 246 | def forward(self, X): 247 | if self._make_dense: 248 | X = self._nonlin(self._fc(X)) 249 | 250 | # Compute the 4 parameters of the quaternion for every primitive 251 | # and add a non-linearity as L2-normalization to enforce the unit 252 | # norm constrain 253 | quats = self._rotation_layer(X)[:, :, 0, 0, 0] 254 | quats = quats.view(-1, self._n_primitives, 4) 255 | rotations = quats / torch.norm(quats, 2, -1, keepdim=True) 256 | rotations = rotations.view(-1, self._n_primitives*4) 257 | 258 | return rotations 259 | 260 | 261 | class Size(nn.Module): 262 | """A layer that predicts the size vector 263 | """ 264 | def __init__(self, n_primitives, input_channels, make_dense=False): 265 | super(Size, self).__init__() 266 | self._n_primitives = n_primitives 267 | 268 | self._make_dense = make_dense 269 | if self._make_dense: 270 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 271 | self._nonlin = nn.LeakyReLU(0.2, True) 272 | 273 | # Layer used to infer the size of each primitive, along each axis, 274 | # namely BxMx3. 275 | self._size_layer = nn.Conv3d( 276 | input_channels, self._n_primitives*3, 1 277 | ) 278 | 279 | def forward(self, X): 280 | if self._make_dense: 281 | X = self._nonlin(self._fc(X)) 282 | 283 | # Bound the sizes so that they won't take values larger than 0.51 and 284 | # smaller than 1e-2 (to avoid numerical instabilities with the 285 | # inside-outside function) 286 | sizes = torch.sigmoid(self._size_layer(X)) * 0.5 + 0.03 287 | sizes = sizes[:, :, 0, 0, 0] 288 | 289 | return sizes 290 | 291 | 292 | class Shape(nn.Module): 293 | """A layer that predicts the shape vector 294 | """ 295 | def __init__(self, n_primitives, input_channels, make_dense=False): 296 | super(Shape, self).__init__() 297 | self._n_primitives = n_primitives 298 | 299 | self._make_dense = make_dense 300 | if self._make_dense: 301 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 302 | self._nonlin = nn.LeakyReLU(0.2, True) 303 | 304 | # Layer used to infer the shape of each primitive, along each axis, 305 | # namely BxMx3. 306 | self._shape_layer = nn.Conv3d( 307 | input_channels, self._n_primitives*2, 1 308 | ) 309 | 310 | def forward(self, X): 311 | if self._make_dense: 312 | X = self._nonlin(self._fc(X)) 313 | 314 | # Bound the predicted shapes to avoid numerical instabilities with 315 | # the inside-outside function 316 | shapes = torch.sigmoid(self._shape_layer(X))*1.1 + 0.4 317 | shapes = shapes[:, :, 0, 0, 0] 318 | 319 | return shapes 320 | 321 | 322 | class Deformation(nn.Module): 323 | """A layer that predicts the deformations 324 | """ 325 | def __init__(self, n_primitives, input_channels, make_dense=False): 326 | super(Deformation, self).__init__() 327 | self._n_primitives = n_primitives 328 | 329 | self._make_dense = make_dense 330 | if self._make_dense: 331 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 332 | self._nonlin = nn.LeakyReLU(0.2, True) 333 | 334 | # Layer used to infer the tapering parameters of each primitive. 335 | self._tapering_layer =\ 336 | nn.Conv3d(input_channels, self._n_primitives*2, 1) 337 | 338 | def forward(self, X): 339 | if self._make_dense: 340 | X = self._nonlin(self._fc(X)) 341 | 342 | # The tapering parameters are from -1 to 1 343 | taperings = torch.tanh(self._tapering_layer(X))*0.9 344 | taperings = taperings[:, :, 0, 0, 0] 345 | 346 | return taperings 347 | 348 | 349 | class Probability(nn.Module): 350 | """A layer that predicts the probabilities 351 | """ 352 | def __init__(self, n_primitives, input_channels, make_dense=False): 353 | super(Probability, self).__init__() 354 | self._n_primitives = n_primitives 355 | 356 | self._make_dense = make_dense 357 | if self._make_dense: 358 | self._fc = nn.Conv3d(input_channels, input_channels, 1) 359 | self._nonlin = nn.LeakyReLU(0.2, True) 360 | 361 | # Layer used to infer the probability of existence for the M 362 | # primitives, namely BxM numbers, where B is the batch size 363 | self._probability_layer = nn.Conv3d( 364 | input_channels, self._n_primitives, 1 365 | ) 366 | 367 | def forward(self, X): 368 | if self._make_dense: 369 | X = self._nonlin(self._fc(X)) 370 | 371 | # Compute the BxM probabilities of existence for the M primitives and 372 | # remove unwanted axis with size 1 373 | probs = torch.sigmoid( 374 | self._probability_layer(X) 375 | ).view(-1, self._n_primitives) 376 | 377 | return probs 378 | 379 | 380 | class PrimitiveParameters(object): 381 | """Represents the \lambda_m.""" 382 | def __init__(self, probs, translations, rotations, sizes, shapes, 383 | deformations): 384 | self.probs = probs 385 | self.translations = translations 386 | self.rotations = rotations 387 | self.sizes = sizes 388 | self.shapes = shapes 389 | self.deformations = deformations 390 | 391 | # Check that everything has a len(shape) > 1 392 | for x in self.members[:-2]: 393 | assert len(x.shape) > 1 394 | 395 | def __getattr__(self, name): 396 | if not name.endswith("_r"): 397 | raise AttributeError() 398 | 399 | prop = getattr(self, name[:-2]) 400 | if not torch.is_tensor(prop): 401 | raise AttributeError() 402 | 403 | return prop.view(self.batch_size, self.n_primitives, -1) 404 | 405 | @property 406 | def members(self): 407 | return ( 408 | self.probs, 409 | self.translations, 410 | self.rotations, 411 | self.sizes, 412 | self.shapes, 413 | self.deformations 414 | ) 415 | 416 | @property 417 | def batch_size(self): 418 | return self.probs.shape[0] 419 | 420 | @property 421 | def n_primitives(self): 422 | return self.probs.shape[1] 423 | 424 | def __len__(self): 425 | return len(self.members) 426 | 427 | def __getitem__(self, i): 428 | return self.members[i] 429 | 430 | 431 | class GeometricPrimitive(nn.Module): 432 | def __init__(self, n_primitives, primitive_params): 433 | super(GeometricPrimitive, self).__init__() 434 | self._n_primitives = n_primitives 435 | self._primitive_params = primitive_params 436 | 437 | self._update_params() 438 | 439 | def _update_params(self): 440 | for i, m in enumerate(self._primitive_params.values()): 441 | self.add_module("layer%d" % (i,), m) 442 | 443 | def forward(self, X): 444 | if "probs" not in self._primitive_params.keys(): 445 | probs = X.new_ones((X.shape[0], self._n_primitives)) 446 | else: 447 | probs = self._primitive_params["probs"].forward(X) 448 | 449 | translations = self._primitive_params["translations"].forward(X) 450 | rotations = self._primitive_params["rotations"].forward(X) 451 | sizes = self._primitive_params["sizes"].forward(X) 452 | 453 | # By default the geometric primitive is a cuboid 454 | if "shapes" not in self._primitive_params.keys(): 455 | shapes = X.new_ones((X.shape[0], self._n_primitives*2)) * 0.25 456 | else: 457 | shapes = self._primitive_params["shapes"].forward(X) 458 | 459 | if "deformations" not in self._primitive_params.keys(): 460 | deformations = X.new_zeros((X.shape[0], self._n_primitives*2)) 461 | else: 462 | deformations = self._primitive_params["deformations"].forward(X) 463 | 464 | return PrimitiveParameters( 465 | probs, translations, rotations, sizes, 466 | shapes, deformations 467 | ) 468 | 469 | 470 | class GaussianNoise(nn.Module): 471 | def __init__(self, mu=0.0, sigma=0.01): 472 | super(GaussianNoise, self).__init__() 473 | # Mean of the distribution 474 | self.mu = mu 475 | # Standard deviation of the distribution 476 | self.sigma = sigma 477 | 478 | def forward(self, X): 479 | if self.training and self.sigma != 0: 480 | n = X.new_zeros(*X.size()).normal_(self.mu, self.sigma) 481 | X = X + n 482 | return X 483 | 484 | 485 | def train_on_batch( 486 | model, 487 | optimizer, 488 | loss_fn, 489 | X, 490 | y_target, 491 | regularizer_terms, 492 | sq_sampler, 493 | loss_options 494 | ): 495 | # Zero the gradient's buffer 496 | optimizer.zero_grad() 497 | y_hat = model(X) 498 | loss, debug_stats = loss_fn( 499 | y_hat, 500 | y_target, 501 | regularizer_terms, 502 | sq_sampler, 503 | loss_options 504 | ) 505 | # Do the backpropagation 506 | loss.backward() 507 | nn.utils.clip_grad_norm_(model.parameters(), 1) 508 | # Do the update 509 | optimizer.step() 510 | 511 | return ( 512 | loss.item(), 513 | [x.data if hasattr(x, "data") else x for x in y_hat], 514 | debug_stats 515 | ) 516 | 517 | 518 | def get_gaussian_noise_layer(add_gaussian_noise, mu=0.0, sigma=0.01): 519 | if add_gaussian_noise: 520 | return GaussianNoise(mu=mu, sigma=sigma) 521 | else: 522 | return GaussianNoise(mu=0.0, sigma=0.0) 523 | 524 | 525 | def optimizer_factory(args, model): 526 | """Based on the input arguments create a suitable optimizer object 527 | """ 528 | if args.probs_only: 529 | params = model._primitive_layer._primitive_params["probs"].parameters() 530 | else: 531 | params = model.parameters() 532 | 533 | if args.optimizer == "SGD": 534 | return optim.SGD( 535 | params, 536 | lr=args.lr, 537 | momentum=args.momentum 538 | ) 539 | elif args.optimizer == "Adam": 540 | return optim.Adam( 541 | params, 542 | lr=args.lr 543 | ) 544 | -------------------------------------------------------------------------------- /learnable_primitives/pointcloud.py: -------------------------------------------------------------------------------- 1 | """Create a PointCloud object by parsing a point cloud from either a .ply file 2 | or an .obj file. 3 | """ 4 | 5 | import sys 6 | 7 | import numpy as np 8 | from sklearn.neighbors import KDTree 9 | 10 | from matplotlib.cm import get_cmap 11 | 12 | 13 | class PLYHeader(object): 14 | """Parse a PLY file header into an object""" 15 | class Element(object): 16 | def __init__(self, name, count, properties): 17 | assert len(properties) > 0 18 | self.name = name 19 | self.count = count 20 | self.properties = properties 21 | 22 | @property 23 | def bytes(self): 24 | return sum(p.bytes for p in self.properties) 25 | 26 | class Property(object): 27 | def __init__(self, name, type): 28 | self.name = name 29 | self.type = type 30 | 31 | @property 32 | def bytes(self): 33 | return { 34 | "float": 4, 35 | "uchar": 1, 36 | "int": 4 37 | }[self.type] 38 | 39 | def __init__(self, fileobj): 40 | assert fileobj.readline().strip() == "ply" 41 | 42 | lines = [] 43 | while True: 44 | l = fileobj.readline() 45 | if "end_header" in l: 46 | break 47 | lines.append(l) 48 | 49 | # Version and format 50 | identifier, format, version = lines[0].split() 51 | assert identifier == "format" 52 | self.is_ascii = "ascii" in format 53 | self.version = float(version) 54 | self.little_endian = "little" in format 55 | lines.pop(0) 56 | 57 | # Comments 58 | self.comments = [ 59 | x.split(" ", 1)[1] 60 | for x in lines 61 | if x.startswith("comment") 62 | ] 63 | 64 | # Elements 65 | lines = [l for l in lines if not l.startswith("comment")] 66 | elements = [] 67 | while lines: 68 | identifier, name, count = lines[0].split() 69 | assert identifier == "element" 70 | count = int(count) 71 | lines.pop(0) 72 | 73 | properties = [] 74 | while lines: 75 | identifier, type, name = lines[0].split() 76 | if identifier != "property": 77 | break 78 | properties.append(self.Property(name, type)) 79 | lines.pop(0) 80 | elements.append(self.Element(name, count, properties)) 81 | self.elements = elements 82 | 83 | 84 | class Pointcloud(object): 85 | """A collection of ND (usually 3D) points that can be searched over and 86 | saved into a file.""" 87 | def __init__(self, points): 88 | assert points.shape[0] == 3 89 | 90 | self._points = points 91 | self._normalize() 92 | 93 | def _normalize(self): 94 | """Normalize the points so that they are in the unit cube.""" 95 | points = self._points 96 | mins = np.min(points, axis=1, keepdims=True) 97 | steps = np.max(points, axis=1, keepdims=True) - mins 98 | points -= mins + steps/2 99 | if steps.max() > 1: 100 | points /= steps.max() 101 | 102 | @property 103 | def points(self): 104 | return self._points 105 | 106 | def sample(self, N): 107 | return self.points[ 108 | :, 109 | np.random.choice(np.arange(self.points.shape[1]), N) 110 | ] 111 | 112 | def _add_header(self, N): 113 | return [ 114 | "ply", 115 | "format binary_%s_endian 1.0" % (sys.byteorder,), 116 | "comment Raynet pointcloud!", 117 | "element vertex %d" % (N,), 118 | "property float x", 119 | "property float y", 120 | "property float z", 121 | "property uchar red", 122 | "property uchar green", 123 | "property uchar blue", 124 | "end_header\n" 125 | ] 126 | 127 | def save_obj(self, file): 128 | with open(file, "w") as f: 129 | f.write("# OBJ file\n") 130 | for v in self.points.T: 131 | f.write("v %.4f %.4f %.4f\n" % tuple(v.tolist())) 132 | 133 | def save_ply(self, file): 134 | N = self.points.shape[1] 135 | with open(file, "w") as f: 136 | header = self._add_header(N) 137 | f.write("\n".join(header[:7] + header[-1:])) 138 | self.points.T.astype(np.float32).tofile(f) 139 | 140 | def save_colored_ply(self, file, intensities, colormap="jet"): 141 | # Get the colormap based on the input 142 | cmap = get_cmap(colormap) 143 | # Based on the selected colormap get the the colors for every point 144 | intensities = intensities / 2 145 | colors = cmap(intensities.ravel())[:, :-1] 146 | # The color values need to be uchar 147 | colors = (colors * 255).astype(np.uint8) 148 | 149 | N = self.points.shape[1] 150 | # idxs = np.arange(N)[intensities.ravel() < 1.0] 151 | idxs = np.arange(N) 152 | with open(file, "w") as f: 153 | f.write("\n".join(self._add_header(len(idxs)))) 154 | cnt = 0 155 | # for point, color in zip(self.points.T, colors): 156 | for i in idxs: 157 | point = self.points.T[i] 158 | color = colors[i] 159 | point.astype(np.float32).tofile(f) 160 | color.tofile(f) 161 | cnt += 1 162 | 163 | def save(self, file): 164 | np.save(file, self.points) 165 | 166 | def filter(self, mask): 167 | self._points = mask.filter(self.points) 168 | 169 | def index(self, leaf_size=40, metric="minkowski"): 170 | if hasattr(self, "_index"): 171 | return 172 | 173 | # NOTE: scikit-learn expects points (samples, features) while we use 174 | # the more traditional (features, samples) 175 | self._index = KDTree(self.points.T, leaf_size, metric) 176 | 177 | def nearest_neighbors(self, X, k=1, return_distances=True): 178 | return self._index.query(X.T, k, return_distances) 179 | 180 | 181 | class PointcloudFromPLY(Pointcloud): 182 | """Create a point cloud from a .PLY file 183 | """ 184 | def __init__(self, ply_file): 185 | self.ply_file = ply_file 186 | self._points = None 187 | 188 | @property 189 | def points(self): 190 | if self._points is None: 191 | with open(self.ply_file, "rb") as f: 192 | header = PLYHeader(f) 193 | assert len(header.elements) == 1 194 | el = header.elements[0] 195 | assert all(p.type == "float" for p in el.properties[:3]) 196 | 197 | # Read the data and place one element per line and skip all the 198 | # extra elements 199 | data = np.fromfile(f, dtype=np.uint8) 200 | data = data.reshape(-1, header.elements[0].bytes) 201 | data =\ 202 | data[:, :sum(p.bytes for p in el.properties[:3])].ravel() 203 | 204 | # Reread in the correct byte-order 205 | order = "<" if header.little_endian else ">" 206 | dtype = order + "f4" 207 | self._points = np.frombuffer(data.data, dtype=dtype).T 208 | self._normalize() 209 | 210 | return self._points 211 | 212 | 213 | class PointcloudFromOBJ(Pointcloud): 214 | """Create a point cloud from a .OBJ file 215 | """ 216 | def __init__(self, obj_file): 217 | self.obj_file = obj_file 218 | self._points = None 219 | 220 | @property 221 | def points(self): 222 | if self._points is None: 223 | # List to keep the unprocessed lines parsed from the file 224 | lines = [] 225 | with open(self.obj_file, "r") as f: 226 | lines = f.readlines() 227 | # Remove lines containing change of line 228 | lines = filter(None, [x.strip("\r\n") for x in lines]) 229 | 230 | # Keep only the lines that start with the letter v, that 231 | # corresponds to vertices 232 | vertices = filter(lambda k: k.startswith("v "), lines) 233 | # Remove the unwanted "v" in front of every row and transform 234 | # it to float 235 | self._points = np.array([ 236 | map(float, k.strip().split(" ")[1:]) for k in vertices 237 | ]).T 238 | self._normalize() 239 | 240 | return self._points 241 | 242 | 243 | class PointcloudFromOFF(Pointcloud): 244 | """Create a point cloud from a .OBJ file 245 | """ 246 | def __init__(self, off_file): 247 | self.off_file = off_file 248 | self._points = None 249 | 250 | @property 251 | def points(self): 252 | if self._points is None: 253 | with open(self.off_file, "r") as fp: 254 | # Read the lines and drop the first to remove the "OFF" 255 | lines = fp.readlines()[1:] 256 | # Parse the number of vertices and the number of faces from the 257 | # first line 258 | n_vertices, n_faces, _ = map(int, lines[0].strip().split()) 259 | vertices = lines[1:n_vertices+1] 260 | assert len(vertices) == n_vertices 261 | self._points = np.array([ 262 | map(float, vi) 263 | for vi in 264 | [v.strip().split() for v in vertices] 265 | ]).T 266 | self._normalize() 267 | 268 | return self._points 269 | -------------------------------------------------------------------------------- /learnable_primitives/regularizers.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import torch 4 | 5 | from .primitives import inside_outside_function, sq_volumes, \ 6 | transform_to_primitives_centric_system, quaternions_to_rotation_matrices, \ 7 | sample_points_inside_primitives 8 | 9 | 10 | def bernoulli(parameters, minimum_number_of_primitives, 11 | maximum_number_of_primitives, w1, w2): 12 | """Ensure that we have at least that many primitives in expectation""" 13 | expected_primitives = parameters[0].sum(-1) 14 | 15 | lower_bound = minimum_number_of_primitives - expected_primitives 16 | upper_bound = expected_primitives - maximum_number_of_primitives 17 | zero = expected_primitives.new_tensor(0) 18 | 19 | t1 = torch.max(lower_bound, zero) 20 | t2 = torch.max(upper_bound, zero) 21 | 22 | return (w1*t1 + w2*t2).mean() 23 | 24 | 25 | def sparsity(parameters, minimum_number_of_primitives, 26 | maximum_number_of_primitives, w1, w2): 27 | expected_primitives = parameters[0].sum(-1) 28 | lower_bound = minimum_number_of_primitives - expected_primitives 29 | upper_bound = expected_primitives - maximum_number_of_primitives 30 | zero = expected_primitives.new_tensor(0) 31 | 32 | t1 = torch.max(lower_bound, zero) * lower_bound**4 33 | t2 = torch.max(upper_bound, zero) * upper_bound**2 34 | 35 | return (w1*t1 + w2*t2).mean() 36 | 37 | 38 | def parsimony(parameters): 39 | """Penalize the use of more primitives""" 40 | expected_primitives = parameters[0].sum(-1) 41 | 42 | return expected_primitives.mean() 43 | 44 | 45 | def entropy_bernoulli(parameters): 46 | """Minimize the entropy of each bernoulli variable pushing them to either 1 47 | or 0""" 48 | probs = parameters[0] 49 | sm = probs.new_tensor(1e-3) 50 | 51 | t1 = torch.log(torch.max(probs, sm)) 52 | t2 = torch.log(torch.max(1 - probs, sm)) 53 | 54 | return torch.mean((-probs * t1 - (1-probs) * t2).sum(-1)) 55 | 56 | 57 | def overlapping(F, X): 58 | """Penalize primitives that are inside other primitives 59 | 60 | Arguments: 61 | ----------- 62 | F: Tensor of shape BxNxM for the X points 63 | X: Tensor of shape BxNxMx3 containing N points transformed in the 64 | primitive's centric coordinate systems 65 | """ 66 | assert F.shape[0] == X.shape[0] 67 | assert F.shape[1] == X.shape[1] 68 | assert F.shape[2] == X.shape[2] 69 | assert X.shape[3] == 3 70 | assert len(F.shape) == 3 71 | assert len(X.shape) == 4 72 | B, N, M = F.shape 73 | 74 | loss = F.new_tensor(0) 75 | for j in range(M): 76 | f = F[F[:, :, j] > 0.5] 77 | if len(f) > 0: 78 | f[:, j] = f.new_tensor(0) 79 | loss += torch.max(f - f.new_tensor(0.5), f.new_tensor(0)).mean() 80 | 81 | return loss 82 | 83 | 84 | def get(regularizer, parameters, F, X_SQ, arguments): 85 | n_primitives = parameters[1].shape[1] 86 | regs = { 87 | "bernoulli_regularizer": partial( 88 | bernoulli, 89 | parameters, 90 | arguments.get("minimum_number_of_primitives", None), 91 | arguments.get("maximum_number_of_primitives", None), 92 | arguments.get("w1", None), 93 | arguments.get("w2", None) 94 | ), 95 | "parsimony_regularizer": partial(parsimony, parameters), 96 | "entropy_bernoulli_regularizer": partial( 97 | entropy_bernoulli, 98 | parameters 99 | ), 100 | "overlapping_regularizer": partial(overlapping, F, X_SQ), 101 | "sparsity_regularizer": partial( 102 | sparsity, 103 | parameters, 104 | arguments.get("minimum_number_of_primitives", None), 105 | arguments.get("maximum_number_of_primitives", None), 106 | arguments.get("w1", None), 107 | arguments.get("w2", None) 108 | ), 109 | } 110 | 111 | # Call the regularizer or return 0 112 | return regs.get(regularizer, lambda: parameters[0].new_tensor(0))() 113 | -------------------------------------------------------------------------------- /learnable_primitives/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/learnable_primitives/utils/__init__.py -------------------------------------------------------------------------------- /learnable_primitives/utils/pcl_voxelization.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def get_voxel_grid(bbox, grid_shape): 5 | """Given a bounding box and the dimensionality of a grid generate a grid of 6 | voxels and return their centers. 7 | 8 | Arguments: 9 | ---------- 10 | bbox: array(shape=(6, 1), dtype=np.float32) 11 | The min and max of the corners of the bbox that encloses the 12 | scene 13 | grid_shape: array(shape(3,), dtype=int32) 14 | The dimensions of the voxel grid used to discretize the 15 | scene 16 | Returns: 17 | -------- 18 | voxel_grid: array(shape=(3,)+grid_shape) 19 | The centers of the voxels 20 | """ 21 | # Make sure that we have the appropriate inputs 22 | assert bbox.shape[0] == 6 23 | assert bbox.shape[1] == 1 24 | 25 | xyz = [ 26 | np.linspace(s, e, c, endpoint=True, dtype=np.float32) 27 | for s, e, c in 28 | zip(bbox[:3], bbox[3:], grid_shape) 29 | ] 30 | bin_size = np.array([xyzi[1]-xyzi[0] for xyzi in xyz]).reshape(3, 1, 1, 1) 31 | return np.stack(np.meshgrid(*xyz, indexing="ij")) + bin_size/2 32 | -------------------------------------------------------------------------------- /learnable_primitives/utils/progbar.py: -------------------------------------------------------------------------------- 1 | 2 | import collections 3 | import sys 4 | import time 5 | 6 | import numpy as np 7 | 8 | 9 | class Progbar(object): 10 | """Displays a progress bar. 11 | # Arguments 12 | target: Total number of steps expected, None if unknown. 13 | width: Progress bar width on screen. 14 | verbose: Verbosity mode, 0 (silent), 1 (verbose), 2 (semi-verbose) 15 | stateful_metrics: Iterable of string names of metrics that 16 | should *not* be averaged over time. Metrics in this list 17 | will be displayed as-is. All others will be averaged 18 | by the progbar before display. 19 | interval: Minimum visual progress update interval (in seconds). 20 | """ 21 | 22 | def __init__(self, target, width=30, verbose=1, interval=0.05, 23 | stateful_metrics=None): 24 | self.target = target 25 | self.width = width 26 | self.verbose = verbose 27 | self.interval = interval 28 | if stateful_metrics: 29 | self.stateful_metrics = set(stateful_metrics) 30 | else: 31 | self.stateful_metrics = set() 32 | 33 | self._dynamic_display = ((hasattr(sys.stdout, 'isatty') and 34 | sys.stdout.isatty()) or 35 | 'ipykernel' in sys.modules) 36 | self._total_width = 0 37 | self._seen_so_far = 0 38 | self._values = collections.OrderedDict() 39 | self._start = time.time() 40 | self._last_update = 0 41 | 42 | def update(self, current, values=None): 43 | """Updates the progress bar. 44 | # Arguments 45 | current: Index of current step. 46 | values: List of tuples: 47 | `(name, value_for_last_step)`. 48 | If `name` is in `stateful_metrics`, 49 | `value_for_last_step` will be displayed as-is. 50 | Else, an average of the metric over time will be displayed. 51 | """ 52 | values = values or [] 53 | for k, v in values: 54 | if k not in self.stateful_metrics: 55 | if k not in self._values: 56 | self._values[k] = [v * (current - self._seen_so_far), 57 | current - self._seen_so_far] 58 | else: 59 | self._values[k][0] += v * (current - self._seen_so_far) 60 | self._values[k][1] += (current - self._seen_so_far) 61 | else: 62 | # Stateful metrics output a numeric value. This representation 63 | # means "take an average from a single value" but keeps the 64 | # numeric formatting. 65 | self._values[k] = [v, 1] 66 | self._seen_so_far = current 67 | 68 | now = time.time() 69 | info = ' - %.0fs' % (now - self._start) 70 | if self.verbose == 1: 71 | if (now - self._last_update < self.interval and 72 | self.target is not None and current < self.target): 73 | return 74 | 75 | prev_total_width = self._total_width 76 | if self._dynamic_display: 77 | sys.stdout.write('\b' * prev_total_width) 78 | sys.stdout.write('\r') 79 | else: 80 | sys.stdout.write('\n') 81 | 82 | if self.target is not None: 83 | numdigits = int(np.floor(np.log10(self.target))) + 1 84 | barstr = '%%%dd/%d [' % (numdigits, self.target) 85 | bar = barstr % current 86 | prog = float(current) / self.target 87 | prog_width = int(self.width * prog) 88 | if prog_width > 0: 89 | bar += ('=' * (prog_width - 1)) 90 | if current < self.target: 91 | bar += '>' 92 | else: 93 | bar += '=' 94 | bar += ('.' * (self.width - prog_width)) 95 | bar += ']' 96 | else: 97 | bar = '%7d/Unknown' % current 98 | 99 | self._total_width = len(bar) 100 | sys.stdout.write(bar) 101 | 102 | if current: 103 | time_per_unit = (now - self._start) / current 104 | else: 105 | time_per_unit = 0 106 | if self.target is not None and current < self.target: 107 | eta = time_per_unit * (self.target - current) 108 | if eta > 3600: 109 | eta_format = ('%d:%02d:%02d' % 110 | (eta // 3600, (eta % 3600) // 60, eta % 60)) 111 | elif eta > 60: 112 | eta_format = '%d:%02d' % (eta // 60, eta % 60) 113 | else: 114 | eta_format = '%ds' % eta 115 | 116 | info = ' - ETA: %s' % eta_format 117 | else: 118 | if time_per_unit >= 1: 119 | info += ' %.0fs/step' % time_per_unit 120 | elif time_per_unit >= 1e-3: 121 | info += ' %.0fms/step' % (time_per_unit * 1e3) 122 | else: 123 | info += ' %.0fus/step' % (time_per_unit * 1e6) 124 | 125 | for k in self._values: 126 | info += ' - %s:' % k 127 | if isinstance(self._values[k], list): 128 | avg = np.mean( 129 | self._values[k][0] / max(1, self._values[k][1])) 130 | if abs(avg) > 1e-3: 131 | info += ' %.4f' % avg 132 | else: 133 | info += ' %.4e' % avg 134 | else: 135 | info += ' %s' % self._values[k] 136 | 137 | self._total_width += len(info) 138 | if prev_total_width > self._total_width: 139 | info += (' ' * (prev_total_width - self._total_width)) 140 | 141 | if self.target is not None and current >= self.target: 142 | info += '\n' 143 | 144 | sys.stdout.write(info) 145 | sys.stdout.flush() 146 | 147 | elif self.verbose == 2: 148 | if self.target is None or current >= self.target: 149 | for k in self._values: 150 | info += ' - %s:' % k 151 | avg = np.mean( 152 | self._values[k][0] / max(1, self._values[k][1])) 153 | if avg > 1e-3: 154 | info += ' %.4f' % avg 155 | else: 156 | info += ' %.4e' % avg 157 | info += '\n' 158 | 159 | sys.stdout.write(info) 160 | sys.stdout.flush() 161 | 162 | self._last_update = now 163 | 164 | def add(self, n, values=None): 165 | self.update(self._seen_so_far + n, values) 166 | -------------------------------------------------------------------------------- /learnable_primitives/volumetric_iou.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from .primitives import transform_to_primitives_centric_system,\ 5 | inside_outside_function, points_to_cuboid_distances 6 | 7 | 8 | def broadcast_cross(a, b): 9 | a1, a2, a3 = a[..., 0], a[..., 1], a[..., 2] 10 | b1, b2, b3 = b[..., 0], b[..., 1], b[..., 2] 11 | 12 | s1 = a2*b3 - a3*b2 13 | s2 = a3*b1 - a1*b3 14 | s3 = a1*b2 - a2*b1 15 | 16 | return torch.stack([s1, s2, s3], dim=-1) 17 | 18 | 19 | def inside_mesh(points, triangles): 20 | """Return a boolean mask that is true for all the points that are inside 21 | the mesh defined by the triangles. 22 | 23 | Arguments 24 | --------- 25 | points: array Nx3 the 3d points 26 | triangles: array Mx3x3 the 3 3d points defining each of the M triangles 27 | """ 28 | # Establish a point that is outside the mesh 29 | out = (triangles.min(0)[0].min(0)[0] - 1).reshape(1, 3) 30 | 31 | # Get a normalized ray direction 32 | rays = points - out 33 | ray_norms = torch.sqrt((rays**2).sum(dim=1, keepdim=True)) 34 | rays /= ray_norms 35 | 36 | # Calculate the edges that share point 0 of the triangle 37 | edges1 = triangles[:, 1] - triangles[:, 0] 38 | edges2 = triangles[:, 2] - triangles[:, 0] 39 | 40 | # Start the calculation of the determinant (see Moeller & Trumbore 1997) 41 | pvec = broadcast_cross(rays.reshape(1, -1, 3), edges2.reshape(-1, 1, 3)) 42 | dets = (edges1.reshape(-1, 1, 3) * pvec).sum(dim=-1) 43 | inv_dets = 1.0 / dets 44 | 45 | # Calculate U 46 | tvec = out - triangles[:, 0] 47 | u = (tvec.reshape(-1, 1, 3) * pvec).sum(dim=-1) * inv_dets 48 | 49 | # Calculate V 50 | qvec = torch.cross(tvec, edges1) 51 | v = (rays.reshape(1, -1, 3) * qvec.reshape(-1, 1, 3)).sum(dim=-1) 52 | v = v * inv_dets 53 | 54 | # Compute all intersections 55 | intersections = (u > 0) * (v > 0) * ((u + v) < 1) 56 | 57 | # Compute the lengths of the intersections 58 | ts = (edges2 * qvec).sum(dim=1).reshape(-1, 1) * inv_dets 59 | 60 | # We need points that starting from `out` the line intersects an odd number 61 | # of triangles. 62 | mask = ((intersections * (ts <= ray_norms.reshape(1, -1))).sum(dim=0) % 2) 63 | mask = mask == 1 64 | 65 | return mask 66 | 67 | 68 | def _test_inside_mesh(): 69 | # Firstly lets make a cube and an in cube function 70 | from scipy.spatial import ConvexHull 71 | 72 | def in_cube(P): 73 | def b(p): 74 | return (0.3 < p) * (p < 0.6) 75 | return b(P[:, 0]) * b(P[:, 1]) * b(P[:, 2]) 76 | cube = torch.FloatTensor([ 77 | [0.3, 0.3, 0.3], 78 | [0.3, 0.3, 0.6], 79 | [0.3, 0.6, 0.3], 80 | [0.3, 0.6, 0.6], 81 | [0.6, 0.3, 0.3], 82 | [0.6, 0.3, 0.6], 83 | [0.6, 0.6, 0.3], 84 | [0.6, 0.6, 0.6] 85 | ]) 86 | h = ConvexHull(cube.numpy()) 87 | triangles = cube[h.simplices.ravel()].reshape(12, 3, 3) 88 | 89 | # Random points in [0, 1]^3 90 | P = torch.rand(100000, 3) 91 | 92 | mask1 = in_cube(P) 93 | mask2 = inside_mesh(P, triangles) 94 | 95 | # The ratio of the volume of the unit cube to the 0.3^3 cube is 0.027. 96 | d1 = 0.027 - mask1.sum().item() / 100000 97 | d2 = 0.027 - mask2.sum().item() / 100000 98 | 99 | assert -0.0005 < d1 < 0.0005 100 | assert -0.0005 < d2 < 0.0005 101 | 102 | 103 | def inside_sqs(P, y_hat, use_cuboids=False, use_sq=False, prob_threshold=0.5): 104 | """ 105 | Arguments: 106 | --------- 107 | P: Tensor with size BxNx3, with N points 108 | y_hat: List of Tensors containing the predictions of the network 109 | use_cuboids: when True use cuboids as primitives 110 | use_sq: when True use SQ as primitives 111 | Returns: 112 | M: Boolean mask with points are inside the SQs 113 | """ 114 | # Make sure that everything has the right shape 115 | assert P.shape[-1] == 3 116 | 117 | # Declare some variables 118 | B = P.shape[0] # batch size 119 | N = P.shape[1] # number of points per sample 120 | M = y_hat[0].shape[1] # number of primitives 121 | 122 | probs = y_hat[0] 123 | translations = y_hat[1].view(B, M, 3) 124 | rotations = y_hat[2].view(B, M, 4) 125 | shapes = y_hat[3].view(B, M, 3) 126 | epsilons = y_hat[4].view(B, M, 2) 127 | 128 | # Transform the 3D points from world-coordinates to primitive-centric 129 | # coordinates 130 | X_transformed = transform_to_primitives_centric_system( 131 | P, 132 | translations, 133 | rotations 134 | ) 135 | assert X_transformed.shape == (B, N, M, 3) 136 | if use_sq: 137 | F = inside_outside_function( 138 | X_transformed, 139 | shapes, 140 | epsilons 141 | ) 142 | inside = F <= 1 143 | elif use_cuboids: 144 | F = points_to_cuboid_distances(X_transformed, shapes) 145 | inside = F <= 0 146 | 147 | probs_mask = probs.unsqueeze(1) > prob_threshold 148 | assert inside.shape == (B, N, M) 149 | inside = inside * probs_mask 150 | 151 | # For every row if a column is 1.0 then that point is inside the SQs 152 | return inside.any(dim=-1) 153 | -------------------------------------------------------------------------------- /learnable_primitives/voxelizers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import uuid 4 | 5 | from .utils.pcl_voxelization import get_voxel_grid 6 | 7 | 8 | class VoxelizerFactory(object): 9 | def __init__(self, voxelizer, output_shape, save_voxels_to=None): 10 | if not isinstance(voxelizer, str) or\ 11 | voxelizer not in ["tsdf_grid", "occupancy_grid", "image"]: 12 | raise AttributeError("The voxelizer is invalid") 13 | if not isinstance(output_shape, np.ndarray): 14 | raise ValueError( 15 | "Output_shape should be a np.ndarray, but is %r" 16 | % (output_shape,) 17 | ) 18 | 19 | self._voxelizer = voxelizer 20 | # Arguments for teh OccupancyGrid 21 | self.output_shape = output_shape 22 | self._save_voxels_to = save_voxels_to 23 | 24 | @property 25 | def voxelizer(self): 26 | if self._voxelizer == "tsdf_grid": 27 | return TSDFPrecomputed(self.output_shape) 28 | elif self._voxelizer == "occupancy_grid": 29 | return OccupancyGrid(self.output_shape, None, self._save_voxels_to) 30 | elif self._voxelizer == "image": 31 | return ImagePrecomputed(self.output_shape) 32 | 33 | 34 | class OccupancyGrid(object): 35 | """OccupancyGrid class is a wrapper for the occupancy grid. 36 | """ 37 | def __init__(self, output_shape, bbox=None, save_voxels_to=None): 38 | self._output_shape = output_shape 39 | 40 | # Array that contains the voxel centers 41 | self._voxel_grid = None 42 | self.save_voxels_to = save_voxels_to 43 | 44 | @property 45 | def output_shape(self): 46 | return ((1,) + tuple(self._output_shape)) 47 | 48 | def _bbox_from_points(self, pts): 49 | bbox = np.array([-0.51]*3 + [0.51]*3) 50 | step = (bbox[3:] - bbox[:3])/self._output_shape 51 | 52 | return bbox, step 53 | 54 | def voxel_grid(self, pts): 55 | if self._voxel_grid is None: 56 | # Get the bounding box from the points 57 | bbox, _ = self._bbox_from_points(pts) 58 | 59 | self._voxel_grid = get_voxel_grid( 60 | bbox.reshape(-1, 1), 61 | self.output_shape 62 | ).astype(np.float32) 63 | return self._voxel_grid 64 | 65 | def voxelize(self, mesh): 66 | """Given a Mesh object we want to estimate the occupancy grid 67 | """ 68 | pts = mesh.points 69 | # Make sure that the points have the correct shape 70 | assert pts.shape[0] == 3 71 | # Make sure that points lie in the unit cube 72 | if (np.any(np.abs(pts.min(axis=-1)) > 0.51) or 73 | np.any(pts.max(axis=-1) > 0.51)): 74 | raise Exception( 75 | "The points do not lie in the unit cube min:%r - max:%r" 76 | % (pts.min(axis=-1), pts.max(axis=-1)) 77 | ) 78 | 79 | bbox, step = self._bbox_from_points(pts) 80 | occupancy_grid = np.zeros(tuple(self._output_shape), dtype=np.float32) 81 | 82 | idxs = ((pts.T - bbox[:3].T)/step).astype(np.int32) 83 | # Array that will contain the occupancy grid, namely if a point lies 84 | # within a voxel then we assign one to this voxel, otherwise we assign 85 | # 0 86 | occupancy_grid[idxs[:, 0], idxs[:, 1], idxs[:, 2]] = 1.0 87 | if self.save_voxels_to is not None: 88 | unique_filename = str(uuid.uuid4()) 89 | import matplotlib 90 | matplotlib.use("agg") 91 | import matplotlib.pyplot as plt 92 | from mpl_toolkits.mplot3d import Axes3D 93 | 94 | fig = plt.figure() 95 | ax = fig.gca(projection='3d') 96 | ax.voxels(occupancy_grid, edgecolor='k') 97 | ax.view_init(elev=10, azim=80) 98 | plt.savefig( 99 | os.path.join(self.save_voxels_to, unique_filename+".png") 100 | ) 101 | plt.close() 102 | return occupancy_grid[np.newaxis] 103 | 104 | def get_occupied_voxel_centers(self, mesh): 105 | occ = self.voxelize(mesh) 106 | voxel_grid = self.voxel_grid(mesh.points) 107 | 108 | return voxel_grid[:, occ == 1] 109 | 110 | def get_X(self, model): 111 | gt_mesh = model.groundtruth_mesh 112 | return self.voxelize(gt_mesh) 113 | 114 | 115 | class PrecomputedVoxelizer(object): 116 | """PrecomputedVoxelizer is a wrapper for precomputed data, such as 117 | precomputed TSDFs and occupancy grids. 118 | """ 119 | def __init__(self, output_shape): 120 | self._output_shape = output_shape 121 | 122 | @property 123 | def output_shape(self): 124 | return tuple(self._output_shape) 125 | 126 | def get_X(self, model): 127 | raise NotImplementedError() 128 | 129 | 130 | class ImagePrecomputed(PrecomputedVoxelizer): 131 | def get_X(self, model): 132 | X = model.random_image 133 | # Transpose image to have the right size for pytorch 134 | return np.transpose(X, (2, 0, 1)) 135 | 136 | 137 | class TSDFPrecomputed(PrecomputedVoxelizer): 138 | @property 139 | def output_shape(self): 140 | return (1,) + tuple(self._output_shape) 141 | 142 | def get_X(self, model): 143 | return model.tsdf[np.newaxis] 144 | -------------------------------------------------------------------------------- /misc/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paschalidoud/superquadric_parsing/19e365f012fb34c5997d2d5c28a5c121228d8063/misc/__init__.py -------------------------------------------------------------------------------- /misc/chamfer_vs_inside_outside_local_minima.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os 4 | 5 | import numpy as np 6 | import torch 7 | from pyquaternion import Quaternion 8 | 9 | from shapes import Shape, Cuboid 10 | from learnable_primitives.loss_functions import\ 11 | euclidean_dual_loss 12 | from learnable_primitives.equal_distance_sampler_sq import\ 13 | EqualDistanceSamplerSQ 14 | from learnable_primitives.primitives import\ 15 | euler_angles_to_rotation_matrices, quaternions_to_rotation_matrices 16 | from learnable_primitives.mesh import MeshFromOBJ 17 | 18 | import matplotlib 19 | matplotlib.use("agg") 20 | import matplotlib.pyplot as plt 21 | from mpl_toolkits.mplot3d import Axes3D 22 | from matplotlib.patches import Rectangle 23 | from matplotlib import rc 24 | #rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']}) 25 | #rc('text', usetex=True) 26 | rc('text', usetex=True) 27 | # rc('font', **{'family': 'serif', 'serif': ['Computer Modern']}) 28 | plt.rc("font", size=10, family="serif") 29 | 30 | def fexp(x, p): 31 | return np.sign(x)*(np.abs(x)**p) 32 | 33 | def sq_surface(a1, a2, a3, e1, e2, eta, omega): 34 | x = a1 * fexp(np.cos(eta), e1) * fexp(np.cos(omega), e2) 35 | y = a2 * fexp(np.cos(eta), e1) * fexp(np.sin(omega), e2) 36 | z = a3 * fexp(np.sin(eta), e1) 37 | return x, y, z 38 | 39 | 40 | def get_sq(a1, a2, a3, e1, e2, R, t, n_samples=100): 41 | """Computes a SQ given a set of parameters and saves it into a np array 42 | """ 43 | assert R.shape == (3, 3) 44 | assert t.shape == (3, 1) 45 | 46 | eta = np.linspace(-np.pi/2, np.pi/2, n_samples, endpoint=True) 47 | omega = np.linspace(-np.pi, np.pi, n_samples, endpoint=True) 48 | eta, omega = np.meshgrid(eta, omega) 49 | x, y, z = sq_surface(a1, a2, a3, e1, e2, eta, omega) 50 | 51 | # Get an array of size 3x10000 that contains the points of the SQ 52 | points = np.stack([x, y, z]).reshape(3, -1) 53 | points_transformed = R.T.dot(points) + t 54 | # print "R:", R 55 | # print "t:", t 56 | # print "e:", [e1, e2] 57 | 58 | x_tr = points_transformed[0].reshape(n_samples, n_samples) 59 | y_tr = points_transformed[1].reshape(n_samples, n_samples) 60 | z_tr = points_transformed[2].reshape(n_samples, n_samples) 61 | 62 | return x_tr, y_tr, z_tr, points_transformed 63 | 64 | 65 | def cubes_inside(x_1, y_1, z_1, x_2, y_2, z_2): 66 | # cube 1 67 | x_min = -x_1 68 | x_max = x_1 69 | y_min = -y_1 70 | y_max = y_1 71 | z_min = -z_1 72 | z_max = z_1 73 | c1 = Cuboid(x_min, x_max, y_min, y_max, z_min, z_max) 74 | c1.translate(np.array([[0.2, 0.2, 0.0]]).T) 75 | 76 | # cube 2 77 | x_min = -x_2 78 | x_max = x_2 79 | y_min = -y_2 80 | y_max = y_2 81 | z_min = -z_2 82 | z_max = z_2 83 | c2 = Cuboid(x_min, x_max, y_min, y_max, z_min, z_max) 84 | 85 | # Denote some useful variables 86 | B = 1 # batch size 87 | M = 2 # number of primitives 88 | shapes = torch.zeros(B, M*3, dtype=torch.float, requires_grad=True) 89 | shapes[0, 0] = 0.1 90 | shapes[0, 1] = 0.1 91 | shapes[0, 2] = z_1 92 | shapes[0, 3] = 0.1 93 | shapes[0, 4] = 0.1 94 | shapes[0, 5] = z_2 95 | # probs 96 | probs = torch.ones(B, M, dtype=torch.float, requires_grad=True) 97 | translations = torch.zeros(B, M*3, dtype=torch.float, requires_grad=True) 98 | translations[0, 0] = 0.2 99 | translations[0, 1] = 0.6 100 | translations[0, 3:5] = -0.5 101 | #translations = 0.5 - torch.randn(B, M*3, dtype=torch.float, requires_grad=True) 102 | quaternions = torch.zeros(B, M*4, dtype=torch.float, requires_grad=True) 103 | quaternions[0, 0] = 1.0 104 | quaternions[0, 4] = 1.0 105 | #for i in range(2): 106 | # q = Quaternion.random() 107 | # for j in range(4): 108 | # quaternions[0, 4*i+j] = q[j] 109 | 110 | epsilons = torch.ones(B, M*2, dtype=torch.float, requires_grad=True) * 0.25 111 | 112 | y_hat = [probs, translations, quaternions, shapes, epsilons] 113 | y_hat = [ 114 | torch.tensor(yi.data, requires_grad=True) 115 | for yi in y_hat 116 | ] 117 | 118 | return c1, c2, y_hat 119 | 120 | 121 | def get_translation(shapes): 122 | p_min = np.min([ 123 | s.points.min(axis=1) 124 | for s in shapes 125 | ], axis=0) 126 | p_max = np.max([ 127 | s.points.max(axis=1) 128 | for s in shapes 129 | ], axis=0) 130 | 131 | return (p_max - p_min) / 2 + p_min 132 | 133 | 134 | def get_rectangle(cube, T, **kwargs): 135 | corner1 = cube.points[:, 0] 136 | corner2 = cube.points[:, -1] 137 | 138 | return Rectangle( 139 | corner1-T, 140 | *(corner2-corner1), 141 | **kwargs 142 | ) 143 | 144 | 145 | if __name__ == "__main__": 146 | c1, c2, y_hat = cubes_inside(0.2, 0.2, 0.1, 0.1, 0.1, 0.1) 147 | T = get_translation([c1, c2]) 148 | probs, translations, quats, shapes, epsilons = y_hat 149 | c = Shape.from_shapes([c1, c2]) 150 | c.save_as_mesh("/tmp/mesh.obj", "obj") 151 | m = MeshFromOBJ("/tmp/mesh.obj") 152 | y_target = torch.from_numpy( 153 | m.sample_faces(1000).astype(np.float32) 154 | ).float().unsqueeze(0) 155 | 156 | # A sampler instance 157 | e = EqualDistanceSamplerSQ(200) 158 | # Compute the loss for the current experiment 159 | l_weights = { 160 | "coverage_loss_weight": 1.0, 161 | "consistency_loss_weight": 1.0, 162 | } 163 | reg_terms = { 164 | "regularizer_type": [], 165 | "shapes_regularizer_weight": 0.0, 166 | "bernoulli_with_sparsity_regularizer_weight": 0.0, 167 | "bernoulli_regularizer_weight": 0.0, 168 | "entropy_bernoulli_weight": 0.0, 169 | "partition_regularizer_weight": 0.0, 170 | "parsimony_regularizer_weight": 0.0, 171 | "overlapping_regularizer_weight": 0.0, 172 | "minimum_number_of_primitives": 0.0 173 | } 174 | 175 | 176 | use_simple_cuboids = False 177 | use_sq = False 178 | lr = 1e-1 # learning rate 179 | n_iters = 300 # number of gradient_descent iterations 180 | optim = torch.optim.SGD([translations, shapes], lr=lr, momentum=0.9) 181 | for k in range(n_iters): 182 | optim.zero_grad() 183 | loss, debug_stats = euclidean_dual_loss( 184 | y_hat, 185 | y_target, 186 | reg_terms, 187 | e, 188 | l_weights, 189 | use_simple_cuboids, 190 | use_sq, 191 | False 192 | ) 193 | loss.backward() 194 | print "It: %d - loss: %f - cnst_loss:%f - cvrg_loss:%f" %( 195 | k, loss, debug_stats[-1], debug_stats[-2] 196 | ) 197 | 198 | if (k % 10) == 0: 199 | fig = plt.figure(figsize=(4, 3)) 200 | axis = plt.gca() 201 | axis.add_patch(get_rectangle(c1, T, edgecolor='g', fill=None, alpha=1, label="target", linewidth=2)) 202 | axis.add_patch(get_rectangle(c2, T, edgecolor='g', fill=None, alpha=1, linewidth=2)) 203 | axis.add_patch( 204 | Rectangle( 205 | (translations[0, 0] - shapes[0, 0], translations[0, 1] - shapes[0, 1]), 206 | 2*shapes[0, 0], 207 | 2*shapes[0, 1], 208 | edgecolor='r', 209 | fill=None, 210 | linestyle='--', alpha=1, linewidth=2, label="primitive 1" 211 | )) 212 | axis.add_patch( 213 | Rectangle((translations[0, 3] - shapes[0, 3], translations[0, 4] - shapes[0, 4]), 214 | 2*shapes[0, 3], 2*shapes[0, 4], edgecolor='b', fill=None, linestyle='--', alpha=1, linewidth=2, label="primitive 2") 215 | ) 216 | plt.xlim((-0.75, 0.75)) 217 | plt.ylim((-0.75, 0.75)) 218 | plt.legend(loc="upper left") 219 | # plt.title(r"\textbf{Chamfer distance (Ours)}") 220 | # plt.savefig("/tmp/empirical_test/chamfer_iter_%05d.png" % (k,), bbox_inches="tight") 221 | # plt.savefig("/tmp/empirical_test/chamfer_iter_%05d.pdf" % (k,), bbox_inches="tight") 222 | plt.title(r"\textbf{Truncated distance (Tulsiani et al.)}") 223 | plt.savefig("/tmp/empirical_test/trunc_iter_%05d.png" % (k,), bbox_inches="tight") 224 | plt.savefig("/tmp/empirical_test/trunc_iter_%05d.pdf" % (k,), bbox_inches="tight") 225 | plt.close() 226 | 227 | # Update everything 228 | optim.step() 229 | shapes.data.clamp_(min=0.01, max=0.5) 230 | -------------------------------------------------------------------------------- /misc/create_simple_shapes_dataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script used to generate a cuboid dataset with cubes and rectangles under 3 | various shapes, rotations, translations following the general format of 4 | ShapeNet. 5 | """ 6 | import argparse 7 | import random 8 | import os 9 | from string import ascii_letters, digits 10 | import sys 11 | 12 | import numpy as np 13 | from progress.bar import Bar 14 | from pyquaternion import Quaternion 15 | 16 | from shapes import Shape, Cuboid, Ellipsoid 17 | 18 | from learnable_primitives.mesh import MeshFromOBJ 19 | 20 | 21 | def get_single_cube(minimum, maximum): 22 | minimum = minimum[0] 23 | maximum = maximum[0] 24 | r = minimum + np.random.rand() * (maximum-minimum) 25 | 26 | return Cuboid(-r, r, -r, r, -r, r) 27 | 28 | 29 | def get_single_rectangle(minimum, maximum): 30 | minimum = np.array(minimum) 31 | maximum = np.array(maximum) 32 | rs = minimum + np.random.rand(3) * (maximum - minimum) 33 | 34 | return Cuboid(-rs[0], rs[0], -rs[1], rs[1], -rs[2], rs[2]) 35 | 36 | 37 | def adjacent_cubes(R): 38 | x_max1, y_max1, z_max1 = tuple(np.random.rand(3)) 39 | x_max2, y_max2, z_max2 = tuple(np.random.rand(3)) 40 | c1 = Cuboid(-x_max1, x_max1, -y_max1, y_max1, -z_max1, z_max1) 41 | c2 = Cuboid(-x_max2, x_max2, -y_max2, y_max2, -z_max2, z_max2) 42 | t1 = np.array([ 43 | [0.0, y_max2 + y_max1, 0.0], 44 | [x_max2 + x_max1, 0.0, 0.0], 45 | [0.0, 0.0, z_max2 + z_max1] 46 | ]) 47 | t = t1[np.random.choice(np.arange(3))].reshape(3, 1) 48 | c2.translate(t) 49 | c1.rotate(R) 50 | c2.rotate(R) 51 | return c1, c2 52 | 53 | 54 | def multiple_cubes(R1, R2, t): 55 | x_max1, y_max1, z_max1 = tuple(np.random.rand(3)) 56 | x_max2, y_max2, z_max2 = tuple(np.random.rand(3)) 57 | c1 = Cuboid(-x_max1, x_max1, -y_max1, y_max1, -z_max1, z_max1) 58 | c2 = Cuboid(-x_max2, x_max2, -y_max2, y_max2, -z_max2, z_max2) 59 | c1.rotate(R1) 60 | c2.translate(t) 61 | c2.rotate(R2) 62 | #c2.translate(R2.dot(t)) 63 | return c1, c2 64 | 65 | 66 | def main(argv): 67 | parser = argparse.ArgumentParser( 68 | description="Generate a cuboid dataset" 69 | ) 70 | parser.add_argument( 71 | "output_directory", 72 | help="Save the dataset in this directory" 73 | ) 74 | parser.add_argument( 75 | "--n_samples", 76 | type=int, 77 | default=10, 78 | help="Number of training samples to be generated" 79 | ) 80 | parser.add_argument( 81 | "--shapes_type", 82 | default="cubes", 83 | choices=[ 84 | "cubes", 85 | "cubes_translated", 86 | "cubes_rotated_translated", 87 | "cubes_rotated", 88 | "rectangles", 89 | "rectangles_translated", 90 | "rectangles_rotated", 91 | "rectangles_rotated_translated", 92 | "ellipsoid", 93 | "random" 94 | ], 95 | help="The type of the shapes in every sample" 96 | ) 97 | parser.add_argument( 98 | "--n_shapes_per_samples", 99 | type=int, 100 | default=1, 101 | help="Number of shapes per sample" 102 | ) 103 | parser.add_argument( 104 | "--maximum", 105 | type=lambda x: tuple(map(float, x.split(","))), 106 | default="0.5,0.5,0.5", 107 | help="Maximum size along every axis" 108 | ) 109 | parser.add_argument( 110 | "--minimum", 111 | type=lambda x: tuple(map(float, x.split(","))), 112 | default="0.13,0.13,0.13", 113 | help="Maximum size along every axis" 114 | ) 115 | args = parser.parse_args(argv) 116 | 117 | # Check if output directory exists and if it doesn't create it 118 | if not os.path.exists(args.output_directory): 119 | os.makedirs(args.output_directory) 120 | 121 | # Create a directory based on the type of the shapes inside the output 122 | # directory 123 | output_directory = os.path.join( 124 | args.output_directory, 125 | args.shapes_type 126 | ) 127 | 128 | ranges = None 129 | if "cubes" in args.shapes_type: 130 | # Make sure that the maximum and minimum range are equal along each 131 | # axis 132 | assert args.maximum[0] == args.maximum[1] 133 | assert args.maximum[1] == args.maximum[2] 134 | assert args.minimum[0] == args.minimum[1] 135 | assert args.minimum[1] == args.minimum[2] 136 | ranges = np.linspace( 137 | args.minimum[0], 138 | args.maximum[0], 139 | 10, 140 | endpoint=False 141 | ) 142 | 143 | # elif "rectangles" in args.shapes_type: 144 | else: 145 | ranges = [ 146 | np.linspace(args.minimum[0], args.maximum[0], 10, endpoint=False), 147 | np.linspace(args.minimum[1], args.maximum[1], 10, endpoint=False), 148 | np.linspace(args.minimum[2], args.maximum[2], 10, endpoint=False), 149 | ] 150 | 151 | bar = Bar("Generating %d cuboids" % (args.n_samples,), max=args.n_samples) 152 | c = None 153 | for i in range(args.n_samples): 154 | if "cubes" in args.shapes_type: 155 | c = get_single_cube(args.minimum, args.maximum) 156 | if "rectangles" in args.shapes_type: 157 | c = get_single_rectangle(args.minimum, args.maximum) 158 | 159 | if "translated" in args.shapes_type: 160 | t = 0.3*np.random.random((3, 1)) 161 | c.translate(t) 162 | 163 | if "rotated" in args.shapes_type: 164 | q = Quaternion.random() 165 | R = q.rotation_matrix 166 | c.rotate(R) 167 | 168 | if "ellipsoid" in args.shapes_type: 169 | abc = np.random.random((3, 1)) 170 | c1 = Ellipsoid(abc[0], abc[1], abc[2]) 171 | c2 = Ellipsoid(abc[0], abc[1], abc[2]) 172 | c3 = Ellipsoid(abc[0], abc[1], abc[2]) 173 | q = Quaternion.random() 174 | R = q.rotation_matrix 175 | c2.rotate(R) 176 | q = Quaternion.random() 177 | R = q.rotation_matrix 178 | c3.rotate(R) 179 | # t = 0.3*np.random.random((3, 1)) 180 | # c1.translate(t) 181 | c = Shape.from_shapes([c1, c2, c3]) 182 | 183 | if "random" in args.shapes_type: 184 | #if random.choice((True, False)): 185 | #if True: 186 | # q = Quaternion.random() 187 | # c1, c2 = adjacent_cubes(q.rotation_matrix) 188 | #else: 189 | if True: 190 | q1 = Quaternion.random() 191 | q2 = Quaternion.random() 192 | c1, c2 = multiple_cubes( 193 | q1.rotation_matrix, 194 | q2.rotation_matrix, 195 | 3.5*np.random.random((3, 1)) 196 | ) 197 | # q = Quaternion.random() 198 | # c1, c2 = adjacent_cubes(q.rotation_matrix) 199 | # q1 = Quaternion.random() 200 | # x_max1, y_max1, z_max1 = tuple(np.random.rand(3)) 201 | # c3 = Cuboid(-x_max1, x_max1, -y_max1, y_max1, -z_max1, z_max1) 202 | # c3.rotate(q1.rotation_matrix) 203 | # c3.translate(np.random.random((3,1)).reshape(3, -1)) 204 | c = Shape.from_shapes([c1, c2]) 205 | 206 | # Create subdirectory to save the sample 207 | folder_name = ''.join([ 208 | random.choice(ascii_letters + digits) for n in xrange(32) 209 | ]) 210 | base_dir = os.path.join(output_directory, folder_name, "models") 211 | if not os.path.exists(base_dir): 212 | os.makedirs(base_dir) 213 | # print base_dir 214 | # Save as obj file 215 | c.save_as_mesh(os.path.join(base_dir, "model_normalized.obj"), "obj") 216 | c.save_as_mesh(os.path.join(base_dir, "model_normalized.ply"), "ply") 217 | c.save_as_pointcloud( 218 | os.path.join(base_dir, "model_normalized_pcl.obj"), "obj" 219 | ) 220 | if "translated" in args.shapes_type: 221 | print os.path.join(base_dir, "model_normalized_pcl.obj"), t.T 222 | if "rotated" in args.shapes_type: 223 | print os.path.join(base_dir, "model_normalized_pcl.obj"), q 224 | bar.next() 225 | 226 | for i in os.listdir(output_directory): 227 | x = os.path.join(output_directory, i, "models/model_normalized.obj") 228 | m = MeshFromOBJ(x) 229 | print x, m.points.max(-1) 230 | 231 | 232 | if __name__ == "__main__": 233 | main(sys.argv[1:]) 234 | -------------------------------------------------------------------------------- /misc/create_spheres_dataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script used to generate a spheres dataset following the same format as 3 | ShapeNet. 4 | """ 5 | import argparse 6 | import os 7 | import sys 8 | 9 | import numpy as np 10 | from pyquaternion import Quaternion 11 | 12 | from shapes import Sphere, Ellipsoid, Shape 13 | 14 | from learnable_primitives.mesh import MeshFromOBJ 15 | from learnable_primitives.utils.progbar import Progbar 16 | 17 | 18 | def build_sequentially_attaching_sheres(N): 19 | """ 20 | Arguments: 21 | --------- 22 | N: number of shapes to be build 23 | """ 24 | def t(s): 25 | return s[0] 26 | 27 | def r(s): 28 | return s[1] 29 | 30 | def overlap(s1, s2): 31 | d = np.sqrt(np.sum((t(s1)-t(s2))**2)) 32 | return d < r(s1) + r(s2) 33 | 34 | def random_point(s, d=0): 35 | v = np.random.randn(3) 36 | v = v / np.sqrt((v**2).sum()) 37 | return t(s) + v*(r(s)+d) 38 | 39 | spheres = [(np.zeros(3), np.random.rand()*0.4 + 0.1)] 40 | while len(spheres) < N: 41 | s1 = spheres[np.random.choice(len(spheres))] 42 | s2r = np.random.rand()*0.4 + 0.1 43 | s2c = random_point(s1, s2r) 44 | s2 = (s2c, s2r) 45 | if not any(overlap(s, s2) for s in spheres): 46 | spheres.append(s2) 47 | 48 | return [Sphere(r(s)).translate(t(s)[:, np.newaxis]) for s in spheres] 49 | 50 | 51 | def main(argv): 52 | parser = argparse.ArgumentParser( 53 | description="Generate a cuboid dataset" 54 | ) 55 | parser.add_argument( 56 | "output_directory", 57 | help="Save the dataset in this directory" 58 | ) 59 | parser.add_argument( 60 | "--n_samples", 61 | type=int, 62 | default=10, 63 | help="Number of training samples to be generated" 64 | ) 65 | parser.add_argument( 66 | "--max_n_shapes_per_samples", 67 | type=int, 68 | default=4, 69 | help="Number of shapes per sample" 70 | ) 71 | args = parser.parse_args(argv) 72 | 73 | # Check if output directory exists and if it doesn't create it 74 | if not os.path.exists(args.output_directory): 75 | os.makedirs(args.output_directory) 76 | 77 | # Create a directory based on the type of the shapes inside the output 78 | # directory 79 | output_directory = os.path.join( 80 | args.output_directory, 81 | "spheres_dataset" 82 | ) 83 | print "Saving models to %s" % (output_directory,) 84 | 85 | prog = Progbar(args.n_samples) 86 | for i in range(args.n_samples): 87 | prims = build_sequentially_attaching_sheres( 88 | np.random.choice(np.arange(2, args.max_n_shapes_per_samples)) 89 | ) 90 | c = Shape.from_shapes(prims) 91 | # Create subdirectory to save the sample 92 | base_dir = os.path.join(output_directory, "%05d" % (i,), "models") 93 | if not os.path.exists(base_dir): 94 | os.makedirs(base_dir) 95 | # print base_dir 96 | # Save as obj file 97 | c.save_as_mesh(os.path.join(base_dir, "model_normalized.obj"), "obj") 98 | c.save_as_mesh(os.path.join(base_dir, "model_normalized.ply"), "ply") 99 | c.save_as_pointcloud( 100 | os.path.join(base_dir, "model_normalized_pcl.obj"), "obj" 101 | ) 102 | prog.update(i + 1) 103 | 104 | 105 | if __name__ == "__main__": 106 | main(sys.argv[1:]) 107 | -------------------------------------------------------------------------------- /misc/shapes.py: -------------------------------------------------------------------------------- 1 | from itertools import permutations 2 | import math 3 | 4 | from csg.core import CSG, Polygon as CSGPolygon, Vertex as CSGVertex 5 | import numpy as np 6 | from pyquaternion import Quaternion 7 | from scipy.spatial import ConvexHull 8 | import trimesh 9 | 10 | from learnable_primitives.pointcloud import Pointcloud 11 | 12 | 13 | class Shape(object): 14 | """A wrapper class for shapes""" 15 | def __init__(self, points, face_idxs): 16 | self._points = points 17 | self._faces_idxs = face_idxs 18 | 19 | @property 20 | def points(self): 21 | if self._points is None: 22 | raise NotImplementedError() 23 | return self._points 24 | 25 | @property 26 | def faces_idxs(self): 27 | if self._faces_idxs is None: 28 | raise NotImplementedError() 29 | return self._faces_idxs 30 | 31 | def save_as_mesh(self, filename, format="ply"): 32 | m = trimesh.Trimesh(vertices=self.points.T, faces=self.faces_idxs) 33 | # Make sure that the face orinetations are ok 34 | trimesh.repair.fix_normals(m, multibody=True) 35 | trimesh.repair.fix_winding(m) 36 | assert m.is_winding_consistent == True 37 | m.export(filename) 38 | 39 | def sample_faces(self, N=1000): 40 | m = trimesh.Trimesh(vertices=self.points.T, faces=self.faces_idxs) 41 | # Make sure that the face orinetations are ok 42 | trimesh.repair.fix_normals(m, multibody=True) 43 | trimesh.repair.fix_winding(m) 44 | assert m.is_winding_consistent == True 45 | P, t = trimesh.sample.sample_surface(m, N) 46 | return np.hstack([ 47 | P, m.face_normals[t, :] 48 | ]) 49 | 50 | def save_as_pointcloud(self, filename, format="ply"): 51 | pcl = Pointcloud(self.points) 52 | if format == "ply": 53 | pcl.save_ply(filename) 54 | else: 55 | pcl.save_obj(filename) 56 | 57 | def rotate(self, R): 58 | """ 3x3 rotation matrix that will rotate the points 59 | """ 60 | # Make sure that the rotation matrix has the right shape 61 | assert R.shape == (3, 3) 62 | self._points = R.T.dot(self.points) 63 | 64 | return self 65 | 66 | def translate(self, t): 67 | # Make sure thate everything has the right shape 68 | assert t.shape[0] == 3 69 | assert t.shape[1] == 1 70 | self._points = self.points + t 71 | 72 | return self 73 | 74 | def to_csg(self): 75 | points = self.points 76 | polygons = [ 77 | CSGPolygon([ 78 | CSGVertex(pos=points[:, i]) 79 | for i in face 80 | ]) 81 | for face in self.faces_idxs 82 | ] 83 | 84 | return CSG.fromPolygons(polygons) 85 | 86 | @classmethod 87 | def from_csg(cls, csg_object): 88 | points, face_idxs, n = csg_object.toVerticesAndPolygons() 89 | triangles = [] 90 | for face in face_idxs: 91 | if len(face) == 3: 92 | triangles.append(face) 93 | else: 94 | for j in range(2, len(face)): 95 | triangles.append((face[0], face[j-1], face[j])) 96 | 97 | return cls( 98 | np.array(points).T, 99 | np.array(triangles) 100 | ) 101 | 102 | @classmethod 103 | def from_shapes(cls, shapes): 104 | # Make sure that the input is a list of shapes 105 | isinstance(shapes, list) 106 | points = [] 107 | triangles = [] 108 | for i, s in enumerate(shapes): 109 | if len(points) == 0: 110 | triangles.append(s.faces_idxs) 111 | else: 112 | triangles.append(s.faces_idxs + i*points[-1].shape[1]) 113 | points.append(s.points) 114 | 115 | return cls( 116 | np.hstack(points), 117 | np.vstack(triangles) 118 | ) 119 | 120 | @staticmethod 121 | def get_orientation_of_face(points, face_idxs): 122 | # face_idxs corresponds to the indices of a single face 123 | assert len(face_idxs.shape) == 1 124 | assert face_idxs.shape[0] == 3 125 | 126 | x = np.vstack([ 127 | points.T[face_idxs, 0].T, 128 | points.T[face_idxs, 1].T, 129 | points.T[face_idxs, 2].T 130 | ]).T 131 | 132 | # Based on the Wikipedia article 133 | # https://en.wikipedia.org/wiki/Curve_orientation 134 | # If the determinant is negative, then the polygon is oriented 135 | # clockwise. If the determinant is positive, the polygon is oriented 136 | # counterclockwise 137 | return np.linalg.det(x) 138 | 139 | @staticmethod 140 | def fix_orientation_of_face(points, face_idxs): 141 | # face_idxs corresponds to the indices of a single face 142 | assert len(face_idxs.shape) == 1 143 | assert face_idxs.shape[0] == 3 144 | 145 | # Iterate over all possible permutations 146 | for item in permutations(face_idxs, face_idxs.shape[0]): 147 | t = np.array(item) 148 | orientation = Shape.get_orientation_of_face(points, t) 149 | if orientation < 0: 150 | pass 151 | else: 152 | return t 153 | 154 | 155 | class ConvexShape(Shape): 156 | """A wrapper class for convex shapes""" 157 | def __init__(self, points): 158 | self._points = points 159 | 160 | # Contains the convexhull of the set of points (see cv property) 161 | self._cv = None 162 | # Contains the faces_idxs (see face_idxs Shape property) 163 | self._faces_idxs = None 164 | 165 | @property 166 | def points(self): 167 | return self._points 168 | 169 | @property 170 | def cv(self): 171 | if self._cv is None: 172 | self._cv = ConvexHull(self.points.T) 173 | return self._cv 174 | 175 | @property 176 | def faces_idxs(self): 177 | if self._faces_idxs is None: 178 | self._faces_idxs = np.array(self.cv.simplices) 179 | self._make_consistent_orientation_of_faces() 180 | return self._faces_idxs 181 | 182 | def _make_consistent_orientation_of_faces(self): 183 | for i, face_idxs in zip(xrange(self.cv.nsimplex), self.faces_idxs): 184 | # Compute the orientation for the current face 185 | orientation = Shape.get_orientation_of_face(self.points, face_idxs) 186 | if orientation < 0: 187 | # if the orientation is negative, permute the face_idxs to make 188 | # it positive 189 | self._faces_idxs[i] =\ 190 | Shape.fix_orientation_of_face(self.points, face_idxs) 191 | 192 | 193 | class Cuboid(ConvexShape): 194 | def __init__(self, x_min, x_max, y_min, y_max, z_min, z_max): 195 | super(Cuboid, self).__init__(self._create_points( 196 | x_min, x_max, y_min, y_max, z_min, z_max 197 | )) 198 | 199 | def _create_points(self, x_min, x_max, y_min, y_max, z_min, z_max): 200 | vertices = np.array([ 201 | [x_min, x_max], 202 | [y_min, y_max], 203 | [z_min, z_max] 204 | ]) 205 | idxs = np.array([ 206 | [i, j, k] 207 | for i in range(2) 208 | for j in range(2) 209 | for k in range(2) 210 | ]) 211 | 212 | return vertices[np.tile(range(3), (8, 1)), idxs].T 213 | 214 | @staticmethod 215 | def keep_points_inside_cube(x, y, z, x_min, x_max, 216 | y_min, y_max, z_min, z_max): 217 | c1 = np.logical_and( 218 | np.logical_and(x >= x_min, x <= x_max), 219 | np.logical_and(y >= y_min, y <= y_max) 220 | ) 221 | c2 = np.logical_and( 222 | np.logical_and(y >= y_min, y <= y_max), 223 | np.logical_and(z >= z_min, z <= z_max) 224 | ) 225 | c3 = np.logical_and( 226 | np.logical_and(x >= x_min, x <= x_max), 227 | np.logical_and(z >= z_min, z <= z_max) 228 | ) 229 | c4 = np.logical_and(c1, z == z_min) 230 | c5 = np.logical_and(c1, z == z_max) 231 | c6 = np.logical_and(c2, x == x_min) 232 | c7 = np.logical_and(c2, x == x_max) 233 | c8 = np.logical_and(c3, y == y_min) 234 | c9 = np.logical_and(c3, y == y_max) 235 | return np.logical_or( 236 | np.logical_or(np.logical_or(c5, c6), np.logical_or(c7, c8)), 237 | np.logical_or(c4, c9) 238 | ) 239 | 240 | 241 | class Sphere(ConvexShape): 242 | def __init__(self, radius): 243 | self._radius = radius 244 | super(Sphere, self).__init__(Sphere.fibonacci_sphere(self._radius)) 245 | 246 | @property 247 | def radius(self): 248 | return self._radius 249 | 250 | @staticmethod 251 | def fibonacci_sphere(radius, samples=100): 252 | # From stackoverflow on How to evenly distribute N points on a sphere 253 | points = [] 254 | offset = 2./samples 255 | increment = math.pi * (3. - math.sqrt(5.)) 256 | 257 | # Point in the unit sphere 258 | for i in range(samples): 259 | y = ((i * offset) - 1) + (offset / 2) 260 | r = math.sqrt(1 - pow(y, 2)) 261 | 262 | phi = ((i + 1) % samples) * increment 263 | 264 | x = math.cos(phi) * r 265 | z = math.sin(phi) * r 266 | 267 | points.append([x, y, z]) 268 | 269 | return np.array(points).T * radius 270 | 271 | 272 | class Ellipsoid(ConvexShape): 273 | def __init__(self, a, b, c): 274 | super(Ellipsoid, self).__init__( 275 | self._create_points(a, b, c) 276 | ) 277 | 278 | def _create_points(self, a, b, c): 279 | theta = np.linspace(-np.pi/2, np.pi/2, 100) 280 | phi = np.linspace(-np.pi, np.pi, 100) 281 | theta, phi = np.meshgrid(theta, phi) 282 | x = a * np.cos(theta) * np.cos(phi) 283 | y = b * np.cos(theta) * np.sin(phi) 284 | z = c * np.sin(theta) 285 | points = np.stack([x, y, z]).reshape(3, -1) 286 | 287 | return points 288 | -------------------------------------------------------------------------------- /misc/visualize_distance_function.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os 4 | import sys 5 | 6 | import matplotlib 7 | matplotlib.use("agg") 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.mplot3d import Axes3D 10 | from matplotlib import cm 11 | 12 | import numpy as np 13 | 14 | def C(points, shapes, epsilons): 15 | assert points.shape[0] == 3 16 | assert shapes.shape[0] == 3 17 | assert epsilons.shape[0] == 2 18 | 19 | a1 = shapes[0] 20 | a2 = shapes[1] 21 | a3 = shapes[2] 22 | e1 = epsilons[0] 23 | e2 = epsilons[1] 24 | 25 | # zeros = points == 0 26 | # points[zeros] = points[zeros] + 1e-6 27 | 28 | F = ((points[0, :] / a1)**2.0)**(1.0/e2) 29 | F += ((points[1, :] / a2)**2.0)**(1.0/e2) 30 | F = F**(e2/e1) 31 | F += ((points[2, :] / a3)**2.0)**(1.0/e1) 32 | return F**e1 - 1.0 33 | 34 | 35 | 36 | def get_C(a1, a2, a3, e1, e2, val, plane="z"): 37 | if val > a3 and plane == "z": 38 | return 39 | elif val > a2 and plane == "y": 40 | return 41 | elif val > a3 and plane == "z": 42 | return 43 | x = np.linspace(-0.5, 0.5, 100) 44 | y = np.linspace(-0.5, 0.5, 100) 45 | xv, yv = np.meshgrid(x, y) 46 | if plane == "z": 47 | points = np.stack([ 48 | xv.ravel(), 49 | yv.ravel(), 50 | np.ones_like(xv.ravel())*val 51 | ]) 52 | elif plane == "y": 53 | points = np.stack([ 54 | xv.ravel(), 55 | np.ones_like(xv.ravel())*val, 56 | yv.ravel() 57 | ]) 58 | elif plane == "x": 59 | points = np.stack([ 60 | np.ones_like(xv.ravel())*val, 61 | xv.ravel(), 62 | yv.ravel() 63 | ]) 64 | z = C( 65 | points, 66 | np.array([[a1, a2, a3]]).ravel(), 67 | np.array([[e1, e2]]).ravel() 68 | ) 69 | return xv, yv, z 70 | 71 | 72 | def plot_C(a1, a2, a3, e1, e2, val, i): 73 | fig, (ax1, ax2, ax3) = plt.subplots(3, figsize=(12,12)) 74 | xv, yv, z = get_C(a1, a2, a3, e1, e2, val[0], "z") 75 | cs = ax1.contourf(xv, yv, z.reshape(100, 100), cmap=cm.PuBu_r, vmin=z.min(), vmax=z.max()) 76 | cbar = fig.colorbar(cs, ax=ax1) 77 | ax1.set_title("Z-plane at %.5f a1: %.2f, a2: %.2f, a3: %.2f, e1: %.2f, e2: %.2f" %( 78 | val[0], a1, a2, a3, e1, e2 79 | )) 80 | xv, yv, z = get_C(a1, a2, a3, e1, e2, val[1], "y") 81 | cs = ax2.contourf(xv, yv, z.reshape(100, 100), cmap=cm.PuBu_r, vmin=z.min(), vmax=z.max()) 82 | ax2.set_title("Y-plane at %.5f a1: %.2f, a2: %.2f, a3: %.2f, e1: %.2f, e2: %.2f" %( 83 | val[1], a1, a2, a3, e1, e2 84 | )) 85 | cbar = fig.colorbar(cs, ax=ax2) 86 | xv, yv, z = get_C(a1, a2, a3, e1, e2, val[1], "x") 87 | cs = ax3.contourf(xv, yv, z.reshape(100,100), cmap=cm.PuBu_r, vmin=z.min(), vmax=z.max()) 88 | ax3.set_title("X-plane at %.5f a1: %.2f, a2: %.2f, a3: %.2f, e1: %.2f, e2: %.2f" %( 89 | val[2], a1, a2, a3, e1, e2 90 | )) 91 | cbar = fig.colorbar(cs, ax=ax3) 92 | plt.savefig("/tmp/C_%03d.png" %(i)) 93 | # plt.show() 94 | plt.close() 95 | 96 | 97 | if __name__ == "__main__": 98 | N = 100 99 | planes = [0.004, 0.006, 0.003] 100 | a1s = np.random.random((N,))*0.5 + 1e-2 101 | a2s = np.random.random((N,))*0.5 + 1e-2 102 | a3s = np.random.random((N,))*0.5 + 1e-2 103 | e1s = np.random.random((N,))*1.6 + 0.2 104 | e2s = np.random.random((N,))*1.6 + 0.2 105 | 106 | for i, a1, a2, a3, e1, e2 in zip(range(N), a1s, a2s, a3s, e1s, e2s): 107 | print i, a1, a2, a3, e1, e2 108 | plot_C(a1, a2, a3, e1, e2, planes, i) 109 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch==0.4.1 2 | torchvision==0.1.8 3 | progress==1.4 4 | trimesh==2.38.42 5 | 6 | numpy 7 | scikit-learn 8 | Cython 9 | Pillow 10 | pyquaternion 11 | backports.functools_lru_cache 12 | sympy 13 | 14 | # Visualization 15 | matplotlib 16 | seaborn 17 | mayavi 18 | -------------------------------------------------------------------------------- /scripts/arguments.py: -------------------------------------------------------------------------------- 1 | 2 | def add_voxelizer_parameters(parser): 3 | parser.add_argument( 4 | "--voxelizer_factory", 5 | choices=[ 6 | "occupancy_grid", 7 | "tsdf_grid", 8 | "image" 9 | ], 10 | default="occupancy_grid", 11 | help="The voxelizer factory to be used (default=occupancy_grid)" 12 | ) 13 | 14 | parser.add_argument( 15 | "--grid_shape", 16 | type=lambda x: tuple(map(int, x.split(","))), 17 | default="32,32,32", 18 | help="The dimensionality of the voxel grid (default=(32, 32, 32)" 19 | ) 20 | parser.add_argument( 21 | "--save_voxels_to", 22 | default=None, 23 | help="Path to save the voxelised input to the network" 24 | ) 25 | parser.add_argument( 26 | "--image_shape", 27 | type=lambda x: tuple(map(int, x.split(","))), 28 | default="3,137,137", 29 | help="The dimensionality of the voxel grid (default=(3,137,137)" 30 | ) 31 | 32 | 33 | def add_training_parameters(parser): 34 | """Add arguments to a parser that are related with the training of the 35 | network. 36 | """ 37 | parser.add_argument( 38 | "--epochs", 39 | type=int, 40 | default=150, 41 | help="Number of times to iterate over the dataset (default=150)" 42 | ) 43 | parser.add_argument( 44 | "--steps_per_epoch", 45 | type=int, 46 | default=500, 47 | help=("Total number of steps (batches of samples) before declaring one" 48 | " epoch finished and starting the next epoch (default=500)") 49 | ) 50 | parser.add_argument( 51 | "--batch_size", 52 | type=int, 53 | default=32, 54 | help="Number of samples in a batch (default=32)" 55 | ) 56 | parser.add_argument( 57 | "--lr", 58 | type=float, 59 | default=1e-3, 60 | help="Learning rate (default 1e-3)" 61 | ) 62 | parser.add_argument( 63 | "--lr_epochs", 64 | type=lambda x: map(int, x.split(",")), 65 | default="500,1000,1500", 66 | help="Training epochs with diminishing learning rate" 67 | ) 68 | parser.add_argument( 69 | "--lr_factor", 70 | type=float, 71 | default=1.0, 72 | help=("Factor according to which the learning rate will be diminished" 73 | " (default=None)") 74 | ) 75 | parser.add_argument( 76 | "--optimizer", 77 | choices=["Adam", "SGD"], 78 | default="Adam", 79 | help="The optimizer to be used (default=Adam)" 80 | ) 81 | parser.add_argument( 82 | "--momentum", 83 | type=float, 84 | default=0.9, 85 | help=("Parameter used to update momentum in case of SGD optimizer" 86 | " (default=0.9)") 87 | ) 88 | 89 | 90 | def add_dataset_parameters(parser): 91 | parser.add_argument( 92 | "--dataset_type", 93 | default="shapenet_quad", 94 | choices=[ 95 | "shapenet_quad", 96 | "shapenet_v1", 97 | "shapenet_v2", 98 | "surreal_bodies", 99 | "dynamic_faust" 100 | ], 101 | help="The type of the dataset type to be used" 102 | ) 103 | parser.add_argument( 104 | "--n_points_from_mesh", 105 | type=int, 106 | default=1000, 107 | help="The maximum number of points sampled from mesh (default=1000)" 108 | ) 109 | parser.add_argument( 110 | "--model_tags", 111 | type=lambda x: x.split(":"), 112 | default=[], 113 | help="The tags to the model to be used for testing", 114 | ) 115 | 116 | 117 | def add_nn_parameters(parser): 118 | """Add arguments to control the design of the neural network architecture. 119 | """ 120 | parser.add_argument( 121 | "--architecture", 122 | choices=["tulsiani", "octnet", "resnet18"], 123 | default="tulsiani", 124 | help="Choose the architecture to train" 125 | ) 126 | parser.add_argument( 127 | "--train_with_bernoulli", 128 | action="store_true", 129 | help="Learn the Bernoulli priors during training" 130 | ) 131 | parser.add_argument( 132 | "--make_dense", 133 | action="store_true", 134 | help="When true use an additional FC before its regressor" 135 | ) 136 | 137 | 138 | def add_regularizer_parameters(parser): 139 | parser.add_argument( 140 | "--regularizer_type", 141 | choices=[ 142 | "bernoulli_regularizer", 143 | "entropy_bernoulli_regularizer", 144 | "parsimony_regularizer", 145 | "overlapping_regularizer", 146 | "sparsity_regularizer" 147 | ], 148 | nargs="+", 149 | default=[], 150 | help=("The type of the regularizer on the shapes to be used" 151 | " (default=None)") 152 | 153 | ) 154 | parser.add_argument( 155 | "--bernoulli_regularizer_weight", 156 | type=float, 157 | default=0.0, 158 | help=("The importance of the regularization term on Bernoulli priors" 159 | " (default=0.0)") 160 | ) 161 | parser.add_argument( 162 | "--maximum_number_of_primitives", 163 | type=int, 164 | default=5000, 165 | help=("The maximum number of primitives in the predicted shape " 166 | " (default=5000)") 167 | ) 168 | parser.add_argument( 169 | "--minimum_number_of_primitives", 170 | type=int, 171 | default=5, 172 | help=("The minimum number of primitives in the predicted shape " 173 | " (default=5)") 174 | ) 175 | parser.add_argument( 176 | "--entropy_bernoulli_regularizer_weight", 177 | type=float, 178 | default=0.0, 179 | help=("The importance of the regularizer term on the entropy of" 180 | " the bernoullis (default=0.0)") 181 | ) 182 | parser.add_argument( 183 | "--sparsity_regularizer_weight", 184 | type=float, 185 | default=0.0, 186 | help="The weight on the sparsity regularizer (default=0.0)" 187 | ) 188 | parser.add_argument( 189 | "--parsimony_regularizer_weight", 190 | type=float, 191 | default=0.0, 192 | help="The weight on the parsimony regularizer (default=0.0)" 193 | ) 194 | parser.add_argument( 195 | "--overlapping_regularizer_weight", 196 | type=float, 197 | default=0.0, 198 | help="The weight on the overlapping regularizer (default=0.0)" 199 | ) 200 | parser.add_argument( 201 | "--enable_regularizer_after_epoch", 202 | type=int, 203 | default=0, 204 | help="Epoch after which regularizer is enabled (default=10)" 205 | ) 206 | parser.add_argument( 207 | "--w1", 208 | type=float, 209 | default=0.005, 210 | help="The weight on the first term of the sparsity regularizer (default=0.005)" 211 | ) 212 | parser.add_argument( 213 | "--w2", 214 | type=float, 215 | default=0.005, 216 | help="The weight on the second term of the sparsity regularizer (default=0.005)" 217 | ) 218 | 219 | 220 | def add_sq_mesh_sampler_parameters(parser): 221 | parser.add_argument( 222 | "--D_eta", 223 | type=float, 224 | default=0.05, 225 | help="Step along the eta (default=0.05)" 226 | ) 227 | parser.add_argument( 228 | "--D_omega", 229 | type=float, 230 | default=0.05, 231 | help="Step along the omega (default=0.05)" 232 | ) 233 | parser.add_argument( 234 | "--n_points_from_sq_mesh", 235 | type=int, 236 | default=180, 237 | help="Number of points to sample from the mesh of the SQ (default=180)" 238 | ) 239 | 240 | 241 | def add_gaussian_noise_layer_parameters(parser): 242 | parser.add_argument( 243 | "--add_gaussian_noise", 244 | action="store_true", 245 | help="Add Gaussian noise in the layers" 246 | ) 247 | parser.add_argument( 248 | "--mu", 249 | type=float, 250 | default=0.0, 251 | help="Mean value of the Gaussian distribution" 252 | ) 253 | parser.add_argument( 254 | "--sigma", 255 | type=float, 256 | default=0.001, 257 | help="Standard deviation of the Gaussian distribution" 258 | ) 259 | 260 | 261 | def add_loss_parameters(parser): 262 | parser.add_argument( 263 | "--prim_to_pcl_loss_weight", 264 | default=1.0, 265 | type=float, 266 | help=("The importance of the primitive-to-pointcloud loss in the " 267 | "final loss (default = 1.0)") 268 | ) 269 | parser.add_argument( 270 | "--pcl_to_prim_loss_weight", 271 | default=1.0, 272 | type=float, 273 | help=("The importance of the pointcloud-to-primitive loss in the " 274 | "final loss (default = 1.0)") 275 | ) 276 | 277 | 278 | def add_loss_options_parameters(parser): 279 | parser.add_argument( 280 | "--use_sq", 281 | action="store_true", 282 | help="Use Superquadrics as geometric primitives" 283 | ) 284 | parser.add_argument( 285 | "--use_cuboids", 286 | action="store_true", 287 | help="Use cuboids as geometric primitives" 288 | ) 289 | parser.add_argument( 290 | "--use_chamfer", 291 | action="store_true", 292 | help="Use the chamfer distance" 293 | ) 294 | 295 | 296 | def voxelizer_shape(args): 297 | if args.voxelizer_factory == "occupancy_grid": 298 | return args.grid_shape 299 | elif args.voxelizer_factory == "image": 300 | return args.image_shape 301 | elif args.voxelizer_factory == "tsdf_grid": 302 | return (args.resolution,)*3 303 | 304 | 305 | def get_loss_weights(args): 306 | args = vars(args) 307 | loss_weights = { 308 | "pcl_to_prim_weight": args.get("pcl_to_prim_loss_weight", 1.0), 309 | "prim_to_pcl_weight": args.get("prim_to_pcl_loss_weight", 1.0), 310 | } 311 | 312 | return loss_weights 313 | 314 | 315 | def get_loss_options(args): 316 | loss_weights = get_loss_weights(args) 317 | 318 | args = vars(args) 319 | # Create a dicitionary with the loss options based on the input arguments 320 | loss_options = { 321 | "use_sq": args.get("use_sq", False), 322 | "use_cuboids": args.get("use_cuboids", False), 323 | "use_chamfer": args.get("use_chamfer", False), 324 | "loss_weights": loss_weights 325 | } 326 | 327 | return loss_options 328 | -------------------------------------------------------------------------------- /scripts/benchmark_sampler.py: -------------------------------------------------------------------------------- 1 | 2 | import time 3 | 4 | import numpy as np 5 | 6 | from learnable_primitives.fast_sampler._sampler import \ 7 | fast_sample_on_batch 8 | 9 | 10 | if __name__ == "__main__": 11 | n_runs = 10 12 | shapes = (np.ones((32, 15, 3))*[[[0.2, 0.1, 0.2]]]).astype(np.float32) 13 | epsilons = (np.ones((32, 15, 2))*[[[0.25, 0.25]]]).astype(np.float32) 14 | 15 | start = time.time() 16 | for i in range(n_runs): 17 | etas, omegas = fast_sample_on_batch(shapes, epsilons, 200) 18 | end = time.time() 19 | 20 | print "Time per sample", (end-start)/n_runs 21 | -------------------------------------------------------------------------------- /scripts/compute_chamfer_loss.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script used to compute the Chamfer distance on a set of pre-trained models 3 | """ 4 | import argparse 5 | import os 6 | import sys 7 | 8 | import numpy as np 9 | import torch 10 | from torch.utils.data import DataLoader 11 | from pyquaternion import Quaternion 12 | 13 | from arguments import add_voxelizer_parameters, add_nn_parameters,\ 14 | add_dataset_parameters, add_gaussian_noise_layer_parameters,\ 15 | voxelizer_shape, add_loss_options_parameters, \ 16 | add_loss_parameters, get_loss_options 17 | 18 | from learnable_primitives.common.dataset import get_dataset_type,\ 19 | compose_transformations 20 | from learnable_primitives.common.model_factory import DatasetBuilder 21 | from learnable_primitives.equal_distance_sampler_sq import\ 22 | EqualDistanceSamplerSQ 23 | from learnable_primitives.models import NetworkParameters 24 | from learnable_primitives.loss_functions import euclidean_dual_loss 25 | from learnable_primitives.voxelizers import VoxelizerFactory 26 | from learnable_primitives.utils.progbar import Progbar 27 | 28 | 29 | def main(argv): 30 | parser = argparse.ArgumentParser( 31 | description="Do the forward pass and estimate a set of primitives" 32 | ) 33 | parser.add_argument( 34 | "dataset_directory", 35 | help="Path to the directory containing the dataset" 36 | ) 37 | parser.add_argument( 38 | "output_directory", 39 | help="Save the output files in that directory" 40 | ) 41 | parser.add_argument( 42 | "--tsdf_directory", 43 | default="", 44 | help="Path to the directory containing the precomputed tsdf files" 45 | ) 46 | parser.add_argument( 47 | "--weight_file", 48 | default=None, 49 | help="The path to the previously trainined model to be used" 50 | ) 51 | 52 | parser.add_argument( 53 | "--n_primitives", 54 | type=int, 55 | default=32, 56 | help="Number of primitives" 57 | ) 58 | parser.add_argument( 59 | "--use_deformations", 60 | action="store_true", 61 | help="Use Superquadrics with deformations as the shape configuration" 62 | ) 63 | parser.add_argument( 64 | "--run_on_gpu", 65 | action="store_true", 66 | help="Use GPU" 67 | ) 68 | 69 | add_dataset_parameters(parser) 70 | add_nn_parameters(parser) 71 | add_voxelizer_parameters(parser) 72 | add_gaussian_noise_layer_parameters(parser) 73 | add_loss_parameters(parser) 74 | add_loss_options_parameters(parser) 75 | args = parser.parse_args(argv) 76 | 77 | # A sampler instance 78 | e = EqualDistanceSamplerSQ(200) 79 | 80 | # Check if output directory exists and if it doesn't create it 81 | if not os.path.exists(args.output_directory): 82 | os.makedirs(args.output_directory) 83 | 84 | if args.run_on_gpu and torch.cuda.is_available(): 85 | device = torch.device("cuda:0") 86 | else: 87 | device = torch.device("cpu") 88 | print "Running code on ", device 89 | 90 | # Create a factory that returns the appropriate voxelizer based on the 91 | # input argument 92 | voxelizer_factory = VoxelizerFactory( 93 | args.voxelizer_factory, 94 | np.array(voxelizer_shape(args)), 95 | args.save_voxels_to 96 | ) 97 | 98 | # Create a dataset instance to generate the samples for training 99 | dataset = get_dataset_type("euclidean_dual_loss")( 100 | (DatasetBuilder() 101 | .with_dataset(args.dataset_type) 102 | .filter_tags(args.model_tags) 103 | .build(args.dataset_directory)), 104 | voxelizer_factory, 105 | args.n_points_from_mesh, 106 | n_bbox=args.n_bbox, 107 | n_surface=args.n_surface, 108 | equal=args.equal, 109 | transform=compose_transformations(voxelizer_factory) 110 | ) 111 | 112 | # TODO: Change batch_size in dataloader 113 | dataloader = DataLoader(dataset, batch_size=1, num_workers=4) 114 | 115 | network_params = NetworkParameters.from_options(args) 116 | # Build the model to be used for testing 117 | model = network_params.network(network_params) 118 | # Move model to device to be used 119 | model.to(device) 120 | if args.weight_file is not None: 121 | # Load the model parameters of the previously trained model 122 | model.load_state_dict( 123 | torch.load(args.weight_file, map_location=device) 124 | ) 125 | model.eval() 126 | 127 | losses = [] 128 | pcl_to_prim_losses = [] 129 | prim_to_pcl_losses = [] 130 | 131 | prog = Progbar(len(dataloader)) 132 | i = 0 133 | for sample in dataloader: 134 | X, y_target = sample 135 | X, y_target = X.to(device), y_target.to(device) 136 | 137 | # Do the forward pass and estimate the primitive parameters 138 | y_hat = model(X) 139 | 140 | reg_terms = { 141 | "regularizer_type": [], 142 | "bernoulli_regularizer_weight": 0.0, 143 | "entropy_bernoulli_regularizer_weight": 0.0, 144 | "parsimony_regularizer_weight": 0.0, 145 | "overlapping_regularizer_weight": 0.0, 146 | "sparsity_regularizer_weight": 0.0, 147 | } 148 | loss, debug_stats = euclidean_dual_loss( 149 | y_hat, 150 | y_target, 151 | reg_terms, 152 | e, 153 | get_loss_options(args) 154 | ) 155 | 156 | if not np.isnan(loss.item()): 157 | losses.append(loss.item()) 158 | pcl_to_prim_losses.append(debug_stats["pcl_to_prim_loss"].item()) 159 | prim_to_pcl_losses.append(debug_stats["prim_to_pcl_loss"].item()) 160 | # Update progress bar 161 | prog.update(i+1) 162 | i += 1 163 | np.savetxt( 164 | os.path.join(args.output_directory, "losses.txt"), 165 | losses 166 | ) 167 | 168 | np.savetxt( 169 | os.path.join(args.output_directory, "pcl_to_prim_losses.txt"), 170 | pcl_to_prim_losses 171 | ) 172 | np.savetxt( 173 | os.path.join(args.output_directory, "prim_to_pcl_losses.txt"), 174 | prim_to_pcl_losses 175 | ) 176 | np.savetxt( 177 | os.path.join(args.output_directory, "mean_std_losses.txt"), 178 | [np.mean(losses), np.std(losses), 179 | np.mean(pcl_to_prim_losses), np.std(pcl_to_prim_losses), 180 | np.mean(prim_to_pcl_losses), np.std(prim_to_pcl_losses)] 181 | ) 182 | 183 | print "loss: %.7f +/- %.7f - pcl_to_prim_loss %.7f +/- %.7f - prim_to_pcl_loss %.7f +/- %.7f" %( 184 | np.mean(losses), 185 | np.std(losses), 186 | np.mean(pcl_to_prim_losses), 187 | np.std(pcl_to_prim_losses), 188 | np.mean(prim_to_pcl_losses), 189 | np.std(prim_to_pcl_losses) 190 | ) 191 | 192 | if __name__ == "__main__": 193 | main(sys.argv[1:]) 194 | -------------------------------------------------------------------------------- /scripts/forward_pass.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script used to perform a forward pass using a previously trained model and 3 | visualize the corresponding primitives 4 | """ 5 | import argparse 6 | import os 7 | import sys 8 | 9 | import numpy as np 10 | import torch 11 | from torch.utils.data import DataLoader 12 | 13 | from arguments import add_voxelizer_parameters, add_nn_parameters, \ 14 | add_dataset_parameters, add_gaussian_noise_layer_parameters, \ 15 | voxelizer_shape, add_loss_options_parameters, add_loss_parameters 16 | from utils import get_colors, store_primitive_parameters 17 | from visualization_utils import points_on_sq_surface, points_on_cuboid, \ 18 | save_prediction_as_ply 19 | 20 | from learnable_primitives.common.dataset import get_dataset_type,\ 21 | compose_transformations 22 | from learnable_primitives.common.model_factory import DatasetBuilder 23 | from learnable_primitives.equal_distance_sampler_sq import\ 24 | EqualDistanceSamplerSQ 25 | from learnable_primitives.models import NetworkParameters 26 | from learnable_primitives.loss_functions import euclidean_dual_loss 27 | from learnable_primitives.primitives import\ 28 | euler_angles_to_rotation_matrices, quaternions_to_rotation_matrices 29 | from learnable_primitives.voxelizers import VoxelizerFactory 30 | 31 | from mayavi import mlab 32 | 33 | 34 | def get_shape_configuration(use_cuboids): 35 | if use_cuboids: 36 | return points_on_cuboid 37 | else: 38 | return points_on_sq_surface 39 | 40 | 41 | def main(argv): 42 | parser = argparse.ArgumentParser( 43 | description="Do the forward pass and estimate a set of primitives" 44 | ) 45 | parser.add_argument( 46 | "dataset_directory", 47 | help="Path to the directory containing the dataset" 48 | ) 49 | parser.add_argument( 50 | "output_directory", 51 | help="Save the output files in that directory" 52 | ) 53 | parser.add_argument( 54 | "--tsdf_directory", 55 | default="", 56 | help="Path to the directory containing the precomputed tsdf files" 57 | ) 58 | parser.add_argument( 59 | "--weight_file", 60 | default=None, 61 | help="The path to the previously trainined model to be used" 62 | ) 63 | 64 | parser.add_argument( 65 | "--n_primitives", 66 | type=int, 67 | default=32, 68 | help="Number of primitives" 69 | ) 70 | parser.add_argument( 71 | "--prob_threshold", 72 | type=float, 73 | default=0.5, 74 | help="Probability threshold" 75 | ) 76 | parser.add_argument( 77 | "--use_deformations", 78 | action="store_true", 79 | help="Use Superquadrics with deformations as the shape configuration" 80 | ) 81 | parser.add_argument( 82 | "--save_prediction_as_mesh", 83 | action="store_true", 84 | help="When true store prediction as a mesh" 85 | ) 86 | parser.add_argument( 87 | "--run_on_gpu", 88 | action="store_true", 89 | help="Use GPU" 90 | ) 91 | parser.add_argument( 92 | "--with_animation", 93 | action="store_true", 94 | help="Add animation" 95 | ) 96 | 97 | add_dataset_parameters(parser) 98 | add_nn_parameters(parser) 99 | add_voxelizer_parameters(parser) 100 | add_gaussian_noise_layer_parameters(parser) 101 | add_loss_parameters(parser) 102 | add_loss_options_parameters(parser) 103 | args = parser.parse_args(argv) 104 | 105 | # A sampler instance 106 | e = EqualDistanceSamplerSQ(200) 107 | 108 | # Check if output directory exists and if it doesn't create it 109 | if not os.path.exists(args.output_directory): 110 | os.makedirs(args.output_directory) 111 | 112 | if args.run_on_gpu and torch.cuda.is_available(): 113 | device = torch.device("cuda:0") 114 | else: 115 | device = torch.device("cpu") 116 | print "Running code on ", device 117 | 118 | # Create a factory that returns the appropriate voxelizer based on the 119 | # input argument 120 | voxelizer_factory = VoxelizerFactory( 121 | args.voxelizer_factory, 122 | np.array(voxelizer_shape(args)), 123 | args.save_voxels_to 124 | ) 125 | 126 | # Create a dataset instance to generate the samples for training 127 | dataset = get_dataset_type("euclidean_dual_loss")( 128 | (DatasetBuilder() 129 | .with_dataset(args.dataset_type) 130 | .filter_tags(args.model_tags) 131 | .build(args.dataset_directory)), 132 | voxelizer_factory, 133 | args.n_points_from_mesh, 134 | transform=compose_transformations(voxelizer_factory) 135 | ) 136 | 137 | # TODO: Change batch_size in dataloader 138 | dataloader = DataLoader(dataset, batch_size=1, num_workers=4) 139 | 140 | network_params = NetworkParameters.from_options(args) 141 | # Build the model to be used for testing 142 | model = network_params.network(network_params) 143 | # Move model to device to be used 144 | model.to(device) 145 | if args.weight_file is not None: 146 | # Load the model parameters of the previously trained model 147 | model.load_state_dict( 148 | torch.load(args.weight_file, map_location=device) 149 | ) 150 | model.eval() 151 | 152 | colors = get_colors(args.n_primitives) 153 | for sample in dataloader: 154 | X, y_target = sample 155 | X, y_target = X.to(device), y_target.to(device) 156 | 157 | # Do the forward pass and estimate the primitive parameters 158 | y_hat = model(X) 159 | 160 | M = args.n_primitives # number of primitives 161 | probs = y_hat[0].to("cpu").detach().numpy() 162 | # Transform the Euler angles to rotation matrices 163 | if y_hat[2].shape[1] == 3: 164 | R = euler_angles_to_rotation_matrices( 165 | y_hat[2].view(-1, 3) 166 | ).to("cpu").detach() 167 | else: 168 | R = quaternions_to_rotation_matrices( 169 | y_hat[2].view(-1, 4) 170 | ).to("cpu").detach() 171 | # get also the raw quaternions 172 | quats = y_hat[2].view(-1, 4).to("cpu").detach().numpy() 173 | translations = y_hat[1].to("cpu").view(args.n_primitives, 3) 174 | translations = translations.detach().numpy() 175 | 176 | shapes = y_hat[3].to("cpu").view(args.n_primitives, 3).detach().numpy() 177 | epsilons = y_hat[4].to("cpu").view( 178 | args.n_primitives, 2 179 | ).detach().numpy() 180 | taperings = y_hat[5].to("cpu").view( 181 | args.n_primitives, 2 182 | ).detach().numpy() 183 | 184 | pts = y_target[:, :, :3].to("cpu") 185 | pts_labels = y_target[:, :, -1].to("cpu").squeeze().numpy() 186 | pts = pts.squeeze().detach().numpy().T 187 | 188 | on_prims = 0 189 | fig = mlab.figure(size=(400, 400), bgcolor=(1, 1, 1)) 190 | mlab.view(azimuth=0.0, elevation=0.0, distance=2) 191 | # Uncomment to visualize the points sampled from the target mesh 192 | # t = np.array([1.2, 0, 0]).reshape(3, -1) 193 | # pts_n = pts + t 194 | # mlab.points3d( 195 | # # pts_n[0], pts_n[1], pts_n[2], 196 | # pts[0], pts[1], pts[2], 197 | # scale_factor=0.03, color=(0.8, 0.8, 0.8) 198 | # ) 199 | 200 | # Keep track of the files containing the parameters of each primitive 201 | primitive_files = [] 202 | for i in range(args.n_primitives): 203 | x_tr, y_tr, z_tr, prim_pts =\ 204 | get_shape_configuration(args.use_cuboids)( 205 | shapes[i, 0], 206 | shapes[i, 1], 207 | shapes[i, 2], 208 | epsilons[i, 0], 209 | epsilons[i, 1], 210 | R[i].numpy(), 211 | translations[i].reshape(-1, 1), 212 | taperings[i, 0], 213 | taperings[i, 1] 214 | ) 215 | 216 | # Dump the parameters of each primitive as a dictionary 217 | store_primitive_parameters( 218 | size=tuple(shapes[i]), 219 | shape=tuple(epsilons[i]), 220 | rotation=tuple(quats[i]), 221 | location=tuple(translations[i]), 222 | tapering=tuple(taperings[i]), 223 | probability=(probs[0, i],), 224 | color=(colors[i % len(colors)]) + (1.0,), 225 | filepath=os.path.join( 226 | args.output_directory, 227 | "primitive_%d.p" %(i,) 228 | ) 229 | ) 230 | if probs[0, i] >= args.prob_threshold: 231 | on_prims += 1 232 | mlab.mesh( 233 | x_tr, 234 | y_tr, 235 | z_tr, 236 | color=tuple(colors[i % len(colors)]), 237 | opacity=1.0 238 | ) 239 | primitive_files.append( 240 | os.path.join(args.output_directory, "primitive_%d.p" % (i,)) 241 | ) 242 | 243 | if args.with_animation: 244 | cnt = 0 245 | for az in range(0, 360, 1): 246 | cnt += 1 247 | mlab.view(azimuth=az, elevation=0.0, distance=2) 248 | mlab.savefig( 249 | os.path.join( 250 | args.output_directory, 251 | "img_%04d.png" % (cnt,) 252 | ) 253 | ) 254 | for i in range(args.n_primitives): 255 | print i, probs[0, i] 256 | 257 | print "Using %d primitives out of %d" % (on_prims, args.n_primitives) 258 | mlab.show() 259 | 260 | if args.save_prediction_as_mesh: 261 | print "Saving prediction as mesh...." 262 | save_prediction_as_ply( 263 | primitive_files, 264 | os.path.join(args.output_directory, "primitives.ply") 265 | ) 266 | print "Saved prediction as ply file in {}".format( 267 | os.path.join(args.output_directory, "primitives.ply") 268 | ) 269 | 270 | 271 | if __name__ == "__main__": 272 | main(sys.argv[1:]) 273 | -------------------------------------------------------------------------------- /scripts/output_logger.py: -------------------------------------------------------------------------------- 1 | from progress.bar import Bar 2 | 3 | 4 | def chamfer_loss_logger(current_epoch, epochs, steps_per_epoch): 5 | bar = Bar( 6 | "Epoch %d/%d" % (current_epoch, epochs), 7 | suffix=("%(index)d/%(max)d - loss: %(loss).7f - " 8 | "pcl_to_prim: %(pcl_to_prim_loss).7f - " 9 | "prim_to_pcl: %(prim_to_pcl_loss).7f - " 10 | # "bern_reg: %(bernoulli_regularizer).8f " 11 | # "entr_bern_reg: %(entropy_bernoulli_regularizer).8f " 12 | # "sp_reg: %(sparsity_regularizer).8f " 13 | # "overl_reg: %(overlapping_regularizer).8f " 14 | # "parsimony_reg: %(parsimony_regularizer).8f " 15 | "exp_n_prims: %(exp_n_prims).4f" 16 | ""), 17 | max=steps_per_epoch, 18 | loss=0.0, 19 | pcl_to_prim_loss=0.0, prim_to_pcl_loss=0.0, 20 | sparsity_regularizer=0.0, 21 | entropy_bernoulli_regularizer=0.0, 22 | bernoulli_regularizer=0.0, 23 | parsimony_regularizer=0.0, 24 | overlapping_regularizer=0.0, 25 | exp_n_prims=0.0 26 | ) 27 | bar.hide_cursor = False 28 | return bar 29 | 30 | 31 | def get_logger(loss_type, current_epoch, epochs, steps_per_epoch): 32 | return { 33 | "euclidean_dual_loss": chamfer_loss_logger 34 | }[loss_type](current_epoch, epochs, steps_per_epoch) 35 | -------------------------------------------------------------------------------- /scripts/train_network.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script used to train the network for representing a 3D obkect as a set of 3 | primitives 4 | """ 5 | import argparse 6 | from itertools import izip 7 | import json 8 | import random 9 | import os 10 | import string 11 | import sys 12 | 13 | import numpy as np 14 | import torch 15 | 16 | from arguments import add_voxelizer_parameters, add_nn_parameters,\ 17 | add_dataset_parameters, add_training_parameters,\ 18 | add_regularizer_parameters, add_sq_mesh_sampler_parameters,\ 19 | add_gaussian_noise_layer_parameters, voxelizer_shape,\ 20 | add_loss_options_parameters, add_loss_parameters, get_loss_options 21 | from output_logger import get_logger 22 | from utils import parse_train_test_splits 23 | 24 | from learnable_primitives.common.dataset import get_dataset_type, \ 25 | compose_transformations 26 | from learnable_primitives.common.model_factory import DatasetBuilder 27 | from learnable_primitives.common.batch_provider import BatchProvider 28 | from learnable_primitives.equal_distance_sampler_sq import get_sampler 29 | from learnable_primitives.models import NetworkParameters, train_on_batch, \ 30 | optimizer_factory 31 | from learnable_primitives.loss_functions import euclidean_dual_loss 32 | from learnable_primitives.voxelizers import VoxelizerFactory 33 | 34 | 35 | def moving_average(prev_val, new_val, b): 36 | return (prev_val*b + new_val) / (b+1) 37 | 38 | 39 | def id_generator(size=6, chars=string.ascii_uppercase + string.digits): 40 | return ''.join(random.choice(chars) for _ in range(size)) 41 | 42 | 43 | def yield_infinite(iterable): 44 | while True: 45 | for item in iterable: 46 | yield item 47 | 48 | 49 | def lr_schedule(optimizer, current_epoch, init_lr, factor, reductions): 50 | def inner(epoch): 51 | for i, e in enumerate(reductions): 52 | if epoch < e: 53 | return init_lr*factor**(-i) 54 | return init_lr*factor**(-len(reductions)) 55 | 56 | for param_group in optimizer.param_groups: 57 | param_group['lr'] = inner(current_epoch) 58 | 59 | return optimizer 60 | 61 | 62 | def get_weight(w, epoch, current_epoch): 63 | if current_epoch < epoch: 64 | return 0.0 65 | else: 66 | return w 67 | 68 | 69 | def get_regularizer_terms(args, current_epoch): 70 | # Create a dictionary with the regularization terms if there are any 71 | regularizer_terms = { 72 | "regularizer_type": args.regularizer_type, 73 | "bernoulli_regularizer_weight": args.bernoulli_regularizer_weight, 74 | "entropy_bernoulli_regularizer_weight": 75 | args.entropy_bernoulli_regularizer_weight, 76 | "parsimony_regularizer_weight": args.parsimony_regularizer_weight, 77 | "sparsity_regularizer_weight": args.sparsity_regularizer_weight, 78 | "overlapping_regularizer_weight": args.overlapping_regularizer_weight, 79 | "minimum_number_of_primitives": args.minimum_number_of_primitives, 80 | "maximum_number_of_primitives": args.maximum_number_of_primitives, 81 | "w1": args.w1, 82 | "w2": args.w2 83 | } 84 | 85 | return regularizer_terms 86 | 87 | 88 | def save_experiment_params(args, experiment_tag, directory): 89 | t = vars(args) 90 | params = {k: str(v) for k, v in t.iteritems()} 91 | 92 | git_head_hash = "foo" 93 | params["git-commit"] = git_head_hash 94 | params["experiment_tag"] = experiment_tag 95 | for k, v in params.items(): 96 | if v == "": 97 | params[k] = None 98 | with open(os.path.join(directory, "params.json"), "w") as f: 99 | json.dump(params, f, indent=4) 100 | 101 | 102 | def main(argv): 103 | parser = argparse.ArgumentParser( 104 | description="Train a network to predict primitives" 105 | ) 106 | parser.add_argument( 107 | "dataset_directory", 108 | help="Path to the directory containing the dataset" 109 | ) 110 | parser.add_argument( 111 | "output_directory", 112 | help="Save the output files in that directory" 113 | ) 114 | 115 | parser.add_argument( 116 | "--tsdf_directory", 117 | default="", 118 | help="Path to the directory containing the precomputed tsdf files" 119 | ) 120 | parser.add_argument( 121 | "--weight_file", 122 | default=None, 123 | help=("The path to a previously trainined model to continue" 124 | " the training from") 125 | ) 126 | parser.add_argument( 127 | "--continue_from_epoch", 128 | default=0, 129 | type=int, 130 | help="Continue training from epoch (default=0)" 131 | ) 132 | parser.add_argument( 133 | "--n_primitives", 134 | type=int, 135 | default=32, 136 | help="Number of primitives" 137 | ) 138 | parser.add_argument( 139 | "--use_deformations", 140 | action="store_true", 141 | help="Use Superquadrics with deformations as the shape configuration" 142 | ) 143 | parser.add_argument( 144 | "--train_test_splits_file", 145 | default=None, 146 | help="Path to the train-test splits file" 147 | ) 148 | parser.add_argument( 149 | "--run_on_gpu", 150 | action="store_true", 151 | help="Use GPU" 152 | ) 153 | parser.add_argument( 154 | "--probs_only", 155 | action="store_true", 156 | help="Optimize only using the probabilities" 157 | ) 158 | 159 | parser.add_argument( 160 | "--experiment_tag", 161 | default=None, 162 | help="Tag that refers to the current experiment" 163 | ) 164 | 165 | parser.add_argument( 166 | "--cache_size", 167 | type=int, 168 | default=2000, 169 | help="The batch provider cache size" 170 | ) 171 | 172 | parser.add_argument( 173 | "--seed", 174 | type=int, 175 | default=27, 176 | help="Seed for the PRNG" 177 | ) 178 | 179 | add_nn_parameters(parser) 180 | add_dataset_parameters(parser) 181 | add_voxelizer_parameters(parser) 182 | add_training_parameters(parser) 183 | add_sq_mesh_sampler_parameters(parser) 184 | add_regularizer_parameters(parser) 185 | add_gaussian_noise_layer_parameters(parser) 186 | # Parameters related to the loss function and the loss weights 187 | add_loss_parameters(parser) 188 | # Parameters related to loss options 189 | add_loss_options_parameters(parser) 190 | args = parser.parse_args(argv) 191 | 192 | if args.train_test_splits_file is not None: 193 | train_test_splits = parse_train_test_splits( 194 | args.train_test_splits_file, 195 | args.model_tags 196 | ) 197 | training_tags = np.hstack([ 198 | train_test_splits["train"], 199 | train_test_splits["val"] 200 | ]) 201 | else: 202 | training_tags = args.model_tags 203 | 204 | #device = torch.device("cuda:0") 205 | if args.run_on_gpu: #and torch.cuda.is_available(): 206 | device = torch.device("cuda:0") 207 | else: 208 | device = torch.device("cpu") 209 | 210 | print "Running code on", device 211 | 212 | # Check if output directory exists and if it doesn't create it 213 | if not os.path.exists(args.output_directory): 214 | os.makedirs(args.output_directory) 215 | 216 | # Create an experiment directory using the experiment_tag 217 | if args.experiment_tag is None: 218 | experiment_tag = id_generator(9) 219 | else: 220 | experiment_tag = args.experiment_tag 221 | 222 | experiment_directory = os.path.join( 223 | args.output_directory, 224 | experiment_tag 225 | ) 226 | if not os.path.exists(experiment_directory): 227 | os.makedirs(experiment_directory) 228 | 229 | # Store the parameters for the current experiment in a json file 230 | save_experiment_params(args, experiment_tag, experiment_directory) 231 | print "Save experiment statistics in %s" %(experiment_tag, ) 232 | 233 | # Create two files to store the training and test evolution 234 | train_stats = os.path.join(experiment_directory, "train.txt") 235 | val_stats = os.path.join(experiment_directory, "val.txt") 236 | if args.weight_file is None: 237 | train_stats_f = open(train_stats, "w") 238 | else: 239 | train_stats_f = open(train_stats, "a+") 240 | train_stats_f.write(( 241 | "epoch loss pcl_to_prim_loss prim_to_pcl_loss bernoulli_regularizer " 242 | "entropy_bernoulli_regularizer parsimony_regularizer " 243 | "overlapping_regularizer sparsity_regularizer\n" 244 | )) 245 | 246 | # Set the random seed 247 | np.random.seed(args.seed) 248 | torch.manual_seed(np.random.randint(np.iinfo(np.int32).max)) 249 | if torch.cuda.is_available(): 250 | torch.cuda.manual_seed_all(np.random.randint(np.iinfo(np.int32).max)) 251 | 252 | # Create an object that will sample points in equal distances on the 253 | # surface of the primitive 254 | sampler = get_sampler( 255 | args.use_cuboids, 256 | args.n_points_from_sq_mesh, 257 | args.D_eta, 258 | args.D_omega 259 | ) 260 | 261 | # Create a factory that returns the appropriate voxelizer based on the 262 | # input argument 263 | voxelizer_factory = VoxelizerFactory( 264 | args.voxelizer_factory, 265 | np.array(voxelizer_shape(args)), 266 | args.save_voxels_to 267 | ) 268 | 269 | # Create a dataset instance to generate the samples for training 270 | training_dataset = get_dataset_type("euclidean_dual_loss")( 271 | (DatasetBuilder() 272 | .with_dataset(args.dataset_type) 273 | .lru_cache(2000) 274 | .filter_tags(training_tags) 275 | .build(args.dataset_directory)), 276 | voxelizer_factory, 277 | args.n_points_from_mesh, 278 | transform=compose_transformations(args.voxelizer_factory) 279 | ) 280 | # Create a batchprovider object to start generating batches 281 | train_bp = BatchProvider( 282 | training_dataset, 283 | batch_size=args.batch_size, 284 | cache_size=args.cache_size 285 | ) 286 | train_bp.ready() 287 | 288 | network_params = NetworkParameters.from_options(args) 289 | # Build the model to be used for training 290 | model = network_params.network(network_params) 291 | 292 | # Move model to the device to be used 293 | model.to(device) 294 | # Check whether there is a weight file provided to continue training from 295 | if args.weight_file is not None: 296 | model.load_state_dict(torch.load(args.weight_file)) 297 | model.train() 298 | 299 | # Build an optimizer object to compute the gradients of the parameters 300 | optimizer = optimizer_factory(args, model) 301 | 302 | # Loop over the dataset multiple times 303 | pcl_to_prim_losses = [] 304 | prim_to_pcl_losses = [] 305 | losses = [] 306 | for i in range(args.epochs): 307 | bar = get_logger( 308 | "euclidean_dual_loss", 309 | i+1, 310 | args.epochs, 311 | args.steps_per_epoch 312 | ) 313 | for b, sample in izip(range(args.steps_per_epoch), yield_infinite(train_bp)): 314 | X, y_target = sample 315 | X, y_target = X.to(device), y_target.to(device) 316 | 317 | # Train on batch 318 | batch_loss, metrics, debug_stats = train_on_batch( 319 | model, 320 | lr_schedule( 321 | optimizer, i, args.lr, args.lr_factor, args.lr_epochs 322 | ), 323 | euclidean_dual_loss, 324 | X, 325 | y_target, 326 | get_regularizer_terms(args, i), 327 | sampler, 328 | get_loss_options(args) 329 | ) 330 | 331 | # Get the regularizer terms 332 | reg_values = debug_stats["regularizer_terms"] 333 | sparsity_regularizer = reg_values["sparsity_regularizer"] 334 | overlapping_regularizer = reg_values["overlapping_regularizer"] 335 | parsimony_regularizer = reg_values["parsimony_regularizer"] 336 | entropy_bernoulli_regularizer = reg_values["entropy_bernoulli_regularizer"] 337 | bernoulli_regularizer = reg_values["bernoulli_regularizer"] 338 | 339 | # The lossess 340 | pcl_to_prim_loss = debug_stats["pcl_to_prim_loss"].item() 341 | prim_to_pcl_loss = debug_stats["prim_to_pcl_loss"].item() 342 | bar.loss = moving_average(bar.loss, batch_loss, b) 343 | bar.pcl_to_prim_loss = \ 344 | moving_average(bar.pcl_to_prim_loss, pcl_to_prim_loss, b) 345 | bar.prim_to_pcl_loss = \ 346 | moving_average(bar.prim_to_pcl_loss, prim_to_pcl_loss, b) 347 | 348 | losses.append(bar.loss) 349 | prim_to_pcl_losses.append(bar.prim_to_pcl_loss) 350 | pcl_to_prim_losses.append(bar.pcl_to_prim_loss) 351 | 352 | bar.bernoulli_regularizer =\ 353 | (bar.bernoulli_regularizer * b + bernoulli_regularizer) / (b+1) 354 | bar.parsimony_regularizer =\ 355 | (bar.parsimony_regularizer * b + parsimony_regularizer) / (b+1) 356 | bar.overlapping_regularizer =\ 357 | (bar.overlapping_regularizer * b + overlapping_regularizer) / (b+1) 358 | bar.entropy_bernoulli_regularizer = \ 359 | (bar.entropy_bernoulli_regularizer * b + 360 | entropy_bernoulli_regularizer) / (b+1) 361 | bar.sparsity_regularizer =\ 362 | (bar.sparsity_regularizer * b + sparsity_regularizer) / (b+1) 363 | 364 | bar.exp_n_prims = metrics[0].sum(-1).mean() 365 | # Update the file that keeps track of the statistics 366 | train_stats_f.write( 367 | ("%d %.8f %.8f %.8f %.6f %.6f %.6f %.6f %.6f") %( 368 | i, bar.loss, bar.pcl_to_prim_loss, bar.prim_to_pcl_loss, 369 | bar.bernoulli_regularizer, 370 | bar.entropy_bernoulli_regularizer, 371 | bar.parsimony_regularizer, 372 | bar.overlapping_regularizer, 373 | bar.sparsity_regularizer 374 | ) 375 | ) 376 | train_stats_f.write("\n") 377 | train_stats_f.flush() 378 | 379 | bar.next() 380 | # Finish the progress bar and save the model after every epoch 381 | bar.finish() 382 | # Stop the batch provider 383 | train_bp.stop() 384 | torch.save( 385 | model.state_dict(), 386 | os.path.join( 387 | experiment_directory, 388 | "model_%d" % (i + args.continue_from_epoch,) 389 | ) 390 | ) 391 | 392 | print [ 393 | sum(losses[args.steps_per_epoch:]) / float(args.steps_per_epoch), 394 | sum(losses[:args.steps_per_epoch]) / float(args.steps_per_epoch), 395 | sum(pcl_to_prim_losses[args.steps_per_epoch:]) / float(args.steps_per_epoch), 396 | sum(pcl_to_prim_losses[:args.steps_per_epoch]) / float(args.steps_per_epoch), 397 | sum(prim_to_pcl_losses[args.steps_per_epoch:]) / float(args.steps_per_epoch), 398 | sum(prim_to_pcl_losses[:args.steps_per_epoch]) / float(args.steps_per_epoch), 399 | ] 400 | 401 | 402 | if __name__ == "__main__": 403 | main(sys.argv[1:]) 404 | -------------------------------------------------------------------------------- /scripts/utils.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import pickle 3 | 4 | import matplotlib 5 | matplotlib.use("agg") 6 | import matplotlib.pyplot as plt 7 | import seaborn as sns 8 | sns.set() 9 | 10 | 11 | def parse_train_test_splits(train_test_splits_file, model_tags): 12 | splits = {} 13 | if not train_test_splits_file.endswith("csv"): 14 | raise Exception("Input file %s is not csv" % (train_test_splits_file,)) 15 | df = pd.read_csv( 16 | train_test_splits_file, 17 | names=["id", "synsetId", "subSynsetId", "modelId", "split"] 18 | ) 19 | keep_from_model = reduce( 20 | lambda a, x: a | (df["synsetId"] in x), 21 | model_tags, 22 | False 23 | ) 24 | # Keep only the rows from the model we want 25 | df_from_model = df[keep_from_model] 26 | 27 | train_idxs = df_from_model["split"] == "train" 28 | splits["train"] = df_from_model[train_idxs].modelId.values.tolist() 29 | test_idxs = df_from_model["split"] == "test" 30 | splits["test"] = df_from_model[test_idxs].modelId.values.tolist() 31 | val_idxs = df_from_model["split"] == "val" 32 | splits["val"] = df_from_model[val_idxs].modelId.values.tolist() 33 | 34 | return splits 35 | 36 | 37 | def get_colors(M): 38 | return sns.color_palette("Paired") 39 | 40 | 41 | def store_primitive_parameters( 42 | size, 43 | shape, 44 | rotation, 45 | location, 46 | tapering, 47 | probability, 48 | color, 49 | filepath 50 | ): 51 | primitive_params = dict( 52 | size=size, 53 | shape=shape, 54 | rotation=rotation, 55 | location=location, 56 | tapering=tapering, 57 | probability=probability, 58 | color=color 59 | ) 60 | pickle.dump( 61 | primitive_params, 62 | open(filepath, "wb") 63 | ) 64 | -------------------------------------------------------------------------------- /scripts/visualization_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pickle 3 | 4 | import trimesh 5 | from pyquaternion import Quaternion 6 | 7 | 8 | def fexp(x, p): 9 | return np.sign(x)*(np.abs(x)**p) 10 | 11 | 12 | def sq_surface(a1, a2, a3, e1, e2, eta, omega): 13 | x = a1 * fexp(np.cos(eta), e1) * fexp(np.cos(omega), e2) 14 | y = a2 * fexp(np.cos(eta), e1) * fexp(np.sin(omega), e2) 15 | z = a3 * fexp(np.sin(eta), e1) 16 | return x, y, z 17 | 18 | 19 | def points_on_sq_surface(a1, a2, a3, e1, e2, R, t, Kx, Ky, n_samples=100): 20 | """Computes a SQ given a set of parameters and saves it into a np array 21 | """ 22 | assert R.shape == (3, 3) 23 | assert t.shape == (3, 1) 24 | 25 | eta = np.linspace(-np.pi/2, np.pi/2, n_samples, endpoint=True) 26 | omega = np.linspace(-np.pi, np.pi, n_samples, endpoint=True) 27 | eta, omega = np.meshgrid(eta, omega) 28 | x, y, z = sq_surface(a1, a2, a3, e1, e2, eta, omega) 29 | 30 | # Apply the deformations 31 | fx = Kx * z / a3 32 | fx += 1 33 | fy = Ky * z / a3 34 | fy += 1 35 | fz = 1 36 | 37 | x = x * fx 38 | y = y * fy 39 | z = z * fz 40 | 41 | # Get an array of size 3x10000 that contains the points of the SQ 42 | points = np.stack([x, y, z]).reshape(3, -1) 43 | points_transformed = R.T.dot(points) + t 44 | # print "R:", R 45 | # print "t:", t 46 | # print "e:", [e1, e2] 47 | # print "K:", [Kx, Ky] 48 | 49 | x_tr = points_transformed[0].reshape(n_samples, n_samples) 50 | y_tr = points_transformed[1].reshape(n_samples, n_samples) 51 | z_tr = points_transformed[2].reshape(n_samples, n_samples) 52 | 53 | return x_tr, y_tr, z_tr, points_transformed 54 | 55 | 56 | def points_on_cuboid(a1, a2, a3, e1, e2, R, t, n_samples=100): 57 | """Computes a cube given a set of parameters and saves it into a np array 58 | """ 59 | assert R.shape == (3, 3) 60 | assert t.shape == (3, 1) 61 | 62 | X = np.array([ 63 | [0, 1, 1, 0, 0, 0, 0, 0, 0], 64 | [0, 1, 1, 0, 0, 1, 1, 1, 1] 65 | ], dtype=np.float32) 66 | X[X == 1.0] = a1 67 | X[X == 0.0] = -a1 68 | 69 | Y = np.array([ 70 | [0, 0, 0, 0, 0, 1, 1, 0, 0], 71 | [1, 1, 1, 1, 1, 1, 1, 0, 0] 72 | ], dtype=np.float32) 73 | Y[Y == 1.0] = a2 74 | Y[Y == 0.0] = -a2 75 | 76 | Z = np.array([ 77 | [1, 1, 0, 0, 1, 1, 0, 0, 1], 78 | [1, 1, 0, 0, 1, 1, 0, 0, 1] 79 | ], dtype=np.float32) 80 | Z[Z == 1.0] = a3 81 | Z[Z == 0.0] = -a3 82 | 83 | points = np.stack([X, Y, Z]).reshape(3, -1) 84 | points_transformed = R.T.dot(points) + t 85 | print "R:", R 86 | print "t:", t 87 | 88 | assert points.shape == (3, 18) 89 | 90 | x_tr = points_transformed[0].reshape(2, 9) 91 | y_tr = points_transformed[1].reshape(2, 9) 92 | z_tr = points_transformed[2].reshape(2, 9) 93 | return x_tr, y_tr, z_tr, points_transformed 94 | 95 | 96 | def _from_primitive_parms_to_mesh(primitive_params): 97 | if not isinstance(primitive_params, dict): 98 | raise Exception( 99 | "Expected dict and got {} as an input" 100 | .format(type(primitive_params)) 101 | ) 102 | # Extract the parameters of the primitives 103 | a1, a2, a3 = primitive_params["size"] 104 | e1, e2 = primitive_params["shape"] 105 | Kx, Ky = primitive_params["tapering"] 106 | t = np.array(primitive_params["location"]).reshape(3, 1) 107 | R = Quaternion(primitive_params["rotation"]).rotation_matrix.reshape(3, 3) 108 | 109 | # Sample points on the surface of its mesh 110 | _, _, _, V = points_on_sq_surface(a1, a2, a3, e1, e2, R, t, Kx, Ky) 111 | assert V.shape[0] == 3 112 | 113 | color = np.array(primitive_params["color"]) 114 | color = (color*255).astype(np.uint8) 115 | 116 | # Build a mesh object using the vertices loaded before and get its 117 | # convex hull 118 | m = trimesh.Trimesh(vertices=V.T).convex_hull 119 | # Apply color 120 | for i in range(len(m.faces)): 121 | m.visual.face_colors[i] = color 122 | 123 | return m 124 | 125 | 126 | def save_primitive_as_ply(primitive_params, filepath): 127 | m = _from_primitive_parms_to_mesh(primitive_params) 128 | # Make sure that the filepath endswith .obj 129 | if not filepath.endswith(".ply"): 130 | raise Exception( 131 | "The filepath should have an .ply suffix, instead we received {}" 132 | .format(filepath) 133 | ) 134 | m.export(filepath, file_type="ply") 135 | 136 | 137 | def save_prediction_as_ply(primitive_files, filepath): 138 | if not isinstance(primitive_files, list): 139 | raise Exception( 140 | "Expected list and got {} as an input" 141 | .format(type(primitive_files)) 142 | ) 143 | m = None 144 | for p in primitive_files: 145 | # Parse the primitive parameters 146 | prim_params = pickle.load(open(p, "r")) 147 | _m = _from_primitive_parms_to_mesh(prim_params) 148 | m = trimesh.util.concatenate(_m, m) 149 | 150 | m.export(filepath, file_type="ply") 151 | -------------------------------------------------------------------------------- /scripts/visualize_sq.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """Script for visualizing superquadrics.""" 3 | import argparse 4 | import os 5 | import sys 6 | 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | from mpl_toolkits.mplot3d import Axes3D 10 | from matplotlib import rc 11 | rc('text', usetex=True) 12 | plt.rc("font", size=10, family="serif") 13 | 14 | from arguments import add_sq_mesh_sampler_parameters 15 | 16 | from learnable_primitives.equal_distance_sampler_sq import \ 17 | EqualDistanceSamplerSQ 18 | from visualization_utils import sq_surface 19 | 20 | 21 | def main(argv): 22 | parser = argparse.ArgumentParser( 23 | description="Visualize superquadrics given a set of parameters" 24 | ) 25 | parser.add_argument( 26 | "--save_image_to", 27 | default=None, 28 | help="Path to save the generated image" 29 | ) 30 | parser.add_argument( 31 | "--size", 32 | type=lambda x: tuple(map(float, x.split(","))), 33 | default="0.55,0.55,0.55", 34 | help="Size of the superquadric" 35 | ) 36 | parser.add_argument( 37 | "--shape", 38 | type=lambda x: tuple(map(float, x.split(","))), 39 | default="0.35,0.35", 40 | help="Shape of the superquadric" 41 | ) 42 | parser.add_argument( 43 | "--n_points_from_sq_mesh", 44 | type=int, 45 | default=1000, 46 | help="Number of points to be sampled from the SQ" 47 | ) 48 | parser.add_argument( 49 | "--with_mesh", 50 | action="store_true", 51 | help="When true visualize the sampled points as a mesh" 52 | ) 53 | args = parser.parse_args(argv) 54 | 55 | # Create an object that will sample points in equal distances on the 56 | # surface of the primitive 57 | e = EqualDistanceSamplerSQ( 58 | args.n_points_from_sq_mesh 59 | ) 60 | a1, a2, a3 = args.size 61 | size = np.array([[[a1, a2, a3]]], dtype=np.float32) 62 | e1, e2, = args.shape 63 | shape = np.array([[[e1, e2]]], dtype=np.float32) 64 | 65 | etas, omegas = e.sample_on_batch(size, shape) 66 | x, y, z = sq_surface(a1, a2, a3, e1, e2, etas.ravel(), omegas.ravel()) 67 | 68 | fig = plt.figure() 69 | ax = fig.add_subplot(111, projection='3d') 70 | ax.scatter(x, y, z) 71 | ax.set_xlim([-0.65, 0.65]) 72 | ax.set_ylim([-0.65, 0.65]) 73 | ax.set_zlim([-0.65, 0.65]) 74 | ax.set_xlabel("x-axis") 75 | ax.set_ylabel("y-axis") 76 | ax.set_zlabel("z-axis") 77 | plt.title("Superquadric with size:(%0.3f, %0.3f, %0.3f) and shape:(%0.3f, %0.3f)" % (a1, a2, a3, e1, e2)) 78 | # Uncomment this if you want to save the SQ as png 79 | if args.save_image_to is not None: 80 | plt.subplots_adjust() 81 | plt.savefig(args.save_image_to) 82 | else: 83 | plt.show() 84 | 85 | if __name__ == "__main__": 86 | main(sys.argv[1:]) 87 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """Setup learnable_primitives""" 2 | 3 | from distutils.core import setup 4 | from Cython.Build import cythonize 5 | from distutils.extension import Extension 6 | 7 | from itertools import dropwhile 8 | import numpy as np 9 | from os import path 10 | 11 | 12 | def collect_docstring(lines): 13 | """Return document docstring if it exists""" 14 | lines = dropwhile(lambda x: not x.startswith('"""'), lines) 15 | doc = "" 16 | for line in lines: 17 | doc += line 18 | if doc.endswith('"""\n'): 19 | break 20 | 21 | return doc[3:-4].replace("\r", "").replace("\n", " ") 22 | 23 | 24 | def collect_metadata(): 25 | meta = {} 26 | with open(path.join("learnable_primitives", "__init__.py")) as f: 27 | lines = iter(f) 28 | meta["description"] = collect_docstring(lines) 29 | for line in lines: 30 | if line.startswith("__"): 31 | key, value = map(lambda x: x.strip(), line.split("=")) 32 | meta[key[2:-2]] = value[1:-1] 33 | 34 | return meta 35 | 36 | 37 | def get_extensions(): 38 | return cythonize([ 39 | Extension( 40 | "learnable_primitives.fast_sampler._sampler", 41 | [ 42 | "learnable_primitives/fast_sampler/_sampler.pyx", 43 | "learnable_primitives/fast_sampler/sampling.cpp" 44 | ], 45 | language="c++11", 46 | libraries=["stdc++"], 47 | include_dirs=[np.get_include()], 48 | extra_compile_args=["-std=c++11", "-O3"] 49 | ) 50 | ]) 51 | 52 | 53 | def get_install_requirements(): 54 | return [ 55 | "numpy", 56 | "scikit-learn", 57 | "trimesh==2.38.42", 58 | "torch==0.4.1", 59 | "torchvision==0.1.8", 60 | "progress==1.4", 61 | "cython", 62 | "Pillow", 63 | "pyquaternion", 64 | "backports.functools_lru_cache", 65 | "sympy", 66 | "matplotlib==2.2.4", 67 | "seaborn", 68 | "mayavi" 69 | ] 70 | 71 | 72 | def setup_package(): 73 | meta = collect_metadata() 74 | setup( 75 | name="learnable_primitives", 76 | version=meta["version"], 77 | maintainer=meta["maintainer"], 78 | maintainer_email=meta["email"], 79 | url=meta["url"], 80 | license=meta["license"], 81 | classifiers=[ 82 | "Intended Audience :: Science/Research", 83 | "Intended Audience :: Developers", 84 | "License :: OSI Approved :: MIT License", 85 | "Topic :: Scientific/Engineering", 86 | "Programming Language :: Python", 87 | "Programming Language :: Python :: 2", 88 | "Programming Language :: Python :: 2.7", 89 | ], 90 | install_requires=get_install_requirements(), 91 | ext_modules=get_extensions() 92 | ) 93 | 94 | 95 | if __name__ == "__main__": 96 | setup_package() 97 | --------------------------------------------------------------------------------