├── images ├── model.png ├── results.png └── pipeline.png ├── saved_models ├── norot_rnn1.pth ├── norot_rnn2.pth ├── norot_rnn3.pth ├── norot_rnn4.pth └── norot_rnn5.pth ├── requirements.txt ├── code ├── inference │ ├── obj2ply.py │ ├── ply2dsm.py │ ├── obj_utils.py │ ├── test.py │ └── plot.py └── core │ ├── hungarian_loss.py │ ├── datasets.py │ ├── pc_utils.py │ ├── train.py │ └── model.py ├── LICENSE ├── .gitignore └── README.md /images/model.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/images/model.png -------------------------------------------------------------------------------- /images/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/images/results.png -------------------------------------------------------------------------------- /images/pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/images/pipeline.png -------------------------------------------------------------------------------- /saved_models/norot_rnn1.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/saved_models/norot_rnn1.pth -------------------------------------------------------------------------------- /saved_models/norot_rnn2.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/saved_models/norot_rnn2.pth -------------------------------------------------------------------------------- /saved_models/norot_rnn3.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/saved_models/norot_rnn3.pth -------------------------------------------------------------------------------- /saved_models/norot_rnn4.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/saved_models/norot_rnn4.pth -------------------------------------------------------------------------------- /saved_models/norot_rnn5.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tejaskhot/primitive_fitting_3d/HEAD/saved_models/norot_rnn5.pth -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | scipy==1.10.1 2 | opencv_python==4.8.1.78 3 | numpy==1.24.3 4 | torch==2.2.0 5 | matplotlib==3.7.3 6 | plyfile==1.0.2 7 | Pillow==10.1.0 -------------------------------------------------------------------------------- /code/inference/obj2ply.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | 4 | 5 | def obj_to_ply(obj_file, ply_file): 6 | cmd = "meshlabserver -i {} -o {} ".format(obj_file, ply_file) 7 | os.system(cmd) 8 | 9 | 10 | for aoi_name in ["wpafb_d1", "wpafb_d2", "ucsd_d3", "jacksonville_d4"]: 11 | obj_primitives = glob.glob( 12 | os.path.join( 13 | "../../outputs/{}/obj_primitives".format(aoi_name), 14 | "{}*.obj".format(aoi_name), 15 | ) 16 | ) 17 | obj_overlayed = glob.glob( 18 | os.path.join( 19 | "../../outputs/{}/obj_overlayed".format(aoi_name), 20 | "{}*.obj".format(aoi_name), 21 | ) 22 | ) 23 | 24 | for i, f in enumerate(obj_primitives): 25 | obj_to_ply( 26 | f, "../../outputs/{}/ply_primitives/{}_{}.ply".format(aoi_name, aoi_name, i) 27 | ) 28 | 29 | for i, f in enumerate(obj_overlayed): 30 | obj_to_ply( 31 | f, "../../outputs/{}/ply_overlayed/{}_{}.ply".format(aoi_name, aoi_name, i) 32 | ) 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Tejas Khot 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | # data 107 | /data 108 | 109 | # outputs 110 | /outputs -------------------------------------------------------------------------------- /code/core/hungarian_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Variable 5 | from scipy.optimize import linear_sum_assignment 6 | 7 | 8 | class HungarianLoss(nn.Module): 9 | 10 | def __init__(self): 11 | super(HungarianLoss, self).__init__() 12 | self.use_cuda = torch.cuda.is_available() 13 | 14 | def forward(self, preds, gts, counts): 15 | batch_size = preds.shape[0] 16 | # P : (batch_size, gts.shape[1], preds.shape[1]) 17 | loss = [] 18 | for b in range(batch_size): 19 | count = counts[b] 20 | P = self.batch_pairwise_dist( 21 | gts[b, :count, :].unsqueeze(0), preds[b, :count, :].unsqueeze(0) 22 | ) 23 | cost = P[0].data.cpu().numpy() 24 | # get optimal assignments by Hungarian algo 25 | row_ind, col_ind = linear_sum_assignment(cost) 26 | # get the assignment indices but retrieve values from 27 | # the GPU tensor; this keeps things differentiable 28 | for i in range(len(row_ind)): 29 | loss.append(P[0, row_ind[i], col_ind[i]]) 30 | return torch.stack(loss).mean() 31 | 32 | def batch_pairwise_dist(self, x, y): 33 | bs, num_points_x, points_dim = x.size() 34 | _, num_points_y, _ = y.size() 35 | xx = torch.bmm(x, x.transpose(2, 1)) 36 | yy = torch.bmm(y, y.transpose(2, 1)) 37 | zz = torch.bmm(x, y.transpose(2, 1)) 38 | if self.use_cuda: 39 | dtype = torch.cuda.LongTensor 40 | else: 41 | dtype = torch.LongTensor 42 | diag_ind_x = torch.arange(0, num_points_x).type(dtype) 43 | diag_ind_y = torch.arange(0, num_points_y).type(dtype) 44 | 45 | rx = xx[:, diag_ind_x, diag_ind_x].unsqueeze(1).expand_as(zz.transpose(2, 1)) 46 | ry = yy[:, diag_ind_y, diag_ind_y].unsqueeze(1).expand_as(zz) 47 | P = rx.transpose(2, 1) + ry - 2 * zz 48 | return P 49 | 50 | 51 | if __name__ == "__main__": 52 | 53 | p1 = Variable(torch.rand(32, 6, 7).uniform_(-0.01, 0.01).cuda(), requires_grad=True) 54 | p2 = Variable(torch.rand(32, 5, 7).uniform_(0.3, 0.7).cuda()) 55 | 56 | hungarian_loss = HungarianLoss() 57 | counts = (torch.ones(32) * 4).type(torch.LongTensor) 58 | counts[:5] = 6 59 | counts[12:14] = 2 60 | counts = Variable(counts.cuda()) 61 | d = hungarian_loss(p1, p2, counts) 62 | print(d) 63 | d.backward() 64 | -------------------------------------------------------------------------------- /code/inference/ply2dsm.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | import sys 4 | 5 | sys.path.append("/usr0/home/tkhot/git_repos/core3d/") 6 | from core3d.mesh import mesh_util 7 | from core3d.texture_mapping import merge 8 | from core3d.texture_mapping.TextureMapper import TextureMapper 9 | import ipdb 10 | 11 | # for aoi_name in ['wpafb_d1', 'wpafb_d2','ucsd_d3', 'jacksonville_d4']: 12 | # print("Constructing DSM from Primitives for AOI : ", aoi_name) 13 | # ply_list = glob.glob(os.path.join('../../outputs/{}/ply_primitives','{}*.ply'.format(aoi_name, aoi_name))) 14 | # dtm_file = os.path.join('../../scoring/data/outputs/{}/buildings_dsm'.format(aoi_name), 'dtm.tif') 15 | # prim_dhm_file = os.path.join('../../scoring/data/outputs/{}/buildings_dsm'.format(aoi_name), 'dhm.tif') 16 | # prim_dsm_file = os.path.join('../../scoring/data/outputs/{}/buildings_dsm'.format(aoi_name), 'dsm.tif') 17 | # mesh_util.ply_list_to_dsm(ply_list, dtm_file, prim_dhm_file, prim_dsm_file) 18 | 19 | 20 | bldg_prim_dir = ( 21 | "/usr0/home/tkhot/Downloads/results/results/jacksonville_d4/buildings_prim" 22 | ) 23 | bridge_prim_dir = ( 24 | "/usr0/home/tkhot/Downloads/results/results/jacksonville_d4/bridge_prim" 25 | ) 26 | 27 | for aoi_name in ["wpafb_d1", "wpafb_d2", "ucsd_d3", "jacksonville_d4"]: 28 | 29 | # add textures 30 | # bldg_fit_dir = os.path.join(bldg_prim_dir, 'fitting_top_roof') 31 | bldg_fit_dir = os.path.join("../../outputs/{}/ply_primitives".format(aoi_name)) 32 | # bridge_fit_dir = os.path.join(bridge_prim_dir, 'fitting_top_roof') 33 | aoi_ply_file = os.path.join( 34 | "../../outputs/{}/ply_primitives".format(aoi_name), "aoi.ply" 35 | ) 36 | # scores_dir = os.path.join(bldg_fit_dir, 'scores') 37 | # os.makedirs(scores_dir, exist_ok=True) 38 | 39 | true_ortho_file = os.path.join( 40 | "../../scoring/data/outputs/{}/buildings_dsm".format(aoi_name), "true_ortho.tif" 41 | ) 42 | # aoi_ply_file = os.path.join(scores_dir, 'aoi.ply') 43 | # buildings_ply_file = os.path.join(scores_dir, 'buildings.ply') 44 | # bridges_ply_file = os.path.join(scores_dir, 'bridges.ply') 45 | textured_ply_file = os.path.join( 46 | "../../outputs/{}/ply_primitives".format(aoi_name), "aoi_textured.ply" 47 | ) 48 | 49 | print("Assembling 3D model") 50 | merge.merge_plys([bldg_fit_dir], aoi_ply_file) 51 | 52 | print("Texturing 3D model") 53 | texture_mapper = TextureMapper(aoi_ply_file, true_ortho_file) 54 | texture_mapper.save(textured_ply_file) 55 | 56 | print("Constructing DSM from Primitives") 57 | dtm_file = os.path.join( 58 | "../../scoring/data/outputs/{}/buildings_dsm".format(aoi_name), "dtm.tif" 59 | ) 60 | prim_dhm_file = os.path.join( 61 | "../../scoring/data/outputs/{}/buildings_dsm".format(aoi_name), "dhm.tif" 62 | ) 63 | prim_dsm_file = os.path.join( 64 | "../../scoring/data/outputs/{}/buildings_dsm".format(aoi_name), "dsm.tif" 65 | ) 66 | mesh_util.ply_list_to_dsm([aoi_ply_file], dtm_file, prim_dhm_file, prim_dsm_file) 67 | -------------------------------------------------------------------------------- /code/inference/obj_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | cube_v = np.array( 4 | [ 5 | [0.0, 0.0, 0.0], 6 | [0.0, 0.0, 1.0], 7 | [0.0, 1.0, 0.0], 8 | [0.0, 1.0, 1.0], 9 | [1.0, 0.0, 0.0], 10 | [1.0, 0.0, 1.0], 11 | [1.0, 1.0, 0.0], 12 | [1.0, 1.0, 1.0], 13 | ] 14 | ) 15 | cube_v = cube_v - 0.5 16 | 17 | cube_f = np.array( 18 | [ 19 | [1, 7, 5], 20 | [1, 3, 7], 21 | [1, 4, 3], 22 | [1, 2, 4], 23 | [3, 8, 7], 24 | [3, 4, 8], 25 | [5, 7, 8], 26 | [5, 8, 6], 27 | [1, 5, 6], 28 | [1, 6, 2], 29 | [2, 6, 8], 30 | [2, 8, 4], 31 | ] 32 | ).astype(np.int_) 33 | 34 | 35 | def append_obj(mf_handle, vertices, faces): 36 | for vx in range(vertices.shape[0]): 37 | mf_handle.write( 38 | "v {:f} {:f} {:f}\n".format( 39 | vertices[vx, 0], vertices[vx, 1], vertices[vx, 2] 40 | ) 41 | ) 42 | for fx in range(faces.shape[0]): 43 | mf_handle.write( 44 | "f {:d} {:d} {:d}\n".format(faces[fx, 0], faces[fx, 1], faces[fx, 2]) 45 | ) 46 | return 47 | 48 | 49 | def append_mtl_obj(mf_handle, vertices, faces, mtl_ids, mtl_name): 50 | mf_handle.write("mtllib {}\n".format(mtl_name)) 51 | for vx in range(vertices.shape[0]): 52 | mf_handle.write( 53 | "v {:f} {:f} {:f}\n".format( 54 | vertices[vx, 0], vertices[vx, 1], vertices[vx, 2] 55 | ) 56 | ) 57 | mts = np.unique(mtl_ids) 58 | for m in mts: 59 | faces_m = faces[mtl_ids == m] 60 | mf_handle.write("usemtl m{}\n".format(m)) 61 | for fx in range(faces_m.shape[0]): 62 | mf_handle.write( 63 | "f {:d} {:d} {:d}\n".format( 64 | faces_m[fx, 0], faces_m[fx, 1], faces_m[fx, 2] 65 | ) 66 | ) 67 | return 68 | 69 | 70 | def append_mtl(mtl_handle, mtl_ids, colors): 71 | for mx in range(len(mtl_ids)): 72 | mtl_handle.write("newmtl m{}\n".format(mtl_ids[mx])) 73 | # The Kd statement specifies the diffuse reflectivity using RGB values. 74 | mtl_handle.write( 75 | "Kd {:f} {:f} {:f}\n".format(colors[mx][0], colors[mx][1], colors[mx][2]) 76 | ) 77 | # The Ka statement specifies the ambient reflectivity using RGB values. 78 | mtl_handle.write("Ka 0 0 0\n") 79 | return 80 | 81 | 82 | def points_to_cubes(points, edge_size=0.05): 83 | """ 84 | Converts an input point cloud to a set of cubes. 85 | 86 | Args: 87 | points: N X 3 array 88 | edge_size: cube edge size 89 | Returns: 90 | vs: vertices 91 | fs: faces 92 | """ 93 | v_counter = 0 94 | tot_points = points.shape[0] 95 | v_all = np.tile(cube_v, [tot_points, 1]) 96 | f_all = np.tile(cube_f, [tot_points, 1]) 97 | f_offset = ( 98 | np.tile(np.linspace(0, 12 * tot_points - 1, 12 * tot_points), 3) 99 | .reshape(3, 12 * tot_points) 100 | .transpose() 101 | ) 102 | f_offset = (f_offset // 12 * 8).astype(np.int_) 103 | f_all += f_offset 104 | for px in range(points.shape[0]): 105 | v_all[v_counter : v_counter + 8, :] *= edge_size 106 | v_all[v_counter : v_counter + 8, :] += points[px, :] 107 | v_counter += 8 108 | 109 | return v_all, f_all 110 | 111 | 112 | def params_to_cubes(params): 113 | """ 114 | Generates a cube from a set of parameters 115 | 116 | Args: 117 | points: N X 3 array 118 | edge_size: cube edge size 119 | Returns: 120 | vs: vertices 121 | fs: faces 122 | """ 123 | tot_points = 1 124 | v_all = cube_v 125 | f_all = cube_f 126 | f_offset = ( 127 | np.tile(np.linspace(0, 12 * tot_points - 1, 12 * tot_points), 3) 128 | .reshape(3, 12 * tot_points) 129 | .transpose() 130 | ) 131 | f_offset = (f_offset // 12 * 8).astype(np.int_) 132 | 133 | S = np.diag(params[3:6] * 2) 134 | v_all = np.dot(v_all, S) 135 | v_all += params[:3] 136 | # print('======\n', f_all) 137 | return v_all, f_all + f_offset 138 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 3D Primitive Fitting to Point Clouds 2 | The goal for the 3D primitive fitting task is to represent buildings and other man-made structures as a collection of parameterized 3D volumetric primitives. The input for this task is a 3D point cloud with only (x,y,z) point tuples and no color information. The expected output is a collection of primitives (for now, cuboids) parameterized by their shape, location and orientation which is represented as a 7 dimensional vector: (width, height, depth, location_x, location_y, location_z, angle_theta). 3 | 4 | ### Overall Pipeline 5 | ![Overall Pipeline](images/pipeline.png) 6 | ### Model Architecture 7 | ![Model Architecture](images/model.png) 8 | ### Results on synthetic data 9 | ![Results on synthetic data](images/results.png) 10 | 11 | ## Getting Started 12 | 13 | The `requirements.txt` file lists the `Python` dependencies for this project. All code was tested using `Python 3.6.6` on an `Ubuntu 16.04` desktop machine. The code organization is as follows: 14 | * `data` : location for reading/writing data from; input is always a `ply` file containing a point cloud. The point cloud is always sampled to be `4096` points with only `(x,y,z)` coordinates. 15 | * `code` : location for code scripts 16 | * `inference` : location for scripts to test the trained model. 17 | * `saved_models` : location for storing pretrained models. 18 | * `images` : images depicting overall pipeline and model architecture 19 | 20 | In the `data` directory, we have: 21 | * `raw` : location for reading the `dhm.tif` and `ndx.tif` files. 22 | * `real` : location for storing the processed `raw` data for real buildings. These files would be used by the model for testing. 23 | * `synthetic` : location for storing samples of synthetic data. 24 | 25 | Within the `code` directory, we have: 26 | * `core` : location for storing common scripts used across the project. 27 | * `inference` : applying the trained model on new data. 28 | * `preprocessing` : preparing data for feeding to the model. 29 | 30 | Note: Only the code has been uploaded here and the structure above describes how to adapt the code for a dataset. 31 | 32 | ## Preprocessing 33 | 34 | **Training** 35 | Synthetic data for training is generated along with the ground truth parameter values using some internal code that is not released here. 36 | 37 | **Testing** 38 | The model consumes a point cloud of size `4096x3` as input which is read from a `ply` file. 39 | 40 | Each of these point clouds is still in the world coordinate frame. We bring the point clouds to a canonical unit cube during training and pass that point cloud to the model. This is done by centering the point cloud to origin and scaling it to make it's longest edge lie along an axis. These normalization steps are part of the `test.py` script and do not change files on disk. 41 | 42 | ## Running the model 43 | 44 | The main execution script is `code/inference/test.py`. This script takes several arguments as can be seen there. Some of the important choices here are: which trained model to use, how many primitives to predict amongst(max count), data and output path. We provide 5 trained models in the `saved_models` directory. These models are named as `norot_rnn*.pth` where `*` can be one of `[1,2,3,4,5]`. Here, the number indicates the maximum count of primitives the model was trained to predict. The models with lower counts give more reasonable outputs for a large number of cases at the cost of loss of detail. Models with more number of primitives capture detail in many cases, but are prone to wrong predictions because of the mistakes made by the counting branch of the model. If the predicted count is incorrect, the primitives are likely to be not a good fit. It is also possible to always predict a fixed number of primitives for each of these models by setting the appropriate flag `var_count=False`. The `norot` in names of these models indicates that they have been trained without having rotation as an output parameter. For all these models, the `num_params` is thus 6. From our experiments we found this to work fine since most instances of the buildings are aligned along some axis after the normalization step. Additionally, predicting rotation alongside leads to some odd cases of saddle points due to symmetries thereby missing out on some of the easier cases. 45 | 46 | This script saves a number of outputs to disk: 47 | * `npy` numpy array of parameters for the predicted primitives (in the canonical unit cube space) 48 | * `jpg` image of the point cloud and the primitive configuration 49 | * `obj` file storing the 3D primitives overlayed on the point cloud 50 | * `obj` file storing the 3D primitives 51 | * (optional) `gif` file of the animated plot of the point cloud and primitives 52 | 53 | While we do not perform any further optimization for fine-tuning the results of the trained model, we apply a simple verification step to remove predicted primitives that lie completely within another primitive. This can be further modified to allow only primitives which have a percentage overlap lower than a threshold, but we don't include code for that here (for now). 54 | 55 | ## Visualization 56 | 57 | The visualization script is called from within `test.py` and stores output files as mentioned above. These `obj` files can be viewed in any viewer application such as `MeshLab`. Optionally, one can convert these to `ply`. Note that the files stored to disk are named based on number of files in the corresponding folder. Also, the primitives in `obj` files are assigned face colors and hence one might need to change `MeshLab` setting if the colors don't show correctly. On Ubuntu, the option is `Render --> Color --> Per Face`. 58 | -------------------------------------------------------------------------------- /code/core/datasets.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | import torch 4 | import torch.utils.data as data 5 | from plyfile import PlyData 6 | import pc_utils 7 | import glob 8 | 9 | 10 | def random_jitter(pcd, sigma=0.02, clip=0.05): 11 | return pcd + np.clip(sigma * np.random.normal(size=pcd.shape), -clip, clip) 12 | 13 | 14 | def Rotate2D(pts, cnt, ang=np.pi / 4): 15 | """pts = {} Rotates points(nx2) about center cnt(2) by angle ang(1) in radian""" 16 | return ( 17 | np.dot( 18 | pts - cnt, 19 | np.array([[np.cos(ang), np.sin(ang)], [-np.sin(ang), np.cos(ang)]]), 20 | ) 21 | + cnt 22 | ) 23 | 24 | 25 | class BuildingDataloader(data.Dataset): 26 | def __init__( 27 | self, 28 | data_path, 29 | params_path, 30 | dtype, 31 | num_data=None, 32 | num_primitives=3, 33 | num_params=6, 34 | num_channels=3, 35 | batch_size=32, 36 | ): 37 | # this is a list of names of all pc files 38 | self.num_primitives = num_primitives 39 | self.num_channels = num_channels 40 | self.num_params = num_params 41 | self.data = [] 42 | for fname in glob.iglob(data_path + dtype + "/*.ply"): 43 | k = fname[len(data_path) + len(dtype) + 1 : -4] 44 | count = int(k.split("_")[1]) 45 | if count > num_primitives: 46 | continue 47 | self.data.append((k, fname)) 48 | self.data = self.data[: (len(self.data) // batch_size) * batch_size] 49 | # self.params is a dict with each key being the name of the pc file 50 | self.params = np.load(params_path).item() 51 | if num_data: 52 | self.data = self.data[:num_data] 53 | 54 | def __getitem__(self, index): 55 | k = self.data[index][0] 56 | fname = self.data[index][1] 57 | points = PlyData.read(fname) 58 | points = np.asarray(points["vertex"].data.tolist(), dtype=np.float32) 59 | idx = np.random.randint(points.shape[0], size=4096) 60 | points = points[idx, : self.num_channels] 61 | # add noise to the point cloud 62 | # points = random_jitter(points) 63 | points = points.astype(np.float32) 64 | 65 | center = (points.max(0) + points.min(0)) / 2 66 | points = points - center 67 | scale = np.max(points.max(0) - points.min(0)) / 2 68 | points = points / scale 69 | 70 | # points = np.random.choice(points,(4096,3),replace=False) 71 | params = self.params[k].astype(np.float32) 72 | params[:, :3] -= center 73 | params[:, :6] /= scale 74 | 75 | probs = len(params) - 1 76 | count = len(params) 77 | # add zero padding 78 | pad_size = self.num_primitives - len(params) 79 | params = np.pad(params, ((0, pad_size), (0, 0)), "constant") 80 | 81 | if self.num_params < 7: 82 | # get rid of rotations 83 | points[:, :2] = Rotate2D(points[:, :2], [0, 0], -params[0][6]) 84 | params[:, :2] = Rotate2D(params[:, :2], [0, 0], -params[0][6]) 85 | params = params[:, : self.num_params] 86 | 87 | return points, params, probs, count 88 | 89 | def __len__(self): 90 | return len(self.data) 91 | 92 | 93 | class BuildingGroupLoader(data.Dataset): 94 | def __init__( 95 | self, data_path, params_path, labels_path, dtype, num_data=None, num_channels=3 96 | ): 97 | # this is a list of names of all pc files 98 | self.data = [] 99 | self.num_channels = num_channels 100 | for fname in glob.iglob(data_path + dtype + "/*.ply"): 101 | k = fname[len(data_path) + len(dtype) + 1 : -4] 102 | self.data.append((k, fname)) 103 | self.labels = {} 104 | for fname in glob.iglob(labels_path + "/*.npy"): 105 | k = fname[len(labels_path) : -4] 106 | self.labels[k] = fname 107 | if num_data: 108 | self.data = self.data[:num_data] 109 | 110 | def __getitem__(self, index): 111 | k = self.data[index][0] 112 | fname = self.data[index][1] 113 | points = PlyData.read(fname) 114 | points = np.asarray(points["vertex"].data.tolist(), dtype=np.float32)[ 115 | :, : self.num_channels 116 | ] 117 | 118 | # add noise to the point cloud 119 | points = random_jitter(points) 120 | points = points.astype(np.float32) 121 | 122 | # group labels 123 | group_labels = np.load(self.labels[k]) 124 | pts_group_label, pts_group_mask, group_counts = ( 125 | pc_utils.convert_groupandcate_to_one_hot( 126 | group_labels, no_batch=True, num_groups=4 127 | ) 128 | ) 129 | 130 | return points, pts_group_label, group_counts 131 | 132 | def __len__(self): 133 | return len(self.data) 134 | 135 | 136 | if __name__ == "__main__": 137 | data_path = "../data/lod/pcd_ply/" 138 | params_path = "../data/lod/params.npy" 139 | dataloader = BuildingDataloader( 140 | data_path, params_path, num_primitives=3, dtype="train" 141 | ) 142 | # labels_path = '../data/instances/labels/' 143 | # dataloader = BuildingGroupLoader(data_path, params_path, labels_path, 'train') 144 | for i in range(10): 145 | points, params, probs, count = dataloader[0] 146 | # points, pts_group_label, group_counts = dataloader[i] 147 | print("len(dataloader) : {}".format(len(dataloader))) 148 | print("points.shape : {}".format(points.shape)) 149 | print("params.shape : {}".format(params.shape)) 150 | # print('pts_group_label.shape : {}'.format(pts_group_label.shape)) 151 | # print('group_counts.shape : {}'.format(group_counts.shape)) 152 | print("=" * 50) 153 | # print('probs.shape : {}'.format(probs.shape)) 154 | for i in range(10): 155 | points, params, probs, count = dataloader[i] 156 | print("len(dataloader) : {}".format(len(dataloader))) 157 | print("points.shape : {}".format(points.shape)) 158 | print("params.shape : {}".format(params.shape)) 159 | -------------------------------------------------------------------------------- /code/inference/test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | import random 5 | import glob 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | from torch.utils.data import DataLoader 10 | from torch.autograd import Variable 11 | import torch.nn.functional as F 12 | 13 | sys.path.append("../core/") 14 | 15 | from model import * 16 | from plot import * 17 | from plyfile import PlyData 18 | import ipdb 19 | 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--num_points", type=np.int_, default=4096, help="number of points") 22 | parser.add_argument( 23 | "--num_primitives", type=np.int_, default=5, help="number of cuboids" 24 | ) 25 | parser.add_argument( 26 | "--num_params", 27 | type=np.int_, 28 | default=6, 29 | help="number of parameters to predict per primitive", 30 | ) 31 | 32 | parser.add_argument( 33 | "--real", 34 | default=True, 35 | type=lambda x: (str(x).lower() == "true"), 36 | help="use real data", 37 | ) 38 | parser.add_argument( 39 | "--save", 40 | default=False, 41 | type=lambda x: (str(x).lower() == "true"), 42 | help="if false, matplotlib plots are displayed", 43 | ) 44 | parser.add_argument( 45 | "--var_count", 46 | default=True, 47 | type=lambda x: (str(x).lower() == "true"), 48 | help="use variable count prediction", 49 | ) 50 | parser.add_argument( 51 | "--cuda", default=True, type=lambda x: (str(x).lower() == "true"), help="use cuda" 52 | ) 53 | parser.add_argument( 54 | "--save_gif", 55 | default=False, 56 | type=lambda x: (str(x).lower() == "true"), 57 | help="save a gif", 58 | ) 59 | parser.add_argument( 60 | "--save_plot", 61 | default=False, 62 | type=lambda x: (str(x).lower() == "true"), 63 | help="save matplotlib plots", 64 | ) 65 | 66 | parser.add_argument("--aoi", type=str, default="", help="aoi name") 67 | parser.add_argument( 68 | "--syn_path", type=str, default="../../data/synthetic/", help="data path" 69 | ) 70 | parser.add_argument( 71 | "--params_path", 72 | type=str, 73 | default="../../data/params_synthetic.npy", 74 | help="params path", 75 | ) 76 | parser.add_argument( 77 | "--file_path", type=str, default="../../data/real/", help="data path" 78 | ) 79 | parser.add_argument("--save_path", type=str, default="../../outputs/", help="save path") 80 | 81 | parser.add_argument( 82 | "--model", 83 | type=str, 84 | default="../../saved_models/norot_rnn5.pth", 85 | help="trained model path", 86 | ) 87 | 88 | opt = parser.parse_args() 89 | print(opt) 90 | 91 | model = RLNet( 92 | num_points=opt.num_points, 93 | out_size=opt.num_params, 94 | num_primitives=opt.num_primitives, 95 | ) 96 | if opt.cuda: 97 | model.cuda() 98 | 99 | if opt.model != "": 100 | model.load_state_dict(torch.load(opt.model)) 101 | model.eval() 102 | 103 | if opt.real: 104 | data_folder = opt.file_path 105 | else: 106 | data_folder = opt.syn_path 107 | 108 | if len(opt.aoi) > 0: 109 | aoi = [opt.aoi] 110 | else: 111 | # aoi = ['wpafb_d1', 'wpafb_d2','ucsd_d3', 'jacksonville_d4'] 112 | aoi = ["COMMERCIALhotel_building_mesh0461"] 113 | 114 | for aoi_name in aoi: 115 | if not os.path.exists(os.path.join(opt.save_path, aoi_name)): 116 | os.makedirs(os.path.join(opt.save_path, aoi_name)) 117 | os.makedirs(os.path.join(opt.save_path, aoi_name, "gifs")) 118 | os.makedirs(os.path.join(opt.save_path, aoi_name, "images")) 119 | os.makedirs(os.path.join(opt.save_path, aoi_name, "obj_overlayed")) 120 | os.makedirs(os.path.join(opt.save_path, aoi_name, "obj_primitives")) 121 | os.makedirs(os.path.join(opt.save_path, aoi_name, "params")) 122 | os.makedirs(os.path.join(opt.save_path, aoi_name, "ply_primitives")) 123 | os.makedirs(os.path.join(opt.save_path, aoi_name, "ply_overlayed")) 124 | print("Processing : ", aoi_name) 125 | print(glob.glob(os.path.join(data_folder, "{}*.ply".format(aoi_name)))) 126 | for fname in glob.glob(os.path.join(data_folder, "{}*.ply".format(aoi_name))): 127 | points = PlyData.read(fname) 128 | print("Processing : ", fname) 129 | points = np.asarray(points["vertex"].data.tolist(), dtype=np.float32) 130 | indices = np.random.randint(points.shape[0], size=4096) 131 | points = points[indices, :3] 132 | # normalize the point cloud to fit in a unit cube with longest side being unit length 133 | center = (points.max(0) + points.min(0)) / 2 134 | points = points - center 135 | scale = np.max(points.max(0) - points.min(0)) / 2 136 | points = points / scale 137 | 138 | points = Variable(torch.Tensor(points)) 139 | points = points.unsqueeze(0) 140 | points = points.transpose(2, 1) 141 | if opt.cuda: 142 | points = points.cuda() 143 | 144 | pred, probs = model(points) 145 | # determine count based on stopping probability 146 | probs_int = (probs > 0.5).long() 147 | idx = torch.argmax(probs_int, 1) 148 | print("=" * 50) 149 | print("num primitives predicted : ", (idx + 1).data.cpu().numpy()[0]) 150 | 151 | pred = torch.reshape(pred, (pred.shape[0], opt.num_primitives, opt.num_params)) 152 | pred = pred.squeeze(0).detach().cpu().numpy() 153 | points = points.squeeze(0).transpose(0, 1).detach().cpu().numpy() 154 | if opt.var_count: 155 | pred = pred[: (idx[0] + 1)] 156 | 157 | # scale result to be in same world coordinates as original 158 | points = points * scale 159 | points = points + center 160 | pred[:, :6] *= scale 161 | pred[:, :3] += center 162 | 163 | # import pdb 164 | # pdb.set_trace() 165 | # make primitives touch the ground 166 | 167 | zmin = np.min(points, 0) 168 | for i in range(len(pred)): 169 | diff = pred[i, 2] - zmin[-1] 170 | pred[i, 2] -= diff 171 | 172 | # save params 173 | name_par = len( 174 | [ 175 | _ 176 | for _ in os.listdir(os.path.join(opt.save_path, aoi_name, "params")) 177 | if "npy" in _ 178 | ] 179 | ) 180 | np.save( 181 | os.path.join( 182 | opt.save_path, 183 | aoi_name, 184 | "params", 185 | "{}_{}.npy".format(aoi_name, name_par), 186 | ), 187 | np.array(pred), 188 | ) 189 | print( 190 | "params saved at : ", 191 | os.path.join( 192 | opt.save_path, "params", "{}_{}.npy".format(aoi_name, name_par) 193 | ), 194 | ) 195 | # save obj files 196 | plot_obj( 197 | points, 198 | pred, 199 | save_path=os.path.join(opt.save_path, aoi_name), 200 | aoi_name=aoi_name, 201 | ) 202 | if opt.save_plot: 203 | # save matplotlib plots 204 | draw_cube_points( 205 | pred, 206 | points, 207 | save=opt.save, 208 | save_path=os.path.join(opt.save_path, aoi_name), 209 | save_gif=opt.save_gif, 210 | aoi_name=aoi_name, 211 | ) 212 | -------------------------------------------------------------------------------- /code/core/pc_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import torch 5 | from torch.autograd import Variable 6 | 7 | NUM_CATEGORY = 13 8 | NUM_GROUPS = 3 9 | 10 | 11 | def var(x, cuda=True): 12 | if cuda: 13 | return Variable(torch.Tensor(x).cuda()) 14 | else: 15 | return Variable(torch.Tensor(x)) 16 | 17 | 18 | def countunique(A, Amax): 19 | res = np.empty(A.shape[1:], A.dtype) 20 | c = np.empty(Amax + 1, A.dtype) 21 | for i in range(A.shape[1]): 22 | for j in range(A.shape[2]): 23 | T = A[:, i, j] 24 | for k in range(c.size): 25 | c[k] = 0 26 | for x in T: 27 | c[x] = 1 28 | res[i, j] = c.sum() 29 | return res 30 | 31 | 32 | def exp_lr_scheduler( 33 | optimizer, global_step, init_lr, decay_steps, decay_rate, lr_clip, staircase=True 34 | ): 35 | """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs.""" 36 | if staircase: 37 | lr = init_lr * decay_rate ** (global_step // decay_steps) 38 | else: 39 | lr = init_lr * decay_rate ** (global_step / decay_steps) 40 | lr = max(lr, lr_clip) 41 | 42 | if global_step % decay_steps == 0: 43 | print("LR is set to {}".format(lr)) 44 | 45 | for param_group in optimizer.param_groups: 46 | param_group["lr"] = lr 47 | 48 | 49 | def convert_seg_to_one_hot(labels, no_batch=False): 50 | # labels:BxN 51 | if no_batch: 52 | labels = np.expand_dims(labels, axis=0) 53 | label_one_hot = np.zeros((labels.shape[0], labels.shape[1], NUM_CATEGORY)) 54 | pts_label_mask = np.zeros((labels.shape[0], labels.shape[1])) 55 | 56 | un, cnt = np.unique(labels, return_counts=True) 57 | label_count_dictionary = dict(zip(un, cnt)) 58 | totalnum = 0 59 | for k_un, v_cnt in label_count_dictionary.items(): 60 | if k_un != -1: 61 | totalnum += v_cnt 62 | 63 | for idx in range(labels.shape[0]): 64 | for jdx in range(labels.shape[1]): 65 | if labels[idx, jdx] != -1: 66 | label_one_hot[idx, jdx, labels[idx, jdx]] = 1 67 | pts_label_mask[idx, jdx] = float(totalnum) / float( 68 | label_count_dictionary[labels[idx, jdx]] 69 | ) # 1. - float(label_count_dictionary[labels[idx, jdx]]) / totalnum 70 | 71 | return label_one_hot.astype(np.float32), pts_label_mask.astype(np.float32) 72 | 73 | 74 | def convert_groupandcate_to_one_hot(grouplabels, no_batch=False, num_groups=NUM_GROUPS): 75 | # grouplabels: BxN 76 | if no_batch: 77 | grouplabels = np.expand_dims(grouplabels, axis=0) 78 | group_one_hot = np.zeros((grouplabels.shape[0], grouplabels.shape[1], num_groups)) 79 | pts_group_mask = np.zeros((grouplabels.shape[0], grouplabels.shape[1])) 80 | # group_counts = np.zeros((grouplabels.shape[0])) 81 | group_counts = [] 82 | 83 | un, cnt = np.unique(grouplabels, return_counts=True) 84 | group_count_dictionary = dict(zip(un, cnt)) 85 | totalnum = 0 86 | for k_un, v_cnt in group_count_dictionary.items(): 87 | if k_un != -1: 88 | totalnum += v_cnt 89 | 90 | for idx in range(grouplabels.shape[0]): 91 | un = np.unique(grouplabels[idx]) 92 | # group_counts[idx] = len(un) 93 | group_counts.append(len(un)) 94 | grouplabel_dictionary = dict(zip(un, range(len(un)))) 95 | for jdx in range(grouplabels.shape[1]): 96 | if grouplabels[idx, jdx] != -1: 97 | group_one_hot[ 98 | idx, jdx, grouplabel_dictionary[grouplabels[idx, jdx]] 99 | ] = 1 100 | pts_group_mask[idx, jdx] = float(totalnum) / float( 101 | group_count_dictionary[grouplabels[idx, jdx]] 102 | ) # 1. - float(group_count_dictionary[grouplabels[idx, jdx]]) / totalnum 103 | 104 | return ( 105 | group_one_hot.astype(np.float32), 106 | pts_group_mask.astype(np.float32), 107 | np.array(group_counts), 108 | ) 109 | 110 | 111 | def get_mask(labels, grouplabels): 112 | # labels:BxN 113 | # grouplabels: BxN 114 | pts_label_mask = np.zeros((labels.shape[0], labels.shape[1])) 115 | pts_group_mask = np.zeros((grouplabels.shape[0], grouplabels.shape[1])) 116 | ####### labels ####### 117 | un, cnt = np.unique(labels, return_counts=True) 118 | label_count_dictionary = dict(zip(un, cnt)) 119 | totalnum = 0 120 | for k_un, v_cnt in label_count_dictionary.items(): 121 | if k_un != -1: 122 | totalnum += v_cnt 123 | 124 | for idx in range(labels.shape[0]): 125 | for jdx in range(labels.shape[1]): 126 | if labels[idx, jdx] != -1: 127 | pts_label_mask[idx, jdx] = float(totalnum) / float( 128 | label_count_dictionary[labels[idx, jdx]] 129 | ) 130 | ####### group labels ####### 131 | un, cnt = np.unique(grouplabels, return_counts=True) 132 | group_count_dictionary = dict(zip(un, cnt)) 133 | totalnum = 0 134 | for k_un, v_cnt in group_count_dictionary.items(): 135 | if k_un != -1: 136 | totalnum += v_cnt 137 | 138 | for idx in range(grouplabels.shape[0]): 139 | un = np.unique(grouplabels[idx]) 140 | grouplabel_dictionary = dict(zip(un, range(len(un)))) 141 | for jdx in range(grouplabels.shape[1]): 142 | if grouplabels[idx, jdx] != -1: 143 | pts_group_mask[idx, jdx] = float(totalnum) / float( 144 | group_count_dictionary[grouplabels[idx, jdx]] 145 | ) 146 | 147 | return pts_label_mask, pts_group_mask 148 | 149 | 150 | def generate_group_mask(pts, grouplabels, labels): 151 | # grouplabels: BxN 152 | # pts: BxNx6 153 | # labels: BxN 154 | 155 | group_mask = np.zeros( 156 | (grouplabels.shape[0], grouplabels.shape[1], grouplabels.shape[1]) 157 | ) 158 | 159 | for idx in range(grouplabels.shape[0]): 160 | for jdx in range(grouplabels.shape[1]): 161 | for kdx in range(grouplabels.shape[1]): 162 | if labels[idx, jdx] == labels[idx, kdx]: 163 | group_mask[idx, jdx, kdx] = 2.0 164 | 165 | if ( 166 | np.linalg.norm( 167 | (pts[idx, jdx, :3] - pts[idx, kdx, :3]) 168 | * (pts[idx, jdx, :3] - pts[idx, kdx, :3]) 169 | ) 170 | < 0.04 171 | ): 172 | if labels[idx, jdx] == labels[idx, kdx]: 173 | group_mask[idx, jdx, kdx] = 5.0 174 | else: 175 | group_mask[idx, jdx, kdx] = 2.0 176 | 177 | return group_mask 178 | 179 | 180 | if __name__ == "__main__": 181 | input_list = "data/train_hdf5_file_list.txt" 182 | train_file_list = provider.getDataFiles(input_list) 183 | num_train_file = len(train_file_list) 184 | train_file_idx = np.arange(0, len(train_file_list)) 185 | np.random.shuffle(train_file_idx) 186 | ## load all data into memory 187 | all_data = [] 188 | all_group = [] 189 | all_seg = [] 190 | for i in range(num_train_file): 191 | cur_train_filename = train_file_list[train_file_idx[i]] 192 | cur_data, cur_group, _, cur_seg = ( 193 | provider.loadDataFile_with_groupseglabel_stanfordindoor(cur_train_filename) 194 | ) 195 | all_data += [cur_data] 196 | all_group += [cur_group] 197 | all_seg += [cur_seg] 198 | 199 | all_data = np.concatenate(all_data, axis=0) 200 | all_group = np.concatenate(all_group, axis=0) 201 | all_seg = np.concatenate(all_seg, axis=0) 202 | 203 | num_data = all_data.shape[0] 204 | -------------------------------------------------------------------------------- /code/core/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import random 4 | import numpy as np 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | from torch.utils.data import DataLoader 9 | from torch.autograd import Variable 10 | import torch.nn.functional as F 11 | from model import * 12 | from datasets import * 13 | from tensorboardX import SummaryWriter 14 | from hungarian_loss import * 15 | import time 16 | 17 | start = time.time() 18 | 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument( 21 | "--server", default=False, type=lambda x: (str(x).lower() == "true") 22 | ) 23 | parser.add_argument("--model_name", type=str, default="norot_rnn1", help="model name") 24 | 25 | parser.add_argument( 26 | "--num_data", type=int, default=None, help="number of input data instances" 27 | ) 28 | parser.add_argument("--num_points", type=int, default=4096, help="number of points") 29 | parser.add_argument("--num_primitives", type=int, default=3, help="number of cuboids") 30 | parser.add_argument( 31 | "--num_params", 32 | type=int, 33 | default=6, 34 | help="number of parameters to predict per primitive", 35 | ) 36 | parser.add_argument( 37 | "--num_channels", type=int, default=3, help="Number of input channels" 38 | ) 39 | parser.add_argument( 40 | "--num_faces", type=int, default=6, help="number of faces per primitive" 41 | ) 42 | parser.add_argument("--batch_size", type=int, default=32, help="input batch size") 43 | 44 | parser.add_argument( 45 | "--workers", type=int, default=10, help="number of data loading workers" 46 | ) 47 | parser.add_argument( 48 | "--nepoch", type=int, default=500, help="number of epochs to train for" 49 | ) 50 | 51 | parser.add_argument("--lr", type=float, default=1e-3, help="learning rate") 52 | parser.add_argument("--alpha", type=float, default=0.1, help="alpha for ce loss weight") 53 | parser.add_argument( 54 | "--grad_norm_limit", type=float, default=1, help="grad norm clipping limit" 55 | ) 56 | 57 | parser.add_argument( 58 | "--save_freq", 59 | type=int, 60 | default=50, 61 | help="save model to file after these many epochs", 62 | ) 63 | parser.add_argument( 64 | "--print_freq", 65 | type=int, 66 | default=10, 67 | help="print intermediate losses after these many epochs", 68 | ) 69 | 70 | parser.add_argument( 71 | "--data_path", type=str, default="/scratch/tkhot/lod/pcd_ply/", help="data path" 72 | ) 73 | parser.add_argument( 74 | "--params_path", 75 | type=str, 76 | default="/scratch/tkhot/lod/params.npy", 77 | help="params path", 78 | ) 79 | parser.add_argument("--outf", type=str, default="", help="output folder") 80 | parser.add_argument("--model", type=str, default="", help="model path") 81 | parser.add_argument("--logs", type=str, default="", help="logs path") 82 | parser.add_argument("--imdir", type=str, default="", help="logs path") 83 | 84 | 85 | opt = parser.parse_args() 86 | green = lambda x: "\033[92m" + x + "\033[00m" 87 | print(opt) 88 | 89 | if not opt.outf: 90 | opt.outf = "../outputs/saved_models" 91 | if not opt.title: 92 | opt.title = opt.model_name 93 | if not opt.logs: 94 | opt.logs = "../outputs/" 95 | if not opt.imdir: 96 | opt.imdir = "../plots/{}".format(opt.model_name) 97 | if opt.num_data: 98 | if opt.num_data < opt.batch_size: 99 | opt.batch_size = opt.num_data 100 | 101 | writer = SummaryWriter(opt.logs + opt.model_name) 102 | 103 | # opt.manualSeed = random.randint(1, 10000) # fix seed 104 | opt.manualSeed = 1 105 | print("Random Seed: ", opt.manualSeed) 106 | random.seed(opt.manualSeed) 107 | torch.manual_seed(opt.manualSeed) 108 | 109 | # DECAY_STEP = 800000. 110 | DECAY_STEP = 15000 111 | DECAY_RATE = 0.1 112 | 113 | dataset_train = BuildingDataloader( 114 | data_path=opt.data_path, 115 | params_path=opt.params_path, 116 | dtype="train", 117 | num_data=opt.num_data, 118 | num_params=opt.num_params, 119 | num_channels=opt.num_channels, 120 | num_primitives=opt.num_primitives, 121 | batch_size=opt.batch_size, 122 | ) 123 | dataloader_train = torch.utils.data.DataLoader( 124 | dataset_train, 125 | batch_size=opt.batch_size, 126 | shuffle=True, 127 | num_workers=int(opt.workers), 128 | pin_memory=True, 129 | ) 130 | dataset_test = BuildingDataloader( 131 | data_path=opt.data_path, 132 | params_path=opt.params_path, 133 | dtype="test", 134 | num_data=opt.num_data, 135 | num_params=opt.num_params, 136 | num_channels=opt.num_channels, 137 | num_primitives=opt.num_primitives, 138 | batch_size=opt.batch_size, 139 | ) 140 | dataloader_test = torch.utils.data.DataLoader( 141 | dataset_test, 142 | batch_size=opt.batch_size, 143 | shuffle=True, 144 | num_workers=int(opt.workers), 145 | pin_memory=True, 146 | ) 147 | 148 | print("len(dataset_train) : ", len(dataset_train)) 149 | print("len(dataset_test) : ", len(dataset_test)) 150 | try: 151 | os.makedirs(opt.outf) 152 | except OSError: 153 | pass 154 | 155 | model = RLNet( 156 | num_points=opt.num_points, 157 | out_size=opt.num_params, 158 | num_primitives=opt.num_primitives, 159 | ) 160 | model.cuda() 161 | 162 | if opt.model != "": 163 | model.load_state_dict(torch.load(opt.model)) 164 | 165 | optimizer = optim.Adam(model.parameters(), lr=opt.lr, amsgrad=True) 166 | 167 | criterion = GeometricLoss( 168 | batch_size=1, 169 | samples_per_face=50, 170 | num_faces=opt.num_faces, 171 | num_params=opt.num_params, 172 | ) 173 | criterion.cuda() 174 | 175 | ce_crit = nn.CrossEntropyLoss() 176 | hungarian_loss = HungarianLoss() 177 | hungarian_loss.cuda() 178 | 179 | 180 | def adjust_lr(optim): 181 | for g in optim.param_groups: 182 | g["lr"] = 0.75 * g["lr"] 183 | print("lr changed to {}", g["lr"]) 184 | 185 | 186 | print("-" * 50) 187 | alpha = 0.5 188 | for epoch in range(opt.nepoch): 189 | 190 | if epoch > 0: 191 | print(green("[{}] Loss : {}".format(epoch + 1, np.mean(total_losses)))) 192 | if ((epoch + 1) % 75) == 0: 193 | adjust_lr(optimizer) 194 | total_losses = [] 195 | for i, data in enumerate(dataloader_train, 0): 196 | points, gt_params, gt_probs, gt_counts = data 197 | points, gt_params, gt_probs, gt_counts = ( 198 | Variable(points), 199 | Variable(gt_params), 200 | Variable(gt_probs.type(torch.LongTensor)), 201 | Variable(gt_counts.type(torch.LongTensor)), 202 | ) 203 | points = points.transpose(2, 1) 204 | points, gt_params, gt_probs, gt_counts = ( 205 | points.cuda(), 206 | gt_params.cuda(), 207 | gt_probs.cuda(), 208 | gt_counts.cuda(), 209 | ) 210 | 211 | del data 212 | 213 | optimizer.zero_grad() 214 | pred, probs = model(points) 215 | out = torch.reshape(pred, (opt.batch_size, opt.num_primitives, opt.num_params)) 216 | mse_loss = hungarian_loss(out, gt_params, gt_counts) 217 | ce_loss = ce_crit(probs, gt_probs) 218 | total_loss = mse_loss + opt.alpha * ce_loss 219 | total_loss.backward() 220 | if opt.grad_norm_limit: 221 | nn.utils.clip_grad_norm_(model.parameters(), opt.grad_norm_limit) 222 | optimizer.step() 223 | total_losses.append(total_loss.data) 224 | 225 | writer.add_scalar("mse_loss", mse_loss, epoch * len(dataloader_train) + i) 226 | writer.add_scalar("ce_loss", ce_loss, epoch * len(dataloader_train) + i) 227 | writer.add_scalar("total_loss", total_loss, epoch * len(dataloader_train) + i) 228 | # for name, param in model.named_parameters(): 229 | # if param.grad is not None: 230 | # writer.add_scalar(name+"_grad", torch.norm(param.grad), epoch*len(dataloader_train)+i) 231 | 232 | # free up some space -- dataloader first allocates memory and then assigns; so deleting ensures we don't need double the memory 233 | del points, gt_params, gt_probs, pred, probs 234 | 235 | if (epoch + 1) % opt.save_freq == 0: 236 | torch.save( 237 | model.state_dict(), 238 | "{}/{}_{}.pth".format(opt.outf, opt.model_name, epoch + 1), 239 | ) 240 | 241 | if opt.nepoch > 0: 242 | torch.save( 243 | model.state_dict(), "{}/{}_{}.pth".format(opt.outf, opt.model_name, epoch + 1) 244 | ) 245 | -------------------------------------------------------------------------------- /code/core/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import torch.nn.functional as F 5 | 6 | 7 | class PointNetfeat(nn.Module): 8 | def __init__(self, num_points=2500, input_dim=3, feat_size=1024): 9 | super(PointNetfeat, self).__init__() 10 | self.feat_size = feat_size 11 | self.conv1 = torch.nn.Conv1d(input_dim, 64, 1) 12 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 13 | self.conv3 = torch.nn.Conv1d(128, self.feat_size, 1) 14 | self.mp1 = torch.nn.MaxPool1d(num_points) 15 | self.num_points = num_points 16 | 17 | def forward(self, x): 18 | x = F.leaky_relu(self.conv1(x)) 19 | x = F.leaky_relu(self.conv2(x)) 20 | x = self.conv3(x) 21 | x = self.mp1(x) 22 | x = x.view(-1, self.feat_size) 23 | return x 24 | 25 | 26 | class PointNet(nn.Module): 27 | def __init__(self, num_points=4096, num_channels=3): 28 | super(PointNet, self).__init__() 29 | self.num_points = num_points 30 | # conv 31 | self.conv1 = torch.nn.Conv1d(num_channels, 64, 1) 32 | self.conv2 = torch.nn.Conv1d(64, 64, 1) 33 | self.conv3 = torch.nn.Conv1d(64, 64, 1) 34 | self.conv4 = torch.nn.Conv1d(64, 128, 1) 35 | self.conv5 = torch.nn.Conv1d(128, 1024, 1) 36 | # max pool 37 | self.mp1 = torch.nn.MaxPool1d(num_points) 38 | # fully connected 39 | self.fc1 = nn.Linear(1024, 256) 40 | self.fc2 = nn.Linear(256, 128) 41 | # conv 42 | self.conv6 = torch.nn.Conv1d(1152, 512, 1) 43 | self.conv7 = torch.nn.Conv1d(512, 256, 1) 44 | 45 | def forward(self, x): 46 | # conv -- bn -- leaky_relu 47 | x = F.leaky_relu(self.conv1(x)) 48 | x = F.leaky_relu(self.conv2(x)) 49 | x = F.leaky_relu(self.conv3(x)) 50 | x = F.leaky_relu(self.conv4(x)) 51 | x = F.leaky_relu(self.conv5(x)) 52 | global_feat = x 53 | # max pool 54 | x = self.mp1(x).squeeze(-1) 55 | # FC layers -- bn -- leaky_relu 56 | x = F.leaky_relu(self.fc1(x)) 57 | x = F.leaky_relu(self.fc2(x)) 58 | # concat 59 | x_expand = x.unsqueeze(-1).repeat(1, 1, self.num_points) 60 | out = torch.cat([global_feat, x_expand], 1) 61 | # conv -- bn -- leaky_relu 62 | out = F.leaky_relu(self.conv6(out)) 63 | out = F.leaky_relu(self.conv7(out)) 64 | return out 65 | 66 | 67 | class PrimitiveNet(nn.Module): 68 | def __init__( 69 | self, num_points=2500, input_dim=3, feat_size=1024, hidden_size=256, out_size=10 70 | ): 71 | super(PrimitiveNet, self).__init__() 72 | self.num_points = num_points 73 | self.input_dim = input_dim 74 | self.feat_size = feat_size 75 | self.hidden_size = hidden_size 76 | self.out_size = out_size 77 | self.feat = PointNetfeat(self.num_points, self.input_dim, self.feat_size) 78 | self.fc1 = nn.Linear(self.feat_size, self.hidden_size) 79 | self.fc2 = nn.Linear(self.hidden_size, self.out_size) 80 | 81 | def forward(self, x): 82 | x = self.feat(x) 83 | x = F.leaky_relu(self.fc1(x)) 84 | x = self.fc2(x) 85 | return x 86 | 87 | 88 | class MetricModel(nn.Module): 89 | def __init__(self, num_points, num_channels, feat_size=128): 90 | super(MetricModel, self).__init__() 91 | self.pointnet = PointNet(num_points, num_channels) 92 | # feature convs 93 | self.conv1 = torch.nn.Conv1d(256, 128, 1) 94 | self.conv2 = torch.nn.Conv1d(256, feat_size, 1) 95 | 96 | def forward(self, x, seg_only=False): 97 | # pointnet feature extraction 98 | x = self.pointnet(x) 99 | 100 | # embedding features 101 | xfeat1 = self.conv2(x) 102 | return {"feat": xfeat1.permute(0, 2, 1)} 103 | 104 | 105 | class RLNet(nn.Module): 106 | def __init__( 107 | self, 108 | num_points=2048, 109 | input_dim=3, 110 | feat_size=1024, 111 | hidden_size=256, 112 | out_size=7, 113 | num_primitives=3, 114 | encoder=None, 115 | ): 116 | super(RLNet, self).__init__() 117 | self.num_points = num_points 118 | self.input_dim = input_dim 119 | self.feat_size = feat_size 120 | self.hidden_size = hidden_size 121 | self.out_size = out_size 122 | self.num_primitives = num_primitives 123 | if encoder is not None: 124 | self.encoder = nn.Sequential( 125 | PointNetfeat(self.num_points), 126 | nn.Linear(self.feat_size, 256), 127 | nn.ReLU(), 128 | nn.Linear(256, 100), 129 | ) 130 | else: 131 | self.encoder = nn.Sequential( 132 | PointNetfeat(self.num_points), 133 | nn.Linear(self.feat_size, 256), 134 | nn.ReLU(), 135 | nn.Linear(256, 100), 136 | ) 137 | self.lstm = nn.LSTMCell(100, 100) 138 | # self.gru = nn.GRUCell(100, 100) 139 | self.fc1 = nn.Linear(100, self.out_size) 140 | self.fc2 = nn.Linear(100, 1) 141 | # initialize bias to be high 142 | # self.fc2.bias.data.fill_(.99) 143 | 144 | def forward(self, x): 145 | x = self.encoder(x) 146 | hx, cx = Variable(torch.zeros(x.size(0), 100).cuda()), Variable( 147 | torch.zeros(x.size(0), 100).cuda() 148 | ) 149 | # hx = Variable(torch.zeros(x.size(0), 100).cuda()) 150 | outputs = [] 151 | probs = [] 152 | for i in range(self.num_primitives): 153 | hx, cx = self.lstm(x, (hx, cx)) 154 | # hx = self.gru(x, hx) 155 | outputs.append(self.fc1(hx)) 156 | probs.append(torch.sigmoid(self.fc2(hx))) 157 | return torch.cat((outputs), 1), torch.cat((probs), 1) 158 | 159 | 160 | class PrimitiveProbNet(nn.Module): 161 | def __init__( 162 | self, 163 | num_points=2500, 164 | input_dim=3, 165 | feat_size=1024, 166 | hidden_size=256, 167 | out_size=10, 168 | num_primitives=5, 169 | ): 170 | super(PrimitiveProbNet, self).__init__() 171 | self.num_points = num_points 172 | self.input_dim = input_dim 173 | self.feat_size = feat_size 174 | self.hidden_size = hidden_size 175 | self.out_size = out_size 176 | self.num_primitives = num_primitives 177 | self.feat = PointNetfeat(self.num_points, self.input_dim, self.feat_size) 178 | # params 179 | self.fc1 = nn.Linear(self.feat_size, self.hidden_size) 180 | self.fc2 = nn.Linear(self.hidden_size, self.out_size) 181 | # probs 182 | self.fc3 = nn.Linear(self.feat_size, self.hidden_size) 183 | self.fc4 = nn.Linear(self.hidden_size, self.num_primitives) 184 | 185 | def forward(self, x): 186 | x = self.feat(x) 187 | probs = torch.sigmoid(self.fc4(F.relu(self.fc3(x)))) 188 | out = self.fc2(F.relu(self.fc1(x))) 189 | return out, probs 190 | 191 | 192 | class PointEncoder(nn.Module): 193 | def __init__(self, num_points=256, feat_size=1024, out_size=10): 194 | super(PointEncoder, self).__init__() 195 | self.num_points = num_points 196 | self.feat = PointNetfeat(num_points=num_points, feat_size=feat_size) 197 | self.fc1 = nn.Linear(feat_size, 128) 198 | self.fc2 = nn.Linear(128, out_size) 199 | 200 | def forward(self, x): 201 | x = self.feat(x) 202 | x = F.relu((self.fc1(x))) 203 | x = self.fc2(x) 204 | return x 205 | 206 | 207 | class PointDecoder(nn.Module): 208 | def __init__(self, num_points=2048, k=2): 209 | super(PointDecoder, self).__init__() 210 | self.num_points = num_points 211 | self.fc1 = nn.Linear(100, 128) 212 | self.fc2 = nn.Linear(128, 256) 213 | self.fc3 = nn.Linear(256, 512) 214 | self.fc4 = nn.Linear(512, 1024) 215 | self.fc5 = nn.Linear(1024, self.num_points * 3) 216 | self.th = nn.Tanh() 217 | 218 | def forward(self, x): 219 | batchsize = x.size()[0] 220 | x = F.relu(self.fc1(x)) 221 | x = F.relu(self.fc2(x)) 222 | x = F.relu(self.fc3(x)) 223 | x = F.relu(self.fc4(x)) 224 | x = self.th(self.fc5(x)) 225 | x = x.view(batchsize, 3, self.num_points) 226 | return x 227 | 228 | 229 | class PointNetAE(nn.Module): 230 | def __init__(self, num_points=2048, k=2): 231 | super(PointNetAE, self).__init__() 232 | self.num_points = num_points 233 | self.encoder = nn.Sequential( 234 | PointNetfeat(num_points), 235 | nn.Linear(1024, 256), 236 | nn.ReLU(), 237 | nn.Linear(256, 100), 238 | ) 239 | self.decoder = PointDecoder(num_points) 240 | 241 | def forward(self, x): 242 | x = self.encoder(x) 243 | x = self.decoder(x) 244 | return x 245 | 246 | 247 | if __name__ == "__main__": 248 | sim_data = Variable(torch.rand(2, 3, 4096)) 249 | 250 | pointfeat = PointNetfeat() 251 | out = pointfeat(sim_data) 252 | print("point feat", out.size()) 253 | 254 | # test PrimitiveNet 255 | model = PrimitiveNet(out_size=3 * 7) 256 | out = model(sim_data) 257 | loss = out.mean() 258 | print("out.size() : ", out.size()) 259 | print(out) 260 | print(loss) 261 | loss.backward() 262 | print("backward pass done") 263 | 264 | print("-" * 50) 265 | # random data 266 | batch_size = 4 267 | num_points = 4096 268 | num_channels = 3 269 | sim_data = Variable(torch.rand(batch_size, num_channels, num_points).cuda()) 270 | print("MetricModel") 271 | metric = MetricModel(num_points=num_points, num_channels=num_channels) 272 | metric.cuda() 273 | output = metric(sim_data) 274 | print("output['feat'].shape : {}".format(output["feat"].shape)) 275 | loss = output["feat"].mean() 276 | loss.backward() 277 | 278 | print("-" * 50) 279 | print("RLNet") 280 | sim_data = Variable(torch.rand(32, 3, 4096).cuda(), requires_grad=True) 281 | model = RLNet(num_points=4096) 282 | model.cuda() 283 | outputs, probs = model(sim_data) 284 | loss = outputs.sum() 285 | loss.backward() 286 | print("sim_data.size() : ", sim_data.size()) 287 | print("outputs.size() : ", outputs.size()) 288 | print("probs.size() : ", probs.size()) 289 | print(loss) 290 | print("backward pass done") 291 | 292 | print("-" * 50) 293 | print("PrimitiveProbNet") 294 | sim_data = Variable(torch.rand(32, 3, 4096).cuda(), requires_grad=True) 295 | model = PrimitiveProbNet(num_points=4096) 296 | model.cuda() 297 | outputs, probs = model(sim_data) 298 | loss = outputs.sum() 299 | loss.backward() 300 | print("sim_data.size() : ", sim_data.size()) 301 | print("outputs.size() : ", outputs.size()) 302 | print("probs.size() : ", probs.size()) 303 | print(loss) 304 | print("backward pass done") 305 | -------------------------------------------------------------------------------- /code/inference/plot.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import math 4 | import matplotlib.pyplot as plt 5 | import matplotlib.cm as cm 6 | from matplotlib import animation 7 | from mpl_toolkits.mplot3d import axes3d, Axes3D 8 | from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection 9 | from obj_utils import * 10 | from plyfile import PlyData 11 | 12 | cmap = cm.get_cmap("tab10") 13 | colors = [cmap(i) for i in np.arange(0, 1, 1 / 10)] 14 | 15 | 16 | def euler2rot(theta): 17 | """Return a 3D rotation matrix given a vector of angles along (X,Y,Z)""" 18 | R_x = np.array( 19 | [ 20 | [1, 0, 0], 21 | [0, np.cos(theta[0]), -np.sin(theta[0])], 22 | [0, np.sin(theta[0]), np.cos(theta[0])], 23 | ] 24 | ) 25 | R_y = np.array( 26 | [ 27 | [np.cos(theta[1]), 0, np.sin(theta[1])], 28 | [0, 1, 0], 29 | [-np.sin(theta[1]), 0, np.cos(theta[1])], 30 | ] 31 | ) 32 | R_z = np.array( 33 | [ 34 | [np.cos(theta[2]), -np.sin(theta[2]), 0], 35 | [np.sin(theta[2]), np.cos(theta[2]), 0], 36 | [0, 0, 1], 37 | ] 38 | ) 39 | R = np.dot(R_z, np.dot(R_y, R_x)) 40 | return R 41 | 42 | 43 | def Rotate2D(pts, cnt, ang=np.pi / 4): 44 | """pts = {} Rotates points(nx2) about center cnt(2) by angle ang(1) in radian""" 45 | return ( 46 | np.dot( 47 | pts - cnt, 48 | np.array([[np.cos(ang), np.sin(ang)], [-np.sin(ang), np.cos(ang)]]), 49 | ) 50 | + cnt 51 | ) 52 | 53 | 54 | def get_processed_cubes(cubes, return_idx=False): 55 | """remove cubes which lie entirely inside of another cube""" 56 | idx = [] 57 | out = [] 58 | for c1, P1 in enumerate(cubes): 59 | flag = True 60 | for c2, P2 in enumerate(cubes): 61 | if c1 == c2: 62 | continue 63 | i = P1[0] - P1[4] 64 | j = P1[5] - P1[4] 65 | k = P1[7] - P1[4] 66 | v = P2 - P1[4] 67 | counts = ( 68 | (0 <= np.dot(v, i)) 69 | & (np.dot(v, i) <= np.dot(i, i)) 70 | & (0 <= np.dot(v, j)) 71 | & (np.dot(v, j) <= np.dot(j, j)) 72 | & (0 <= np.dot(v, k)) 73 | & (np.dot(v, k) <= np.dot(k, k)) 74 | ) 75 | if sum(counts) == 8: 76 | flag = False 77 | idx.append(c2) 78 | out = [cubes[i] for i in range(len(cubes)) if i not in idx] 79 | if len(out) < len(cubes): 80 | print( 81 | "Removed {} cubes for being within another cube".format( 82 | len(cubes) - len(out) 83 | ) 84 | ) 85 | if return_idx: 86 | return out, idx 87 | return out 88 | 89 | 90 | def draw_cube_points( 91 | cubes=None, pts=None, save=False, save_path="", save_gif=False, aoi_name="" 92 | ): 93 | 94 | def animate(i): 95 | ax.view_init(elev=30, azim=i) 96 | return (fig,) 97 | 98 | fig = plt.figure() 99 | ax = Axes3D(fig) 100 | minmax = {} 101 | minmax["xmin"] = 1e15 102 | minmax["ymin"] = 1e15 103 | minmax["zmin"] = 1e15 104 | minmax["xmax"] = -1e15 105 | minmax["ymax"] = -1e15 106 | minmax["zmax"] = -1e15 107 | 108 | cube_cords = [] 109 | if cubes is not None: 110 | if len(cubes.shape) == 1: 111 | cubes = cubes.reshape((1, -1)) 112 | for i, cube in enumerate(cubes): 113 | if len(cube) > 6: 114 | R = euler2rot([0, 0, cube[-1]]) 115 | else: 116 | R = euler2rot([0, 0, 0]) 117 | R = R.T 118 | S = np.diag(cube[3:6] * 2) 119 | 120 | # unit cube at center 121 | P1 = np.array( 122 | [ 123 | [-0.5, -0.5, -0.5], 124 | [0.5, -0.5, -0.5], 125 | [0.5, 0.5, -0.5], 126 | [-0.5, 0.5, -0.5], 127 | [-0.5, -0.5, 0.5], 128 | [0.5, -0.5, 0.5], 129 | [0.5, 0.5, 0.5], 130 | [-0.5, 0.5, 0.5], 131 | ] 132 | ) 133 | # apply transforms 134 | P2 = np.dot(P1, S) 135 | P = np.dot(P2, R) + cube[0:3] 136 | 137 | # add offset to heights based on points 138 | zmin = pts.min(0)[-1] 139 | diff = P.min(0)[-1] - zmin 140 | P[:, -1][:4] -= diff 141 | cube_cords.append(P) 142 | 143 | for i, P in enumerate(get_processed_cubes(cube_cords)): 144 | sides = [ 145 | [P[0], P[1], P[2], P[3]], 146 | [P[4], P[5], P[6], P[7]], 147 | [P[0], P[1], P[5], P[4]], 148 | [P[1], P[2], P[6], P[5]], 149 | [P[4], P[7], P[3], P[0]], 150 | [P[2], P[3], P[7], P[6]], 151 | ] 152 | collection = Poly3DCollection(sides, linewidths=1, alpha=1.0) 153 | collection.set_edgecolor("k") 154 | collection.set_facecolor(colors[i]) 155 | ax.add_collection3d(collection) 156 | minmax["xmin"] = min(minmax["xmin"], P.min(0)[0]) 157 | minmax["ymin"] = min(minmax["ymin"], P.min(0)[1]) 158 | minmax["zmin"] = min(minmax["zmin"], P.min(0)[2]) 159 | minmax["xmax"] = max(minmax["xmax"], P.max(0)[0]) 160 | minmax["ymax"] = max(minmax["ymax"], P.max(0)[1]) 161 | minmax["zmax"] = max(minmax["zmax"], P.max(0)[2]) 162 | 163 | if pts is not None: 164 | ax.scatter3D(pts[:, 0], pts[:, 1], pts[:, 2], s=1, c="r", depthshade=False) 165 | 166 | # Hide grid lines 167 | ax.grid(False) 168 | plt.axis("off") 169 | # Hide axes ticks 170 | # ax.set_xticks([]) 171 | # ax.set_yticks([]) 172 | # ax.set_zticks([]) 173 | # ax.set_aspect("equal") 174 | if save: 175 | if save_gif: 176 | # Animate 177 | anim = animation.FuncAnimation( 178 | fig, animate, frames=360, interval=20, blit=True 179 | ) 180 | # Save 181 | # anim.save('{}.mp4'.format(name), fps=30, extra_args=['-vcodec', 'libx264']) 182 | name_gif = len( 183 | [_ for _ in os.listdir(os.path.join(save_path, "gifs")) if "gif" in _] 184 | ) 185 | anim.save( 186 | os.path.join(save_path, "gifs/{}_{}.gif".format(aoi_name, name_gif)), 187 | writer="imagemagick", 188 | fps=30, 189 | ) 190 | print( 191 | "gif saved at : ", 192 | os.path.join(save_path, "gifs/{}_{}.gif".format(aoi_name, name_gif)), 193 | ) 194 | name_img = len( 195 | [_ for _ in os.listdir(os.path.join(save_path, "images")) if "jpg" in _] 196 | ) 197 | plt.savefig( 198 | os.path.join(save_path, "images/{}_{}.jpg".format(aoi_name, name_img)) 199 | ) 200 | print( 201 | "img saved at : ", 202 | os.path.join(save_path, "images/{}_{}.jpg".format(aoi_name, name_img)), 203 | ) 204 | plt.close() 205 | else: 206 | plt.show() 207 | 208 | 209 | def plot_obj(points, params, save_path, name_obj="", aoi_name=""): 210 | """save ply files output""" 211 | cube_cords = [] 212 | for i, cube in enumerate(params): 213 | if len(cube) > 6: 214 | R = euler2rot([0, 0, cube[-1]]) 215 | else: 216 | R = euler2rot([0, 0, 0]) 217 | R = R.T 218 | S = np.diag(cube[3:6] * 2) 219 | 220 | # unit cube at center 221 | P1 = np.array( 222 | [ 223 | [-0.5, -0.5, -0.5], 224 | [0.5, -0.5, -0.5], 225 | [0.5, 0.5, -0.5], 226 | [-0.5, 0.5, -0.5], 227 | [-0.5, -0.5, 0.5], 228 | [0.5, -0.5, 0.5], 229 | [0.5, 0.5, 0.5], 230 | [-0.5, 0.5, 0.5], 231 | ] 232 | ) 233 | # apply transforms 234 | P2 = np.dot(P1, S) 235 | P = np.dot(P2, R) + cube[0:3] 236 | 237 | # add offset to heights based on points 238 | # zmin = points.min(0)[-1] 239 | # diff = P.min(0)[-1] - zmin 240 | # P[:,-1][:4] -= diff 241 | cube_cords.append(P) 242 | 243 | _, idx = get_processed_cubes(cube_cords, return_idx=True) 244 | params = [params[i] for i in range(len(cube_cords)) if i not in idx] 245 | 246 | vs = [] 247 | fs = [] 248 | mtls = [] 249 | for i, cube in enumerate(params): 250 | vs1, fs1 = params_to_cubes(cube) 251 | mtls1 = i + np.zeros(len(fs1)).astype(np.int32) 252 | if len(fs) > 0: 253 | fs1 += np.max(fs) 254 | fs = np.vstack((fs, fs1)) 255 | vs = np.vstack((vs, vs1)) 256 | mtls = np.hstack((mtls, mtls1)) 257 | else: 258 | vs = np.array(vs1) 259 | fs = np.array(fs1) 260 | mtls = mtls1 261 | cols = [colors[i][:3] for i in range(len(params))] 262 | 263 | # plot only primitives 264 | if not name_obj: 265 | name_obj = len( 266 | [ 267 | _ 268 | for _ in os.listdir(os.path.join(save_path, "obj_primitives")) 269 | if "obj" in _ 270 | ] 271 | ) 272 | fout1 = open( 273 | os.path.join( 274 | save_path, "obj_primitives", "{}_{}.obj".format(aoi_name, name_obj) 275 | ), 276 | "w", 277 | ) 278 | fout2 = open( 279 | os.path.join( 280 | save_path, "obj_primitives", "{}_{}.mtl".format(aoi_name, name_obj) 281 | ), 282 | "w", 283 | ) 284 | 285 | append_mtl_obj(fout1, vs, fs, mtls, "{}_{}.mtl".format(aoi_name, name_obj)) 286 | append_mtl(fout2, list(set(mtls)), cols) 287 | fout1.close() 288 | fout2.close() 289 | print( 290 | "obj primitives saved at : ", 291 | os.path.join( 292 | save_path, "obj_primitives", "{}_{}.obj".format(aoi_name, name_obj) 293 | ), 294 | ) 295 | 296 | vs1, fs1 = points_to_cubes(points, 0.005) 297 | fs1 += np.max(fs) 298 | fs = np.vstack((fs, fs1)) 299 | vs = np.vstack((vs, vs1)) 300 | mtls1 = len(params) + np.zeros(len(fs1)).astype(np.int32) 301 | mtls = np.hstack((mtls, mtls1)) 302 | # cols.append(colors[len(cols)][:3]) 303 | cols.append([0, 0, 0]) 304 | 305 | # plot points and primitives overlayed 306 | if not name_obj: 307 | name_obj = len( 308 | [ 309 | _ 310 | for _ in os.listdir(os.path.join(save_path, "obj_overlayed")) 311 | if "obj" in _ 312 | ] 313 | ) 314 | fout1 = open( 315 | os.path.join( 316 | save_path, "obj_overlayed", "{}_{}.obj".format(aoi_name, name_obj) 317 | ), 318 | "w", 319 | ) 320 | fout2 = open( 321 | os.path.join( 322 | save_path, "obj_overlayed", "{}_{}.mtl".format(aoi_name, name_obj) 323 | ), 324 | "w", 325 | ) 326 | 327 | append_mtl_obj(fout1, vs, fs, mtls, "{}_{}.mtl".format(aoi_name, name_obj)) 328 | append_mtl(fout2, list(set(mtls)), cols) 329 | fout1.close() 330 | fout2.close() 331 | print( 332 | "obj (points+primitives) saved at : ", 333 | os.path.join( 334 | save_path, "obj_overlayed", "{}_{}.obj".format(aoi_name, name_obj) 335 | ), 336 | ) 337 | 338 | 339 | if __name__ == "__main__": 340 | data = PlyData.read("../../data/synthetic/34184_2.ply") 341 | data = np.asarray(data["vertex"].data.tolist(), dtype=np.float32) 342 | cubes = np.load("../../data/params_synthetic.npy").item() 343 | cube = cubes["34184_2"] 344 | 345 | # center = (data.max(0)+data.min(0))/2 346 | # data = data - center 347 | # scale = np.max(data.max(0)-data.min(0))/2 348 | # data = data/scale 349 | 350 | # cube[:,:3] -= center 351 | # cube[:,:6] /= scale 352 | # data[:,:2] = Rotate2D(data[:,:2], [0,0], -cube[0][6]) 353 | # cube[:,:2] = Rotate2D(cube[:,:2], [0,0], -cube[0][6]) 354 | 355 | # data = data + center 356 | # data = data*scale 357 | 358 | # cube[:,:3] += center 359 | # cube[:,:6] *= scale 360 | 361 | # draw_cube_points(cube, data, save=False, save_path='../../outputs/') 362 | # plot_obj(data, cube, save_path='../../outputs/') 363 | # import glob 364 | # import ipdb 365 | # for f in glob.iglob('../../data/synthetic/*5.ply'): 366 | # data = PlyData.read(f) 367 | # data = np.asarray(data['vertex'].data.tolist(), dtype=np.float32) 368 | # cubes = np.load('../../data/params_synthetic.npy').item() 369 | # cube = cubes[f[21:-4]] 370 | # draw_cube_points(cube, data, save=False, save_path='../../outputs/') 371 | 372 | data = PlyData.read("/usr0/home/tkhot/Downloads/slides.ply") 373 | data = np.asarray(data["vertex"].data.tolist(), dtype=np.float32) 374 | cube = np.load("/usr0/home/tkhot/Downloads/slides.npy") 375 | # cube = cubes['34184_2'] 376 | draw_cube_points(cube, data, save=False, save_path="../../outputs/") 377 | --------------------------------------------------------------------------------