├── LICENSE ├── README.md ├── assets ├── pipeline_neuralImplicit.png └── twobranch.png ├── configs ├── README.md ├── ifr-fusion-lr-kt0.yaml ├── ifr-fusion-lr-kt1.yaml ├── ifr-fusion-lr-kt2.yaml ├── ifr-fusion-lr-kt3.yaml ├── ifr-fusion-replica.yaml ├── test_kitti.yaml └── train.yaml ├── data_generator.py ├── dataset ├── kitti_dataset.py ├── production │ ├── __init__.py │ └── icl_nuim.py └── training │ ├── lif_dataset.py │ ├── main_generator.py │ ├── shapenet_model.py │ ├── shapenet_valid_list.json │ └── simple_shape.py ├── dataset_ptam.py ├── exp_transform.py ├── ifr ├── ifr_exp_transform.py └── ifr_main.py ├── im2mesh ├── __init__.py ├── checkpoints.py ├── common.py ├── config.py ├── data │ ├── __init__.py │ ├── core.py │ ├── fields.py │ ├── real.py │ └── transforms.py ├── dmc │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ │ ├── __init__.py │ │ ├── decoder.py │ │ └── encoder.py │ ├── ops │ │ ├── __init__.py │ │ ├── cpp_modules │ │ │ ├── old │ │ │ │ ├── commons.cpp │ │ │ │ ├── commons.h │ │ │ │ ├── pred_to_mesh.cpp │ │ │ │ └── pred_to_mesh.h │ │ │ ├── pred_to_mesh_.cpp │ │ │ └── setup.py │ │ ├── curvature_constraint.py │ │ ├── grid_pooling.py │ │ ├── occupancy_connectivity.py │ │ ├── occupancy_to_topology.py │ │ ├── point_triangle_distance.py │ │ ├── setup.py │ │ ├── src │ │ │ ├── curvature_constraint_kernel.cu │ │ │ ├── extension.cpp │ │ │ ├── grid_pooling_kernel.cu │ │ │ ├── kernels.h │ │ │ ├── occupancy_connectivity_kernel.cu │ │ │ ├── occupancy_to_topology_kernel.cu │ │ │ └── point_triangle_distance_kernel.cu │ │ ├── table.py │ │ └── tests │ │ │ ├── loss_autograd.py │ │ │ ├── test_curvature.py │ │ │ ├── test_distance.py │ │ │ ├── test_gridpooling.py │ │ │ ├── test_occupancy_connectivity.py │ │ │ ├── test_occupancy_connectivity_yiyi.py │ │ │ └── test_occupancy_to_topology.py │ ├── training.py │ └── utils │ │ ├── __init__.py │ │ ├── config.py │ │ ├── pointTriangleDistance.py │ │ ├── pred2mesh.py │ │ ├── util.py │ │ └── visualize.py ├── encoder │ ├── __init__.py │ ├── conv.py │ ├── pix2mesh_cond.py │ ├── pointnet.py │ ├── psgn_cond.py │ ├── r2n2.py │ ├── vnn.py │ ├── vnn2.py │ ├── vnn_tnet.py │ └── voxels.py ├── eval.py ├── layers.py ├── layers_equi.py ├── onet │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ │ ├── __init__.py │ │ ├── decoder.py │ │ ├── decoder_inner.py │ │ ├── encoder_latent.py │ │ └── legacy.py │ └── training.py ├── pix2mesh │ ├── __init__.py │ ├── config.py │ ├── ellipsoid │ │ ├── face1.obj │ │ ├── face2.obj │ │ ├── face3.obj │ │ └── info_ellipsoid.dat │ ├── generation.py │ ├── layers.py │ ├── models │ │ ├── __init__.py │ │ └── decoder.py │ └── training.py ├── preprocess.py ├── psgn │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ │ ├── __init__.py │ │ ├── decoder.py │ │ └── psgn_2branch.py │ └── training.py ├── r2n2 │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ │ ├── __init__.py │ │ └── decoder.py │ └── training.py ├── training.py ├── utils │ ├── __init__.py │ ├── binvox_rw.py │ ├── icp.py │ ├── io.py │ ├── libkdtree │ │ ├── .gitignore │ │ ├── LICENSE.txt │ │ ├── MANIFEST.in │ │ ├── README │ │ ├── README.rst │ │ ├── __init__.py │ │ ├── pykdtree │ │ │ ├── __init__.py │ │ │ ├── _kdtree_core.c │ │ │ ├── _kdtree_core.c.mako │ │ │ ├── kdtree.c │ │ │ ├── kdtree.pyx │ │ │ ├── render_template.py │ │ │ └── test_tree.py │ │ └── setup.cfg │ ├── libmcubes │ │ ├── .gitignore │ │ ├── LICENSE │ │ ├── README.rst │ │ ├── __init__.py │ │ ├── exporter.py │ │ ├── marchingcubes.cpp │ │ ├── marchingcubes.h │ │ ├── mcubes.pyx │ │ ├── pyarray_symbol.h │ │ ├── pyarraymodule.h │ │ ├── pywrapper.cpp │ │ └── pywrapper.h │ ├── libmesh │ │ ├── .gitignore │ │ ├── __init__.py │ │ ├── inside_mesh.py │ │ └── triangle_hash.pyx │ ├── libmise │ │ ├── .gitignore │ │ ├── __init__.py │ │ ├── mise.pyx │ │ └── test.py │ ├── libsimplify │ │ ├── Simplify.h │ │ ├── __init__.py │ │ ├── simplify_mesh.pyx │ │ └── test.py │ ├── libvoxelize │ │ ├── .gitignore │ │ ├── __init__.py │ │ ├── tribox2.h │ │ └── voxelize.pyx │ ├── mesh.py │ ├── visualize.py │ └── voxels.py └── vnn_onet │ ├── __init__.py │ ├── config.py │ ├── generation.py │ ├── models │ ├── __init__.py │ ├── decoder.py │ ├── decoder_inner.py │ ├── encoder_latent.py │ └── legacy.py │ └── training.py ├── network ├── criterion.py ├── di_decoder.py ├── di_encoder.py ├── di_vnn_encoder.py ├── utility.py └── vnn.py ├── network_trainer.py ├── pose_fmt.py ├── pose_fmt_kitti.py ├── sampler_cuda ├── CMakeLists.txt ├── PreprocessMesh.cu ├── ShaderProgram.cpp ├── Utils.cu └── Utils.h ├── system ├── ext │ ├── __init__.py │ ├── imgproc │ │ ├── common.cuh │ │ ├── imgproc.cpp │ │ ├── imgproc.cu │ │ └── photometric.cu │ ├── indexing │ │ ├── indexing.cpp │ │ └── indexing.cu │ ├── marching_cubes │ │ ├── mc.cpp │ │ ├── mc_data.cuh │ │ └── mc_interp_kernel.cu │ └── pcproc │ │ ├── cuda_kdtree.cu │ │ ├── cuda_kdtree.cuh │ │ ├── cutil_math.h │ │ ├── pcproc.cpp │ │ └── pcproc.cu ├── map.py ├── map_dictindexer.py ├── map_outer_J.py └── tracker.py ├── trainer ├── lr_schedule.py └── main.py ├── transform ├── transform.py └── utils.py ├── treasure ├── encoder_600.pth.tar ├── hyper.json └── model_600.pth.tar └── utils ├── coding.py ├── draw.py ├── exp_util.py ├── gradient.py ├── index.py ├── motion_util.py ├── pt_util.py └── vis_util.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Yijun Yuan 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 | -------------------------------------------------------------------------------- /assets/pipeline_neuralImplicit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jarrome/IMT_Mapping/974547684a4474be4c281a9380d4a15150f4bf48/assets/pipeline_neuralImplicit.png -------------------------------------------------------------------------------- /assets/twobranch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jarrome/IMT_Mapping/974547684a4474be4c281a9380d4a15150f4bf48/assets/twobranch.png -------------------------------------------------------------------------------- /configs/README.md: -------------------------------------------------------------------------------- 1 | ## How to configure 2 | 3 | For example, ```./config/ifr-fusion-lr-kt0.yaml``` contains several settings: 4 | 5 | ``` 6 | dataset_type: "TUM" # for ICL-NUIM 's TUM format 7 | scene: "0n" # scene 0 with noise 8 | use_gt: False # not use groundtruth trajectory 9 | pose_folder: "./treasure/orbslam2_record/lrkt0n/" # the predicted pose stream from orbslam2 10 | outdir: "./res/lrkt0n_ours/" # where we stores the intermediary mesh 11 | calib: [481.2, 480.0, 319.50, 239.50, 5000.0] # calibration of images 12 | 13 | sequence_kwargs: # this is called with within sequence namespace 14 | path: "./data/ICL_NUIM/lr_kt0n/" 15 | start_frame: 0 16 | end_frame: -1 # Run all frames 17 | first_tq: [-1.4, 1.5, 1.5, 0.0, -1.0, 0.0, 0.0] # Starting pose 18 | 19 | # Network parameters 20 | training_hypers: "./treasure/hyper.json" 21 | using_epoch: 600 22 | 23 | # Enable visualization 24 | vis: True 25 | resolution: 4 26 | 27 | # meshing 28 | max_n_triangles: 4e6 29 | max_std: 0.15 # 0.06 30 | 31 | # These two define the range of depth observations to be cropped. Unit is meter. 32 | depth_cut_min: 0.5 33 | depth_cut_max: 5.0 34 | 35 | # not exactly used in code, please follow real implementation 36 | meshing_interval: 20 37 | integrate_interval: 20 38 | 39 | # Mapping parameters 40 | mapping: 41 | # Bound of the scene to be reconstructed 42 | bound_min: [-3.5, -0.5, -2.5] 43 | bound_max: [4.5, 3.5, 5.5] 44 | voxel_size: 0.1 45 | # Prune observations if detected as noise. 46 | prune_min_vox_obs: 16 47 | ignore_count_th: 16.0 48 | encoder_count_th: 600.0 49 | ``` 50 | -------------------------------------------------------------------------------- /configs/ifr-fusion-lr-kt0.yaml: -------------------------------------------------------------------------------- 1 | # Sequence parameters 2 | sequence_type: "icl_nuim.ICLNUIMSequence" 3 | 4 | dataset_type: "TUM" 5 | scene: "0n" 6 | use_gt: False 7 | pose_folder: "./treasure/orbslam2_record/lrkt0n/" 8 | outdir: "./res/lrkt0n_ours/" 9 | 10 | calib: [481.2, 480.0, 319.50, 239.50, 5000.0] 11 | 12 | sequence_kwargs: 13 | path: "./data/ICL_NUIM/lr_kt0n/" 14 | start_frame: 0 15 | end_frame: -1 # Run all frames 16 | first_tq: [-1.4, 1.5, 1.5, 0.0, -1.0, 0.0, 0.0] # Starting pose 17 | 18 | # Network parameters (network structure, etc. will be inherited from the training config) 19 | training_hypers: "./treasure/hyper.json" 20 | using_epoch: 600 21 | 22 | # Separate tracking and meshing. 23 | run_async: false 24 | # Enable visualization 25 | vis: True 26 | resolution: 4 27 | 28 | # meshing 29 | max_n_triangles: 4e6 30 | max_std: 0.15 # 0.06 31 | 32 | 33 | # These two define the range of depth observations to be cropped. Unit is meter. 34 | depth_cut_min: 0.5 35 | depth_cut_max: 5.0 36 | 37 | meshing_interval: 20 38 | integrate_interval: 20 39 | 40 | # Mapping parameters 41 | mapping: 42 | # Bound of the scene to be reconstructed 43 | bound_min: [-3.5, -0.5, -2.5] 44 | bound_max: [4.5, 3.5, 5.5] 45 | voxel_size: 0.1 46 | # Prune observations if detected as noise. 47 | prune_min_vox_obs: 16 48 | ignore_count_th: 16.0 49 | encoder_count_th: 600.0 50 | 51 | # Tracking parameters 52 | tracking: 53 | # An array defining how the camera pose is optimized. 54 | # Each element is a dictionary: 55 | # For example {"n": 2, "type": [['sdf'], ['rgb', 1]]} means to optimize the summation of sdf term and rgb term 56 | # at the 1st level pyramid for 2 iterations. 57 | iter_config: 58 | - {"n": 10, "type": [['rgb', 2]]} 59 | - {"n": 10, "type": [['sdf'], ['rgb', 1]]} 60 | - {"n": 50, "type": [['sdf'], ['rgb', 0]]} 61 | sdf: 62 | robust_kernel: "huber" 63 | robust_k: 5.0 64 | subsample: 0.5 65 | rgb: 66 | weight: 500.0 67 | robust_kernel: null 68 | robust_k: 0.01 69 | min_grad_scale: 0.0 70 | max_depth_delta: 0.2 71 | -------------------------------------------------------------------------------- /configs/ifr-fusion-lr-kt1.yaml: -------------------------------------------------------------------------------- 1 | # Sequence parameters 2 | sequence_type: "icl_nuim.ICLNUIMSequence" 3 | 4 | dataset_type: "TUM" 5 | scene: "1n" 6 | use_gt: False 7 | pose_folder: "./treasure/orbslam2_record/lrkt1n/" 8 | outdir: "./res/lrkt1n_ours/" 9 | 10 | calib: [481.2, 480.0, 319.50, 239.50, 5000.0] 11 | 12 | sequence_kwargs: 13 | path: "./data/ICL_NUIM/lr_kt1n/" 14 | start_frame: 0 15 | end_frame: -1 16 | first_tq: [-0.1, 1.3, 0.1, 0.0, -1.0, 0.0, 0.0] 17 | 18 | # Network parameters (network structure, etc. will be inherited from the training config) 19 | training_hypers: "./treasure/hyper.json" 20 | using_epoch: 600 21 | 22 | # Separate tracking and meshing. 23 | run_async: false 24 | # Enable visualization 25 | vis: True 26 | resolution: 4 27 | 28 | 29 | # meshing 30 | max_n_triangles: 4e6 31 | max_std: 0.15 # 0.06 32 | 33 | 34 | depth_cut_min: 0.5 35 | depth_cut_max: 3.0 36 | 37 | meshing_interval: 20 38 | integrate_interval: 20 39 | 40 | # Mapping parameters (See README.md for details) 41 | mapping: 42 | bound_min: [-3.5, -0.5, -2.5] 43 | bound_max: [4.5, 3.5, 5.5] 44 | voxel_size: 0.1 45 | prune_min_vox_obs: 16 46 | ignore_count_th: 16.0 47 | encoder_count_th: 600.0 48 | 49 | # Tracking parameters (See README.md for details) 50 | tracking: 51 | iter_config: 52 | - {"n": 10, "type": [['rgb', 2]]} 53 | - {"n": 10, "type": [['sdf'], ['rgb', 1]]} 54 | - {"n": 50, "type": [['sdf'], ['rgb', 0]]} 55 | sdf: 56 | robust_kernel: "huber" 57 | robust_k: 5.0 58 | subsample: 0.5 59 | rgb: 60 | weight: 500.0 61 | robust_kernel: null 62 | robust_k: 0.01 63 | min_grad_scale: 0.0 64 | max_depth_delta: 0.2 65 | 66 | -------------------------------------------------------------------------------- /configs/ifr-fusion-lr-kt2.yaml: -------------------------------------------------------------------------------- 1 | # Sequence parameters 2 | sequence_type: "icl_nuim.ICLNUIMSequence" 3 | 4 | dataset_type: "TUM" 5 | scene: "2n" 6 | use_gt: False 7 | pose_folder: "./treasure/orbslam2_record/lrkt2n/" 8 | outdir: "./res/lrkt2n_ours/" 9 | 10 | calib: [481.2, 480.0, 319.50, 239.50, 5000.0] 11 | 12 | 13 | sequence_kwargs: 14 | path: "./data/ICL_NUIM/lr_kt2n/" 15 | start_frame: 0 16 | end_frame: -1 # Run all frames 17 | first_tq: [-0.7, 1.3, 1.5, 0.0, -1.0, 0.0, 0.0] # Starting pose 18 | 19 | # Network parameters (network structure, etc. will be inherited from the training config) 20 | training_hypers: "./treasure/hyper.json" 21 | using_epoch: 600 22 | 23 | # Separate tracking and meshing. 24 | run_async: false 25 | # Enable visualization 26 | vis: True 27 | resolution: 4 28 | 29 | # meshing 30 | max_n_triangles: 4e6 31 | max_std: 0.15 # 0.06 32 | 33 | 34 | # These two define the range of depth observations to be cropped. Unit is meter. 35 | depth_cut_min: 0.5 36 | depth_cut_max: 3.0 37 | 38 | meshing_interval: 20 39 | integrate_interval: 20 40 | 41 | # Mapping parameters 42 | mapping: 43 | # Bound of the scene to be reconstructed 44 | bound_min: [-3.5, -0.5, -2.5] 45 | bound_max: [4.5, 3.5, 5.5] 46 | voxel_size: 0.1 47 | # Prune observations if detected as noise. 48 | prune_min_vox_obs: 16 49 | ignore_count_th: 16.0 50 | encoder_count_th: 600.0 51 | 52 | # Tracking parameters 53 | tracking: 54 | # An array defining how the camera pose is optimized. 55 | # Each element is a dictionary: 56 | # For example {"n": 2, "type": [['sdf'], ['rgb', 1]]} means to optimize the summation of sdf term and rgb term 57 | # at the 1st level pyramid for 2 iterations. 58 | iter_config: 59 | - {"n": 10, "type": [['sdf'], ['rgb', 1]]} 60 | - {"n": 50, "type": [['sdf'], ['rgb', 0]]} 61 | sdf: 62 | robust_kernel: "huber" 63 | robust_k: 5.0 64 | subsample: 0.5 65 | rgb: 66 | weight: 50.0 67 | robust_kernel: null 68 | robust_k: 0.01 69 | min_grad_scale: 0.0 70 | max_depth_delta: 0.2 71 | -------------------------------------------------------------------------------- /configs/ifr-fusion-lr-kt3.yaml: -------------------------------------------------------------------------------- 1 | # Sequence parameters 2 | sequence_type: "icl_nuim.ICLNUIMSequence" 3 | 4 | dataset_type: "TUM" 5 | scene: "3n" 6 | use_gt: False 7 | pose_folder: "./treasure/orbslam2_record/lrkt3n/" 8 | outdir: "./res/lrkt3n_ours/" 9 | 10 | calib: [481.2, 480.0, 319.50, 239.50, 5000.0] 11 | 12 | sequence_kwargs: 13 | path: "./data/ICL_NUIM/lr_kt3n/" 14 | start_frame: 0 15 | end_frame: -1 # Run all frames 16 | first_tq: [-1.2, 1.3, 1.0, 0.0, -1.0, 0.0, 0.0] # Starting pose 17 | 18 | # Network parameters (network structure, etc. will be inherited from the training config) 19 | training_hypers: "./treasure/hyper.json" 20 | using_epoch: 600 21 | 22 | # Separate tracking and meshing. 23 | run_async: false 24 | # Enable visualization 25 | vis: True 26 | resolution: 4 27 | 28 | # meshing 29 | max_n_triangles: 4e6 30 | max_std: 0.15 # 0.06 31 | 32 | 33 | # These two define the range of depth observations to be cropped. Unit is meter. 34 | depth_cut_min: 0.5 35 | depth_cut_max: 5.0 36 | 37 | meshing_interval: 20 38 | integrate_interval: 20 39 | 40 | # Mapping parameters 41 | mapping: 42 | # Bound of the scene to be reconstructed 43 | bound_min: [-3.5, -0.5, -2.5] 44 | bound_max: [4.5, 3.5, 5.5] 45 | voxel_size: 0.1 46 | # Prune observations if detected as noise. 47 | prune_min_vox_obs: 16 48 | ignore_count_th: 16.0 49 | encoder_count_th: 600.0 50 | 51 | # Tracking parameters 52 | tracking: 53 | # An array defining how the camera pose is optimized. 54 | # Each element is a dictionary: 55 | # For example {"n": 2, "type": [['sdf'], ['rgb', 1]]} means to optimize the summation of sdf term and rgb term 56 | # at the 1st level pyramid for 2 iterations. 57 | iter_config: 58 | - {"n": 10, "type": [['sdf'], ['rgb', 1]]} 59 | - {"n": 50, "type": [['sdf'], ['rgb', 0]]} 60 | sdf: 61 | robust_kernel: "huber" 62 | robust_k: 5.0 63 | subsample: 0.5 64 | rgb: 65 | weight: 50.0 66 | robust_kernel: null 67 | robust_k: 0.01 68 | min_grad_scale: 0.0 69 | max_depth_delta: 0.2 70 | -------------------------------------------------------------------------------- /configs/ifr-fusion-replica.yaml: -------------------------------------------------------------------------------- 1 | # Sequence parameters 2 | sequence_type: "icl_nuim.ICLNUIMSequence" 3 | 4 | dataset_type: "replica" 5 | scene: "office0" 6 | use_gt: False 7 | pose_folder: "./treasure/orbslam2_record/imap_office0/" 8 | outdir: "./res/replica_office0_ours/" 9 | 10 | calib: [600., 600., 599.5, 339.5, 6553.5] 11 | 12 | sequence_kwargs: 13 | path: "./data/replica/office0" 14 | start_frame: 0 15 | end_frame: -1 # Run all frames 16 | #first_tq: [-1.2, 1.3, 1.0, 0.0, -1.0, 0.0, 0.0] # Starting pose 17 | first_tq: [0, 0, 0.0, 0.0, -1.0, 0.0, 0.0] 18 | 19 | # Network parameters (network structure, etc. will be inherited from the training config) 20 | training_hypers: "./treasure/hyper.json" 21 | using_epoch: 600 22 | 23 | # Separate tracking and meshing. 24 | run_async: false 25 | # Enable visualization 26 | vis: false 27 | resolution: 3 28 | 29 | # These two define the range of depth observations to be cropped. Unit is meter. 30 | depth_cut_min: 0.5 31 | depth_cut_max: 5.0 32 | 33 | meshing_interval: 20 34 | integrate_interval: 20 35 | 36 | # Mapping parameters 37 | mapping: 38 | # Bound of the scene to be reconstructed 39 | bound_min: [-10.5, -5.5, -10.] 40 | bound_max: [10.5, 5.5, 10.5] 41 | 42 | voxel_size: 0.1 43 | # Prune observations if detected as noise. 44 | prune_min_vox_obs: 16 45 | ignore_count_th: 16.0 46 | encoder_count_th: 600.0 47 | 48 | # Tracking parameters 49 | tracking: 50 | # An array defining how the camera pose is optimized. 51 | # Each element is a dictionary: 52 | # For example {"n": 2, "type": [['sdf'], ['rgb', 1]]} means to optimize the summation of sdf term and rgb term 53 | # at the 1st level pyramid for 2 iterations. 54 | iter_config: 55 | - {"n": 10, "type": [['sdf'], ['rgb', 1]]} 56 | - {"n": 50, "type": [['sdf'], ['rgb', 0]]} 57 | sdf: 58 | robust_kernel: "huber" 59 | robust_k: 5.0 60 | subsample: 0.5 61 | rgb: 62 | weight: 50.0 63 | robust_kernel: null 64 | robust_k: 0.01 65 | min_grad_scale: 0.0 66 | max_depth_delta: 0.2 67 | -------------------------------------------------------------------------------- /configs/test_kitti.yaml: -------------------------------------------------------------------------------- 1 | # data 2 | basedir: "./data/kitti_odometry/dataset/sequences" 3 | sequence: 0 4 | 5 | # Network parameters (network structure, etc. will be inherited from the training config) 6 | training_hypers: "./treasure/hyper.json" 7 | using_epoch: 600 8 | 9 | # Separate tracking and meshing. 10 | run_async: false 11 | # Enable visualization 12 | vis: false 13 | resolution: 4 14 | 15 | # These two define the range of depth observations to be cropped. Unit is meter. 16 | depth_cut_min: 0.5 17 | depth_cut_max: 5.0 18 | 19 | meshing_interval: 20 20 | integrate_interval: 5 21 | 22 | # Mapping parameters 23 | mapping: 24 | # Bound of the scene to be reconstructed 25 | bound_min: [-200., -800., -100.] 26 | bound_max: [800., 800., 200.] 27 | voxel_size: 4. 28 | # Prune observations if detected as noise. 29 | prune_min_vox_obs: 0 30 | ignore_count_th: 100 31 | encoder_count_th: 600.0 32 | 33 | vis: False 34 | 35 | 36 | # Tracking parameters 37 | tracking: 38 | # An array defining how the camera pose is optimized. 39 | # Each element is a dictionary: 40 | # For example {"n": 2, "type": [['sdf'], ['rgb', 1]]} means to optimize the summation of sdf term and rgb term 41 | # at the 1st level pyramid for 2 iterations. 42 | iter_config: 43 | - {"n": 10, "type": [['sdf']]} 44 | - {"n": 50, "type": [['sdf']]} 45 | sdf: 46 | robust_kernel: "huber" 47 | robust_k: 5.0 48 | subsample: 0.5 49 | -------------------------------------------------------------------------------- /configs/train.yaml: -------------------------------------------------------------------------------- 1 | run_name: "default" 2 | 3 | num_epochs: 600 4 | batch_size: 64 5 | batch_split: 1 6 | samples_per_lif: 4096 7 | min_context_points: 16 8 | 9 | lr_schedule: 10 | # For decoder parameters 11 | - { "Type" : "Step", "Initial" : 0.001, "Interval" : 80, "Factor" : 0.4 } 12 | # For encoder parameters 13 | - { "Type" : "Step", "Initial" : 0.001, "Interval" : 80, "Factor" : 0.4 } 14 | 15 | # Dataset. 16 | train_set: 17 | - { "data_path": "../di-datasets/shapenet_plivoxs", "augment_rotation": 'Y', "num_surface_sample": 128, "augment_noise": [0.025, 40.0] } 18 | 19 | # Code specification 20 | code_bound: null 21 | code_length: 9 22 | 23 | # Decoder specification 24 | network_name: "di_decoder" 25 | network_specs: 26 | dims: [ 128, 128, 128, 128 ] 27 | dropout: [0, 1, 2, 3, 4, 5] 28 | dropout_prob: 0.2 29 | norm_layers: [0, 1, 2, 3, 4, 5] 30 | latent_in: [3] 31 | weight_norm: true 32 | 33 | # Encoder specification 34 | encoder_name: "di_vnn_encoder" 35 | encoder_specs: 36 | per_point_feat: [ 6, 32, 64, 256 ] 37 | bn: {"class": "BatchNorm"} 38 | 39 | # Snapshots saving parameters 40 | snapshot_frequency: 100 41 | additional_snapshots: [50] 42 | 43 | # SDF samples 44 | training_loss: 45 | types: [ "neg_log_likelihood", "reg_loss" ] 46 | enforce_minmax: true 47 | clamping_distance: 0.2 48 | code_reg_lambda: 1.0e-2 49 | -------------------------------------------------------------------------------- /data_generator.py: -------------------------------------------------------------------------------- 1 | dataset/training/main_generator.py -------------------------------------------------------------------------------- /dataset/production/__init__.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | class FrameIntrinsic: 5 | def __init__(self, fx, fy, cx, cy, dscale): 6 | self.cx = cx 7 | self.cy = cy 8 | self.fx = fx 9 | self.fy = fy 10 | self.dscale = dscale 11 | 12 | def to_K(self): 13 | return np.asarray([ 14 | [self.fx, 0.0, self.cx], 15 | [0.0, self.fy, self.cy], 16 | [0.0, 0.0, 1.0] 17 | ]) 18 | 19 | 20 | class FrameData: 21 | def __init__(self): 22 | self.rgb = None 23 | self.depth = None 24 | self.gt_pose = None 25 | self.calib = None 26 | 27 | 28 | class RGBDSequence: 29 | def __init__(self): 30 | self.frame_id = 0 31 | 32 | def __iter__(self): 33 | return self 34 | 35 | def __len__(self): 36 | raise NotImplementedError 37 | 38 | def __next__(self): 39 | raise NotImplementedError 40 | -------------------------------------------------------------------------------- /dataset/production/icl_nuim.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import torch 4 | from dataset.production import * 5 | from pyquaternion import Quaternion 6 | from pathlib import Path 7 | from utils import motion_util 8 | import pdb 9 | 10 | 11 | class ICLNUIMSequence(RGBDSequence): 12 | def __init__(self, path: str, start_frame: int = 0, end_frame: int = -1, first_tq: list = None, load_gt: bool = False): 13 | super().__init__() 14 | self.path = Path(path) 15 | self.color_names = sorted([f"rgb/{t}" for t in os.listdir(self.path / "rgb")], key=lambda t: int(t[4:].split(".")[0])) 16 | self.depth_names = [f"depth/{t}.png" for t in range(len(self.color_names))] 17 | self.calib = [481.2, 480.0, 319.50, 239.50, 5000.0] 18 | #self.calib = [525.0, 525.0, 319.5, 239.5, 5000] 19 | if first_tq is not None: 20 | self.first_iso = motion_util.Isometry(q=Quaternion(array=first_tq[3:]), t=np.array(first_tq[:3])) 21 | else: 22 | self.first_iso = motion_util.Isometry(q=Quaternion(array=[0.0, -1.0, 0.0, 0.0])) 23 | 24 | if end_frame == -1: 25 | end_frame = len(self.color_names) 26 | 27 | self.color_names = self.color_names[start_frame:end_frame] 28 | self.depth_names = self.depth_names[start_frame:end_frame] 29 | 30 | if load_gt: 31 | gt_traj_path = (list(self.path.glob("*.freiburg")) + list(self.path.glob("groundtruth.txt")))[0] 32 | self.gt_trajectory = self._parse_traj_file(gt_traj_path) 33 | self.gt_trajectory = self.gt_trajectory[start_frame:end_frame] 34 | change_iso = self.first_iso.dot(self.gt_trajectory[0].inv()) 35 | self.gt_trajectory = [change_iso.dot(t) for t in self.gt_trajectory] 36 | assert len(self.gt_trajectory) == len(self.color_names) 37 | else: 38 | self.gt_trajectory = None 39 | 40 | def _parse_traj_file(self, traj_path): 41 | camera_ext = {} 42 | traj_data = np.genfromtxt(traj_path) 43 | cano_quat = motion_util.Isometry(q=Quaternion(axis=[0.0, 0.0, 1.0], degrees=180.0)) 44 | for cur_p in traj_data: 45 | cur_q = Quaternion(imaginary=cur_p[4:7], real=cur_p[-1]).rotation_matrix 46 | cur_t = cur_p[1:4] 47 | cur_q[1] = -cur_q[1] 48 | cur_q[:, 1] = -cur_q[:, 1] 49 | cur_t[1] = -cur_t[1] 50 | cur_iso = motion_util.Isometry(q=Quaternion(matrix=cur_q), t=cur_t) 51 | camera_ext[cur_p[0]] = cano_quat.dot(cur_iso) 52 | camera_ext[0] = camera_ext[1] 53 | return [camera_ext[t] for t in range(len(camera_ext))] 54 | 55 | def __len__(self): 56 | return len(self.color_names) 57 | 58 | def __next__(self): 59 | if self.frame_id >= len(self): 60 | raise StopIteration 61 | 62 | depth_img_path = self.path / self.depth_names[self.frame_id] 63 | rgb_img_path = self.path / self.color_names[self.frame_id] 64 | 65 | # Convert depth image into point cloud. 66 | frame_data = FrameData() 67 | depth_data = cv2.imread(str(depth_img_path), cv2.IMREAD_UNCHANGED) 68 | frame_data.depth_cv2 = depth_data 69 | depth_data = torch.from_numpy(depth_data.astype(np.float32)).cuda() / self.calib[4] 70 | 71 | rgb_data = cv2.imread(str(rgb_img_path)) 72 | rgb_data = cv2.cvtColor(rgb_data, cv2.COLOR_BGR2RGB) 73 | frame_data.rgb_cv2 = rgb_data 74 | rgb_data = torch.from_numpy(rgb_data).cuda().float() / 255. 75 | 76 | frame_data.gt_pose = self.gt_trajectory[self.frame_id] if self.gt_trajectory is not None else None 77 | frame_data.calib = FrameIntrinsic(self.calib[0], self.calib[1], self.calib[2], self.calib[3], self.calib[4]) 78 | frame_data.depth = depth_data 79 | frame_data.rgb = rgb_data 80 | 81 | self.frame_id += 1 82 | return frame_data 83 | -------------------------------------------------------------------------------- /dataset/training/shapenet_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import json 4 | import random 5 | import logging 6 | import numpy as np 7 | from utils import motion_util 8 | from pathlib import Path 9 | 10 | 11 | class ShapeNetGenerator: 12 | """ 13 | Use ShapeNet core to generate data. 14 | """ 15 | VALID_LIST_PATH = Path(__file__).parent / "shapenet_valid_list.json" 16 | 17 | def __init__(self, shapenet_path, categories, shapes_per_category, scale): 18 | self.categories = categories 19 | self.shapes_per_category = shapes_per_category 20 | self.scale = scale 21 | 22 | # Sample objects 23 | self.data_sources = [] 24 | self.data_scales = [] 25 | with self.VALID_LIST_PATH.open("r") as f: 26 | valid_list_data = json.load(f) 27 | 28 | for category_name, category_shape_count, category_scale in zip(self.categories, self.shapes_per_category, self.scale): 29 | category_path = Path(shapenet_path) / category_name 30 | if category_name in valid_list_data["ShapeNetV2"].keys(): 31 | logging.info(f"Category {category_name} is found in plist file") 32 | sampled_objects = valid_list_data["ShapeNetV2"][category_name] 33 | else: 34 | logging.info(f"Category {category_name} is not found in plist file") 35 | sampled_objects = os.listdir(category_path) 36 | if category_shape_count != -1: 37 | sampled_objects = random.sample(sampled_objects, category_shape_count) 38 | self.data_sources += [category_path / s for s in sampled_objects] 39 | self.data_scales += [category_scale for _ in sampled_objects] 40 | 41 | def __len__(self): 42 | return len(self.data_sources) 43 | 44 | @staticmethod 45 | def _equidist_point_on_sphere(samples): 46 | points = [] 47 | phi = math.pi * (3. - math.sqrt(5.)) 48 | 49 | for i in range(samples): 50 | y = 1 - (i / float(samples - 1)) * 2 51 | radius = math.sqrt(1 - y * y) 52 | theta = phi * i 53 | 54 | x = math.cos(theta) * radius 55 | z = math.sin(theta) * radius 56 | points.append((x, y, z)) 57 | 58 | return np.asarray(points) 59 | 60 | def get_source(self, data_id): 61 | return str(self.data_sources[data_id]) 62 | 63 | def __getitem__(self, idx): 64 | data_source = self.data_sources[idx] 65 | data_scale = self.data_scales[idx] 66 | obj_path = data_source / "models" / "model_normalized.obj" 67 | 68 | vp_camera = self._equidist_point_on_sphere(300) 69 | camera_ext = [] 70 | for camera_i in range(vp_camera.shape[0]): 71 | iso = motion_util.Isometry.look_at(vp_camera[camera_i], np.zeros(3,)) 72 | camera_ext.append(iso) 73 | camera_int = [0.8, 0.0, 2.5] # (window-size-half, z-min, z-max) under ortho-proj. 74 | 75 | return str(obj_path), [camera_int, camera_ext], None, data_scale 76 | 77 | def clean(self, data_id): 78 | pass 79 | -------------------------------------------------------------------------------- /exp_transform.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | import time 5 | from itertools import chain 6 | from collections import defaultdict 7 | from scipy.spatial.transform import Rotation as R 8 | 9 | import ifr.ifr_exp_transform as ifr 10 | from ifr.ifr_exp_transform import vis_param 11 | 12 | import glob 13 | import pdb 14 | 15 | 16 | 17 | from pyquaternion import Quaternion 18 | 19 | # namespace 20 | args = vis_param.args 21 | sequence = vis_param.sequence 22 | # scene name 23 | scene = args.scene 24 | # where we store the incremental result for demonstration 25 | args.outdir_transform = args.outdir[:-1]+'_transform' 26 | os.makedirs(args.outdir_transform,exist_ok=True) 27 | 28 | 29 | 30 | if __name__ == '__main__': 31 | import os 32 | import sys 33 | 34 | from dataset_ptam import TUMRGBDDataset, ICLNUIMDataset 35 | 36 | 37 | args.dataset = args.dataset_type 38 | 39 | if 'tum' in args.dataset.lower(): 40 | dataset = TUMRGBDDataset(sequence.path) 41 | else: 42 | assert "Not supported data type" 43 | 44 | ''' 45 | load gt traj to check correctness 46 | ''' 47 | GT = True 48 | if GT: 49 | gt_traj = np.genfromtxt(str(sequence.path)+'/livingRoom'+scene+'.gt.freiburg') 50 | gt_poses = [] 51 | 52 | 53 | 54 | durations = [] 55 | data_i = 0 56 | #for i in range(len(dataset))[:]: 57 | kf_idx = [] 58 | def run_algo(vis): 59 | global data_i 60 | i = data_i#data_next() 61 | data_i += 1 62 | 63 | 64 | if i % 20 == 0:# 65 | is_keyframe = True 66 | kf_idx.append(i) 67 | else: 68 | is_keyframe = False 69 | if dataset.timestamps is None: 70 | timestamp = i / 20. 71 | else: 72 | timestamp = dataset.timestamps[i] 73 | 74 | time_start = time.time() 75 | 76 | # 0. check if current keyframe 77 | if is_keyframe: 78 | #gt_pose = g2o.Isometry3d(g2o.Quaternion(gt_traj[i,-1],gt_traj[i,4],gt_traj[i,5],gt_traj[i,6]), gt_traj[i,1:4]) 79 | gt_pose = gt_traj[i,:] if GT else None 80 | # 1. prepare current frame to get torch frame_data 81 | frame_data = (dataset.rgb[i],dataset.depth[i]) 82 | 83 | # 2. get all the poses of keyframe 84 | new_poses = [] 85 | if not GT: 86 | assert(False) 87 | poses = read_elasticfusion_file(i, kf_idx) 88 | new_poses= poses 89 | else: 90 | gt_poses.append(gt_pose) 91 | new_poses = gt_poses 92 | 93 | # 3.2 if some pose changed, update map 94 | ifr.refresh(frame_data, new_poses, frame_id = i, vis=vis, ptam_p = not GT) 95 | else: 96 | return 97 | 98 | 99 | duration = time.time() - time_start 100 | durations.append(duration) 101 | print('duration', duration) 102 | print() 103 | print() 104 | 105 | if ifr.engine: 106 | ifr.engine.register_animation_callback(callback_func = run_algo) 107 | vis_ph = ifr.vis_util.wireframe_bbox([-4., -4., -4.], [4., 4., 4.]) 108 | ifr.engine.add_geometry(vis_ph) 109 | ifr.engine.remove_geometry(vis_ph, reset_bounding_box=False) 110 | ifr.engine.run() 111 | ifr.engine.destroy_window() 112 | else: 113 | try: 114 | while True: 115 | run_algo(None) 116 | except Exception as e: 117 | print(e) 118 | 119 | 120 | -------------------------------------------------------------------------------- /im2mesh/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jarrome/IMT_Mapping/974547684a4474be4c281a9380d4a15150f4bf48/im2mesh/__init__.py -------------------------------------------------------------------------------- /im2mesh/checkpoints.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib 3 | import torch 4 | from torch.utils import model_zoo 5 | 6 | 7 | class CheckpointIO(object): 8 | ''' CheckpointIO class. 9 | 10 | It handles saving and loading checkpoints. 11 | 12 | Args: 13 | checkpoint_dir (str): path where checkpoints are saved 14 | ''' 15 | def __init__(self, checkpoint_dir='./chkpts', **kwargs): 16 | self.module_dict = kwargs 17 | self.checkpoint_dir = checkpoint_dir 18 | if not os.path.exists(checkpoint_dir): 19 | os.makedirs(checkpoint_dir) 20 | 21 | def register_modules(self, **kwargs): 22 | ''' Registers modules in current module dictionary. 23 | ''' 24 | self.module_dict.update(kwargs) 25 | 26 | def save(self, filename, **kwargs): 27 | ''' Saves the current module dictionary. 28 | 29 | Args: 30 | filename (str): name of output file 31 | ''' 32 | if not os.path.isabs(filename): 33 | filename = os.path.join(self.checkpoint_dir, filename) 34 | 35 | outdict = kwargs 36 | for k, v in self.module_dict.items(): 37 | outdict[k] = v.state_dict() 38 | torch.save(outdict, filename) 39 | 40 | def load(self, filename): 41 | '''Loads a module dictionary from local file or url. 42 | 43 | Args: 44 | filename (str): name of saved module dictionary 45 | ''' 46 | if is_url(filename): 47 | return self.load_url(filename) 48 | else: 49 | return self.load_file(filename) 50 | 51 | def load_file(self, filename): 52 | '''Loads a module dictionary from file. 53 | 54 | Args: 55 | filename (str): name of saved module dictionary 56 | ''' 57 | 58 | if not os.path.isabs(filename): 59 | filename = os.path.join(self.checkpoint_dir, filename) 60 | 61 | if os.path.exists(filename): 62 | print(filename) 63 | print('=> Loading checkpoint from local file...') 64 | state_dict = torch.load(filename) 65 | scalars = self.parse_state_dict(state_dict) 66 | return scalars 67 | else: 68 | raise FileExistsError 69 | 70 | def load_url(self, url): 71 | '''Load a module dictionary from url. 72 | 73 | Args: 74 | url (str): url to saved model 75 | ''' 76 | print(url) 77 | print('=> Loading checkpoint from url...') 78 | state_dict = model_zoo.load_url(url, progress=True) 79 | scalars = self.parse_state_dict(state_dict) 80 | return scalars 81 | 82 | def parse_state_dict(self, state_dict): 83 | '''Parse state_dict of model and return scalars. 84 | 85 | Args: 86 | state_dict (dict): State dict of model 87 | ''' 88 | 89 | for k, v in self.module_dict.items(): 90 | if k in state_dict: 91 | v.load_state_dict(state_dict[k]) 92 | else: 93 | print('Warning: Could not find %s in checkpoint!' % k) 94 | scalars = {k: v for k, v in state_dict.items() 95 | if k not in self.module_dict} 96 | return scalars 97 | 98 | def is_url(url): 99 | scheme = urllib.parse.urlparse(url).scheme 100 | return scheme in ('http', 'https') -------------------------------------------------------------------------------- /im2mesh/data/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from im2mesh.data.core import ( 3 | Shapes3dDataset, collate_remove_none, worker_init_fn 4 | ) 5 | from im2mesh.data.fields import ( 6 | IndexField, CategoryField, ImagesField, PointsField, 7 | VoxelsField, PointCloudField, MeshField, 8 | ) 9 | from im2mesh.data.transforms import ( 10 | PointcloudNoise, SubsamplePointcloud, 11 | SubsamplePoints 12 | ) 13 | from im2mesh.data.real import ( 14 | KittiDataset, OnlineProductDataset, 15 | ImageDataset, 16 | ) 17 | 18 | 19 | __all__ = [ 20 | # Core 21 | Shapes3dDataset, 22 | collate_remove_none, 23 | worker_init_fn, 24 | # Fields 25 | IndexField, 26 | CategoryField, 27 | ImagesField, 28 | PointsField, 29 | VoxelsField, 30 | PointCloudField, 31 | MeshField, 32 | # Transforms 33 | PointcloudNoise, 34 | SubsamplePointcloud, 35 | SubsamplePoints, 36 | # Real Data 37 | KittiDataset, 38 | OnlineProductDataset, 39 | ImageDataset, 40 | ] 41 | -------------------------------------------------------------------------------- /im2mesh/dmc/__init__.py: -------------------------------------------------------------------------------- 1 | from im2mesh.dmc import ( 2 | config, generation, training, models 3 | ) 4 | 5 | __all__ = [ 6 | config, generation, training, models 7 | ] 8 | -------------------------------------------------------------------------------- /im2mesh/dmc/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | from im2mesh.dmc import models, training, generation 3 | from im2mesh import data 4 | 5 | 6 | def get_model(cfg, device=None, **kwargs): 7 | encoder = cfg['model']['encoder'] 8 | decoder = cfg['model']['decoder'] 9 | c_dim = cfg['model']['c_dim'] 10 | encoder_kwargs = cfg['model']['encoder_kwargs'] 11 | decoder_kwargs = cfg['model']['decoder_kwargs'] 12 | 13 | encoder = models.encoder_dict[encoder]( 14 | **encoder_kwargs 15 | ) 16 | 17 | decoder = models.decoder_dict[decoder]( 18 | **decoder_kwargs 19 | ) 20 | 21 | model = models.DMC(decoder, encoder) 22 | model = model.to(device) 23 | return model 24 | 25 | 26 | def get_trainer(model, optimizer, cfg, device, **kwargs): 27 | input_type = cfg['data']['input_type'] 28 | out_dir = cfg['training']['out_dir'] 29 | vis_dir = os.path.join(out_dir, 'vis') 30 | num_voxels = cfg['model']['num_voxels'] 31 | weight_prior = cfg['model']['dmc_weight_prior'] 32 | 33 | trainer = training.Trainer( 34 | model, optimizer, device=device, input_type=input_type, 35 | vis_dir=vis_dir, num_voxels=num_voxels, 36 | weight_prior=weight_prior, 37 | ) 38 | return trainer 39 | 40 | 41 | def get_generator(model, cfg, device, **kwargs): 42 | num_voxels = cfg['model']['num_voxels'] 43 | 44 | generator = generation.Generator3D( 45 | model, device=device, num_voxels=num_voxels 46 | ) 47 | return generator 48 | 49 | 50 | def get_data_fields(split, cfg, **kwargs): 51 | with_transforms = cfg['data']['with_transforms'] 52 | # TODO: put this into config 53 | pointcloud_n = 3000 54 | pointcloud_transform = data.SubsamplePointcloud(pointcloud_n) 55 | 56 | fields = {} 57 | fields['pointcloud'] = data.PointCloudField( 58 | cfg['data']['pointcloud_file'], pointcloud_transform, 59 | with_transforms=with_transforms 60 | ) 61 | 62 | return fields 63 | -------------------------------------------------------------------------------- /im2mesh/dmc/generation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import trimesh 4 | from im2mesh.dmc.utils.pred2mesh import pred_to_mesh_max 5 | from im2mesh.dmc.ops.occupancy_to_topology import OccupancyToTopology 6 | from im2mesh.dmc.ops.table import get_accept_topology 7 | 8 | 9 | class Generator3D(object): 10 | def __init__(self, model, device=None, num_voxels=32): 11 | self.model = model.to(device) 12 | self.device = device 13 | self.num_voxels = num_voxels 14 | self.vis_topology = torch.LongTensor(get_accept_topology(4)) 15 | 16 | def generate_mesh(self, data): 17 | self.model.eval() 18 | device = self.device 19 | 20 | inputs = data.get('inputs', torch.empty(1, 0)).to(device) 21 | 22 | inputs = self.num_voxels * (inputs / 1.2 + 0.5) 23 | 24 | with torch.no_grad(): 25 | offset, topology, occupancy = self.model(inputs) 26 | 27 | offset = offset.squeeze() 28 | topology = topology.squeeze() 29 | topology = topology[:, self.vis_topology] 30 | 31 | vertices, faces = pred_to_mesh_max(offset, topology) 32 | faces = faces.astype(np.int64) 33 | 34 | vertices = 1.2 * (vertices / self.num_voxels - 0.5) 35 | mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) 36 | return mesh 37 | 38 | 39 | -------------------------------------------------------------------------------- /im2mesh/dmc/models/__init__.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from im2mesh.dmc.models import encoder, decoder 3 | 4 | 5 | decoder_dict = { 6 | 'unet': decoder.UNetDecoder 7 | } 8 | 9 | encoder_dict = { 10 | 'pointnet_local': encoder.PointNetLocal, 11 | } 12 | 13 | class DMC(nn.Module): 14 | def __init__(self, decoder, encoder): 15 | super().__init__() 16 | self.decoder = decoder 17 | self.encoder = encoder 18 | 19 | def forward(self, x): 20 | c = self.encoder(x) 21 | offset, topology, occupancy = self.decoder(c) 22 | 23 | return offset, topology, occupancy 24 | -------------------------------------------------------------------------------- /im2mesh/dmc/models/encoder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | 4 | from im2mesh.dmc.ops.grid_pooling import GridPooling 5 | 6 | 7 | class PointNetLocal(nn.Module): 8 | ''' Point Net Local Conditional Network from the Deep Marching Cubes paper. 9 | 10 | It applies two fully connected layers to the input points (dim 3) in a 11 | 1D Convolutional Layer fashion to avoid to specify the number of 12 | incoming points 13 | ''' 14 | def __init__(self, c_dim=256, out_dim=16, cell_W=16, cell_H=16, cell_D=16): 15 | super().__init__() 16 | self.cell_W = cell_W 17 | self.cell_H = cell_H 18 | self.cell_D = cell_D 19 | 20 | # TODO change gridpooling input to be compatible to single values of W H D 21 | self.gridshape = torch.cuda.LongTensor([cell_W, cell_H, cell_D]) 22 | actvn = nn.ReLU() 23 | self.grid_pool = GridPooling(self.gridshape) 24 | self.conv1 = nn.Sequential( 25 | nn.Conv1d(3, c_dim, 1), actvn 26 | ) 27 | #self.conv2 = nn.Sequential( 28 | # nn.Conv1d(c_dim, out_dim, 1), actvn 29 | #) 30 | self.conv2 = nn.Conv1d(c_dim, out_dim, 1) 31 | 32 | def forward(self, x): 33 | pts = x 34 | feats = x.transpose(1, 2) # b_size x 3 x num_points 35 | feats = self.conv1(feats) # b_size x c_dim x num_points 36 | feats = self.conv2(feats) # b_size x out_dim x num_points 37 | feats = feats.transpose(1, 2) # b_size x num_points x out_dim 38 | 39 | out = self.point_to_cell(pts, feats, self.cell_W, self.cell_H, self.cell_D) 40 | return out 41 | 42 | def point_to_cell(self, pts, feat, W, H, D, expand=1): 43 | """ perform maxpool on points in every cell set zero vector if cell is 44 | empty if expand=1 then return (N+1)x(N+1)x(N+1), for dmc xpand=0 then 45 | return NxNxN, for occupancy/sdf baselines 46 | """ 47 | batchsize = feat.size()[0] 48 | C = feat.size()[2] 49 | 50 | feat_cell = [] 51 | # grid_shape = torch.LongTensor([W, H, D]) 52 | for k in range(batchsize): 53 | feat_cell.append(self.grid_pool(feat[k, :, :], pts[k, :, :])) 54 | 55 | feat_cell = torch.stack(feat_cell, dim=0) 56 | 57 | # TODO check if this view is compatible to output of grid pool 58 | feat_cell = torch.transpose(feat_cell, 1, 2).contiguous().view( 59 | -1, C, W, H, D) 60 | if expand == 0: 61 | return feat_cell 62 | 63 | # expand to (W+1)x(H+1) 64 | curr_size = feat_cell.size() 65 | feat_cell_exp = torch.zeros( 66 | curr_size[0], curr_size[1], curr_size[2]+1, curr_size[3]+1, 67 | curr_size[4]+1).to(pts.device) 68 | feat_cell_exp[:, :, :-1, :-1, :-1] = feat_cell 69 | return feat_cell_exp 70 | -------------------------------------------------------------------------------- /im2mesh/dmc/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jarrome/IMT_Mapping/974547684a4474be4c281a9380d4a15150f4bf48/im2mesh/dmc/ops/__init__.py -------------------------------------------------------------------------------- /im2mesh/dmc/ops/cpp_modules/old/commons.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "commons.h" 3 | 4 | /** 5 | * convert vertex displacement field to vertices locations 6 | * params: 7 | * offset vertex displacement field, 3xWxHxD 8 | * x indice of a cell in the full grid 9 | * y indice of a cell in the full grid 10 | * z indice of a cell in the full grid 11 | * return: 12 | * vertices the location of 12 vertices for the specific cell, 3x12 13 | * 14 | */ 15 | THFloatTensor* offset_to_vertices(THFloatTensor *offset, int x, int y, int z){ 16 | THFloatTensor *vertices = THFloatTensor_newWithSize2d(3, 12); 17 | 18 | // #0 19 | THFloatTensor_set2d(vertices, 0, 0, 0.5-THFloatTensor_get4d(offset, 0, x+1, y+1, z )); 20 | THFloatTensor_set2d(vertices, 1, 0, 1.0); 21 | THFloatTensor_set2d(vertices, 2, 0, 0.0); 22 | // #1 23 | THFloatTensor_set2d(vertices, 0, 1, 1.0); 24 | THFloatTensor_set2d(vertices, 1, 1, 0.5-THFloatTensor_get4d(offset, 1, x+1, y+1, z )); 25 | THFloatTensor_set2d(vertices, 2, 1, 0.0); 26 | // #2 27 | THFloatTensor_set2d(vertices, 0, 2, 0.5-THFloatTensor_get4d(offset, 0, x+1, y , z )); 28 | THFloatTensor_set2d(vertices, 1, 2, 0.0); 29 | THFloatTensor_set2d(vertices, 2, 2, 0.0); 30 | // #3 31 | THFloatTensor_set2d(vertices, 0, 3, 0.0); 32 | THFloatTensor_set2d(vertices, 1, 3, 0.5-THFloatTensor_get4d(offset, 1, x , y+1, z )); 33 | THFloatTensor_set2d(vertices, 2, 3, 0.0); 34 | 35 | // #4 36 | THFloatTensor_set2d(vertices, 0, 4, 0.5-THFloatTensor_get4d(offset, 0, x+1, y+1, z+1)); 37 | THFloatTensor_set2d(vertices, 1, 4, 1.0); 38 | THFloatTensor_set2d(vertices, 2, 4, 1.0); 39 | // #5 40 | THFloatTensor_set2d(vertices, 0, 5, 1.0); 41 | THFloatTensor_set2d(vertices, 1, 5, 0.5-THFloatTensor_get4d(offset, 1, x+1, y+1, z+1)); 42 | THFloatTensor_set2d(vertices, 2, 5, 1.0); 43 | // #6 44 | THFloatTensor_set2d(vertices, 0, 6, 0.5-THFloatTensor_get4d(offset, 0, x+1, y , z+1)); 45 | THFloatTensor_set2d(vertices, 1, 6, 0.0); 46 | THFloatTensor_set2d(vertices, 2, 6, 1.0); 47 | // #7 48 | THFloatTensor_set2d(vertices, 0, 7, 0.0); 49 | THFloatTensor_set2d(vertices, 1, 7, 0.5-THFloatTensor_get4d(offset, 1, x , y+1, z+1)); 50 | THFloatTensor_set2d(vertices, 2, 7, 1.0); 51 | 52 | // #8 53 | THFloatTensor_set2d(vertices, 0, 8, 0.0); 54 | THFloatTensor_set2d(vertices, 1, 8, 1.0); 55 | THFloatTensor_set2d(vertices, 2, 8, 0.5-THFloatTensor_get4d(offset, 2, x , y+1, z+1)); 56 | // #9 57 | THFloatTensor_set2d(vertices, 0, 9, 1.0); 58 | THFloatTensor_set2d(vertices, 1, 9, 1.0); 59 | THFloatTensor_set2d(vertices, 2, 9, 0.5-THFloatTensor_get4d(offset, 2, x+1, y+1, z+1)); 60 | // #10 61 | THFloatTensor_set2d(vertices, 0, 10, 1.0); 62 | THFloatTensor_set2d(vertices, 1, 10, 0.0); 63 | THFloatTensor_set2d(vertices, 2, 10, 0.5-THFloatTensor_get4d(offset, 2, x+1, y , z+1)); 64 | // #11 65 | THFloatTensor_set2d(vertices, 0, 11, 0.0); 66 | THFloatTensor_set2d(vertices, 1, 11, 0.0); 67 | THFloatTensor_set2d(vertices, 2, 11, 0.5-THFloatTensor_get4d(offset, 2, x , y , z+1)); 68 | return vertices; 69 | } 70 | 71 | /** 72 | * get points in a specific cell 73 | * params: 74 | * points all points in the grid, Nx3 75 | * i the offset of the specific cell 76 | * j the offset of the specific cell 77 | * k the offset of the specific cell 78 | * return: 79 | * indices a binary 1D tensor indicating if a point is in a specific cell or not, N 80 | * 81 | */ 82 | THLongTensor* points_in_grid(THFloatTensor *points, float i, float j, float k){ 83 | int N=THFloatTensor_size(points, 0); 84 | THLongTensor *indices = THLongTensor_new(); 85 | 86 | THByteTensor *mask = THByteTensor_newWithSize1d(N); 87 | THByteTensor_zero(mask); 88 | for (int p=0; p= i && THFloatTensor_get2d(points, p, 0) < i+1 && 90 | THFloatTensor_get2d(points, p, 1) >= j && THFloatTensor_get2d(points, p, 1) < j+1 && 91 | THFloatTensor_get2d(points, p, 2) >= k && THFloatTensor_get2d(points, p, 2) < k+1) 92 | THByteTensor_set1d(mask, p, 1); 93 | } 94 | 95 | THByteTensor_nonzero(indices, mask); 96 | 97 | THLongTensor_squeeze(indices, indices); 98 | THByteTensor_free(mask); 99 | return indices; 100 | } 101 | -------------------------------------------------------------------------------- /im2mesh/dmc/ops/cpp_modules/old/pred_to_mesh.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "commons.h" 5 | 6 | 7 | // considered all topologies with 4 triangles during visualization 8 | static int visTopology[2][140]={{0, 1, 2, 3, 4, 6, 7, 8, 9, 11, 12, 13, 14, 15, 16, 17, 19, 23, 25, 27, 29, 31, 32, 34, 35, 38, 39, 43, 46, 47, 48, 49, 50, 51, 54, 55, 57, 59, 63, 64, 68, 70, 71, 76, 77, 78, 79, 95, 96, 98, 99, 100, 102, 103, 108, 110, 111, 112, 113, 114, 115, 116, 118, 119, 123, 126, 127, 128, 136, 137, 139, 140, 141, 142, 143, 144, 145, 147, 152, 153, 155, 156, 157, 159, 175, 176, 177, 178, 179, 183, 184, 185, 187, 189, 191, 192, 196, 198, 200, 201, 204, 205, 206, 207, 208, 209, 212, 216, 217, 219, 220, 221, 222, 223, 224, 226, 228, 230, 231, 232, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255}, 9 | {0, 1, 1, 2, 1, 2, 3, 1, 2, 3, 2, 3, 3, 2, 1, 2, 3, 4, 3, 4, 4, 3, 1, 2, 3, 3, 4, 4, 4, 3, 2, 3, 3, 2, 4, 3, 4, 3, 2, 1, 2, 3, 4, 3, 4, 4, 3, 4, 2, 3, 4, 3, 2, 3, 4, 3, 2, 3, 4, 4, 3, 4, 3, 2, 4, 4, 1, 1, 2, 3, 4, 3, 4, 4, 3, 2, 3, 4, 3, 2, 3, 4, 3, 2, 4, 3, 4, 4, 3, 4, 4, 3, 2, 4, 1, 2, 3, 4, 3, 4, 2, 3, 3, 2, 3, 4, 4, 4, 3, 4, 3, 2, 4, 1, 3, 4, 4, 3, 4, 4, 3, 4, 2, 1, 2, 3, 3, 2, 3, 4, 2, 1, 3, 2, 4, 1, 2, 1, 1, 0}}; 10 | 11 | /** 12 | * convert the topology probability and vertex displacement field to a mesh by 13 | * selecting the topology with maximum probability in every cell 14 | * params: 15 | * offset vertex displacement field 16 | * topology topology probabilities 17 | * vertices_all vertices locations for all triangles in topologies with maximum probabilities 18 | * note there might be duplications and the unique vertices will be extracted afterwards 19 | * faces_all faces represented by the indices in vertices_all 20 | * vertice_number record the number of vertices as we initialzed the vertices_all with a fixed length 21 | * face_number record the number of faces as we initialized the faces_all with a fixed length 22 | */ 23 | int pred_to_mesh(THFloatTensor offset, THLongTensor *topology, THFloatTensor *vertices_all, THFloatTensor *faces_all, THLongTensor *vertice_number, THLongTensor *face_number){ 24 | // data format check 25 | if (THFloatTensor_nDimension(offset)!=4 || THLongTensor_nDimension(topology)!=3 ){ 26 | printf("Invalid nDimension!\n"); 27 | printf("Expected 4, 3, received %d, %d \n", THFloatTensor_nDimension(offset), THLongTensor_nDimension(topology)); 28 | return 0; 29 | } 30 | int W,H,D; 31 | W = THFloatTensor_size(offset,1)-1; 32 | H = THFloatTensor_size(offset,2)-1; 33 | D = THFloatTensor_size(offset,3)-1; 34 | 35 | int vertice_cnt=0; 36 | int face_cnt=0; 37 | 38 | for (int i=0; i 2 | #include 3 | #include 4 | #include 5 | #include 6 | __constant__ float grid_size=1.0; 7 | namespace{ 8 | /** 9 | * perform max-pooling within the cells 10 | * parallel over each cell and each feature dimension 11 | */ 12 | template 13 | __global__ void grid_pooling_kernel( 14 | const scalar_t *point, 15 | const scalar_t *feat_points, 16 | scalar_t *feat_cell, 17 | int *indices, 18 | const int n ){ 19 | // cell indices 20 | int i = blockIdx.x; 21 | int j = blockIdx.y; 22 | int k = blockIdx.z; 23 | // cell size 24 | // int W = gridDim.x; 25 | int H = gridDim.y; 26 | int D = gridDim.z; 27 | int ind = i*H*D + j*D + k; 28 | int c = threadIdx.x; 29 | int C = blockDim.x; 30 | 31 | for (int p=0; p= i && px < i+grid_size && py >= j && py < j+grid_size && pz >= k && pz < k+grid_size){ 37 | // max-pooling, update feat_cell if the feature is larger than the current feat_cell 38 | // can be async for max operation 39 | if ( feat_points[p*C + c] > feat_cell[ind*C + c] ){ 40 | feat_cell[ind*C + c] = feat_points[p*C + c]; 41 | indices[ind*C + c] = p; 42 | } 43 | } 44 | } 45 | } 46 | /** 47 | * back-propagate the loss from the max-pooled feature to point features 48 | * parallel over each cell and each feature dimension 49 | */ 50 | template 51 | __global__ void grad_grid_pooling_kernel( 52 | const scalar_t *grad_output, 53 | const int *indices, 54 | scalar_t *grad_feat_points){ 55 | // cell indices 56 | int i = blockIdx.x; 57 | int j = blockIdx.y; 58 | int k = blockIdx.z; 59 | // cell size 60 | // int W = gridDim.x; 61 | int H = gridDim.y; 62 | int D = gridDim.z; 63 | int ind = i*H*D + j*D + k; 64 | int c = threadIdx.x; 65 | int C = blockDim.x; 66 | long int p = indices[ind*C + c]; 67 | if (p < 0) return; 68 | 69 | grad_feat_points[p*C + c] = grad_output[ind*C + c]; 70 | } 71 | } //namespace 72 | 73 | /* 74 | * Forward function, project the point features to cells, perform max pooling in every cell 75 | * params: 76 | * point input, all points, Nx3 77 | * feat_points input, feature of all points, NxC 78 | * shape input, size of the grid [W, H, D], 3 79 | * feat_cell output, feature of all cells, (WxHxD)xC 80 | * indices output, indices of max pooling, saved for back propagation, (WxHxD)xC 81 | * 82 | */ 83 | 84 | void grid_pooling_kernel_forward( 85 | at::Tensor point, 86 | at::Tensor feat_points, 87 | at::Tensor shape, 88 | at::Tensor feat_cell, 89 | at::Tensor indices){ 90 | int W = shape.data()[0]; 91 | int H = shape.data()[1]; 92 | int D = shape.data()[2]; 93 | int C = feat_cell.size(1); 94 | 95 | dim3 dimGrid(W, H, D); 96 | dim3 dimBlock(C, 1, 1); 97 | // lauch the kernel 98 | int n = point.size(0); 99 | grid_pooling_kernel<<< dimGrid, dimBlock>>>( 100 | point.data(), 101 | feat_points.data(), 102 | feat_cell.data(), 103 | indices.data(), 104 | n); 105 | } 106 | 107 | /* 108 | * Backward function, back-propagate the loss to the point features 109 | * params: 110 | * grad_output input, gradient on the output feature, WxHxC 111 | * shape input, size of the grid [W, H, D], 3 112 | * indices input, indices of max pooling, WxHxC 113 | * grad_feat_points output, gradient on the features, NxC 114 | * 115 | */ 116 | void grid_pooling_kernel_backward( 117 | at::Tensor grad_output, 118 | at::Tensor shape, 119 | at::Tensor indices, 120 | at::Tensor grad_feat_points){ 121 | int W = shape.data()[0]; 122 | int H = shape.data()[1]; 123 | int D = shape.data()[2]; 124 | int C = grad_output.size(1); 125 | dim3 dimGrid(W, H, D); 126 | dim3 dimBlock(C, 1, 1); 127 | // lauch the kernel 128 | grad_grid_pooling_kernel<<< dimGrid, dimBlock>>>( 129 | grad_output.data(), 130 | indices.data(), 131 | grad_feat_points.data()); 132 | } 133 | 134 | 135 | -------------------------------------------------------------------------------- /im2mesh/dmc/ops/src/kernels.h: -------------------------------------------------------------------------------- 1 | #ifndef __DMC_CUDA_KERNELS__ 2 | #define __DMC_CUDA_KERNELS__ 3 | 4 | #include 5 | #include 6 | 7 | 8 | // Curvature constraint 9 | void curvature_constraint_kernel_forward( 10 | at::Tensor offset, 11 | at::Tensor topology, 12 | at::Tensor xTable, 13 | at::Tensor yTable, 14 | at::Tensor zTable, 15 | at::Tensor innerTable, 16 | at::Tensor loss_x, 17 | at::Tensor loss_y, 18 | at::Tensor loss_z, 19 | at::Tensor loss_inner); 20 | 21 | 22 | void curvature_constraint_kernel_backward( 23 | at::Tensor grad_output, 24 | at::Tensor offset, 25 | at::Tensor topology, 26 | at::Tensor xTable, 27 | at::Tensor yTable, 28 | at::Tensor zTable, 29 | at::Tensor innerTable, 30 | at::Tensor grad_offset); 31 | 32 | // Grid pooling 33 | void grid_pooling_kernel_forward( 34 | at::Tensor point, 35 | at::Tensor feat_points, 36 | at::Tensor shape, 37 | at::Tensor feat_cell, 38 | at::Tensor indices); 39 | 40 | void grid_pooling_kernel_backward( 41 | at::Tensor grad_output, 42 | at::Tensor shape, 43 | at::Tensor indices, 44 | at::Tensor grad_feat_points); 45 | 46 | // Occ2Topo 47 | void occupancy_to_topology_kernel_forward( 48 | at::Tensor occupancy, 49 | at::Tensor topology ); 50 | 51 | void occupancy_to_topology_kernel_backward( 52 | at::Tensor grad_output, 53 | at::Tensor occupancy, 54 | at::Tensor topology, 55 | at::Tensor grad_occupancy); 56 | 57 | // OccConstraint 58 | void occupancy_connectivity_kernel_forward( 59 | at::Tensor occupancy, 60 | at::Tensor loss); 61 | 62 | void occupancy_connectivity_kernel_backward( 63 | at::Tensor grad_output, 64 | at::Tensor occupancy, 65 | at::Tensor grad_occupancy); 66 | 67 | // Points Triangle distance 68 | void point_topology_distance_kernel_forward( 69 | at::Tensor offset, 70 | at::Tensor points, 71 | at::Tensor distances, 72 | at::Tensor indices_all); 73 | 74 | void point_topology_distance_kernel_backward( 75 | at::Tensor grad_output, 76 | at::Tensor offset, 77 | at::Tensor points, 78 | at::Tensor indices_all, 79 | at::Tensor grad_offset); 80 | 81 | #endif -------------------------------------------------------------------------------- /im2mesh/dmc/ops/tests/test_curvature.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import sys 5 | 6 | sys.path.append('../../../..') 7 | from im2mesh.dmc.ops.tests.loss_autograd import LossAutoGrad 8 | from im2mesh.dmc.ops.curvature_constraint import CurvatureConstraint 9 | import torch.nn.functional as F 10 | import numpy as np 11 | import time 12 | 13 | # check the cuda extension or c extension 14 | 15 | print ("Testing CUDA extension...") 16 | dtype = torch.cuda.FloatTensor 17 | 18 | 19 | # autograd loss 20 | num_cells = 4 21 | len_cell = 1.0 22 | W = H = D = num_cells 23 | 24 | loss_autograd = LossAutoGrad(num_cells, len_cell) 25 | 26 | 27 | # cffi loss 28 | class SmoothLoss(nn.Module): 29 | def __init__(self): 30 | super(SmoothLoss, self).__init__() 31 | self.smoothLoss = CurvatureConstraint() 32 | 33 | def forward(self, offset, topology): 34 | return self.smoothLoss(offset, topology) 35 | 36 | 37 | if __name__ == '__main__': 38 | 39 | # generate offset and topology with relatively low-dimension 40 | print ("=========== Input =============") 41 | T = 96 42 | W = num_cells 43 | H = num_cells 44 | D = num_cells 45 | offset = Variable((torch.rand(3, W+1, H+1, D+1)).type(dtype) * 0.1, requires_grad=True) 46 | topology = Variable(torch.rand(W*H*D, T).type(dtype), requires_grad=True) 47 | #print (offset) 48 | #print (topology) 49 | 50 | loss_cffi = SmoothLoss() 51 | l = loss_cffi(offset, F.softmax(topology, dim=1)) 52 | l.backward() 53 | offset.grad.data.zero_() 54 | 55 | # evaluating the running time of the cffi extension 56 | print ("============= cffi ============") 57 | tf_c = time.time() 58 | l = loss_cffi(offset, F.softmax(topology, dim=1)) 59 | print ("cffi loss:") 60 | print (l) 61 | tf_c = time.time()-tf_c 62 | 63 | 64 | tb_c = time.time() 65 | l.backward() 66 | print ("cffi gradient:") 67 | print( offset.grad) 68 | tb_c = time.time()-tb_c 69 | grad_np = np.copy(offset.grad.data.cpu().numpy()) 70 | 71 | 72 | # evaluating the running time of the autograd version 73 | print ("============= auto ============") 74 | tf_py = time.time() 75 | l_auto = loss_autograd.loss_on_curvature_autograd(offset, topology) 76 | print ("auto loss:") 77 | print (l_auto) 78 | tf_py = time.time()-tf_py 79 | 80 | offset.grad.data.zero_() 81 | tb_py = time.time() 82 | l_auto.backward() 83 | print ("auto grad:") 84 | print (offset.grad) 85 | tb_py = time.time()-tb_py 86 | grad_auto_np = np.copy(offset.grad.data.cpu().numpy()) 87 | assert np.sum(np.abs(grad_auto_np)) and np.sum(np.abs(grad_np)) != 0.0 88 | # print the loss and grad difference and the time comparison 89 | print ("========== summary ===========") 90 | print ("Forward difference between cffi and auto: ", (l-l_auto).data.cpu().numpy()) 91 | print ("Backward difference between cffi and auto: ", np.sum(np.abs(grad_np-grad_auto_np))) 92 | print ("Backward difference between cffi and auto: ", np.mean(np.abs(grad_np-grad_auto_np))) 93 | 94 | print ("cffi forward time: %f, backward time: %f, full time: %f " % (tf_c, tb_c, tf_c+tb_c)) 95 | print ("auto forward time: %f, backward time: %f, full time: %f " % (tf_py, tb_py, tf_py+tb_py)) 96 | print ("ratio: ", (tf_py+tb_py)/(tf_c + tb_c)) -------------------------------------------------------------------------------- /im2mesh/dmc/ops/tests/test_distance.py: -------------------------------------------------------------------------------- 1 | 2 | import sys 3 | sys.path.append('../../../..') 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.autograd import Variable 8 | import time 9 | import numpy as np 10 | import resource 11 | 12 | from im2mesh.dmc.ops.tests.loss_autograd import LossAutoGrad 13 | from im2mesh.dmc.ops.point_triangle_distance import PointTriangleDistance 14 | 15 | 16 | print("Testing CUDA extension...") 17 | dtype = torch.cuda.FloatTensor 18 | dtype_long = torch.cuda.LongTensor 19 | num_cells = 4 20 | # autograd loss 21 | loss_autograd = LossAutoGrad(num_cells, 1.0) 22 | 23 | multiGrids = PointTriangleDistance() 24 | 25 | 26 | if __name__ == '__main__': 27 | 28 | print("=========== Input =============") 29 | point = Variable(torch.rand(10, 3).view(-1,3).type(dtype) * 0.9) * num_cells 30 | offset = Variable(torch.zeros(3, num_cells+1, num_cells+1, num_cells+1).type(dtype)*0.5, requires_grad=True) 31 | print(point.shape) 32 | print(offset.shape) 33 | 34 | print("============= cuda extension ============") 35 | # forward 36 | tf_c = time.time() 37 | distance = multiGrids.forward(offset, point) 38 | tf_c = time.time() - tf_c 39 | distance_np = distance.data.cpu().numpy() 40 | print("cffi distance:") 41 | print(distance_np.shape) 42 | 43 | weight_rnd = Variable(torch.rand(distance.size()).type(dtype), requires_grad=False) 44 | distance_sum = torch.sum(torch.mul(distance, weight_rnd)) 45 | 46 | # backward 47 | tb_c = time.time() 48 | grad = distance_sum.backward() 49 | tb_c = time.time() - tb_c 50 | offset_np = np.copy(offset.grad.data.cpu().numpy()) 51 | 52 | print("cffi grad:") 53 | print(offset_np.shape) 54 | 55 | print("============= auto ============") 56 | # forward 57 | tf_py = time.time() 58 | distance_auto = loss_autograd.loss_point_to_mesh_distance_autograd(offset, point) 59 | tf_py = time.time()-tf_py 60 | distance_auto_np = distance_auto.data.cpu().numpy() 61 | print("auto distance:") 62 | print(distance_auto_np.shape) 63 | weight_rnd = Variable(weight_rnd.data) 64 | distance_sum_auto = torch.sum(torch.mul(distance_auto, weight_rnd)) 65 | 66 | # backward 67 | offset.grad.data.zero_() 68 | 69 | tb_py = time.time() 70 | distance_sum_auto.backward() 71 | tb_py = time.time() - tb_py 72 | print("auto grad: ") 73 | offset_auto_np = np.copy(offset.grad.data.cpu().numpy()) 74 | print(offset_auto_np.shape) 75 | 76 | print("========== summary ===========") 77 | print("Forward difference between cffi and auto: "+str(np.sum(np.abs(distance_np[:,:-1]-distance_auto_np[:,:-1])))) 78 | print("Backward difference between cffi and auto: "+str(np.sum(np.abs(offset_np-offset_auto_np)))) -------------------------------------------------------------------------------- /im2mesh/dmc/ops/tests/test_occupancy_connectivity.py: -------------------------------------------------------------------------------- 1 | 2 | import sys 3 | sys.path.append('../../../..') 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.autograd import Variable 8 | 9 | import time 10 | import numpy as np 11 | from im2mesh.dmc.ops.occupancy_connectivity import OccupancyConnectivity 12 | #from loss import Loss 13 | #from loss_autograd import LossAutoGrad 14 | #from parse_args import parse_args 15 | 16 | # check the cuda extension or c extension 17 | 18 | def loss_on_smoothness_autograd( occupancy): 19 | """ Compute the smoothness loss using pytorch, 20 | implemented for gradient check of the c/c++ extensions 21 | """ 22 | 23 | Wo=occupancy.size()[0] 24 | Ho=occupancy.size()[1] 25 | Do=occupancy.size()[2] 26 | 27 | loss = 0 28 | for x_ in range(Wo): 29 | for y_ in range(Ho): 30 | for z_ in range(Do): 31 | # horizontal direction 32 | if x_