├── reconstruction ├── __init__.py ├── src │ ├── __init__.py │ ├── samplers.py │ ├── tf_utils.py │ ├── ae_templates.py │ └── neural_net.py ├── autoencoder │ ├── __init__.py │ └── train_ae.py ├── external │ ├── __init__.py │ ├── sampling │ │ ├── __init__.py │ │ ├── .gitignore │ │ ├── tf_sampling_compile.sh │ │ └── tf_sampling.py │ ├── grouping │ │ ├── test │ │ │ ├── __init__.py │ │ │ ├── compile.sh │ │ │ ├── selection_sort.cu │ │ │ ├── selection_sort_const.cu │ │ │ ├── selection_sort.cpp │ │ │ ├── query_ball_point.cpp │ │ │ ├── query_ball_point_block.cu │ │ │ ├── query_ball_point.cu │ │ │ └── query_ball_point_grid.cu │ │ ├── __init__.py │ │ ├── .gitignore │ │ ├── tf_grouping_compile.sh │ │ ├── tf_grouping_op_test.py │ │ ├── tf_grouping.py │ │ └── tf_grouping_g.cu │ ├── python_plyfile │ │ ├── __init__.py │ │ └── .gitignore │ └── structural_losses │ │ ├── __init__.py │ │ ├── .gitignore │ │ ├── tf_approxmatch_compile.sh │ │ ├── tf_nndistance_compile.sh │ │ ├── tf_nndistance.py │ │ ├── tf_nndistance_g.cu │ │ └── tf_approxmatch.py ├── sampler │ ├── __init__.py │ ├── .gitignore │ └── train_samplenet.py ├── .gitignore ├── compile_ops.sh ├── download_data.sh ├── runner_samplenet.sh ├── runner_samplenet_progressive.sh └── README.md ├── classification ├── grouping │ ├── test │ │ ├── __init__.py │ │ ├── compile.sh │ │ ├── selection_sort.cu │ │ ├── selection_sort_const.cu │ │ ├── selection_sort.cpp │ │ ├── query_ball_point.cpp │ │ ├── query_ball_point_block.cu │ │ ├── query_ball_point.cu │ │ └── query_ball_point_grid.cu │ ├── __init__.py │ ├── .gitignore │ ├── tf_grouping_compile.sh │ ├── tf_grouping_op_test.py │ ├── tf_grouping.py │ └── tf_grouping_g.cu ├── structural_losses │ ├── .gitignore │ ├── __init__.py │ ├── tf_approxmatch_compile.sh │ ├── tf_nndistance_compile.sh │ ├── tf_nndistance.py │ ├── tf_nndistance_g.cu │ └── tf_approxmatch.py ├── compile_ops.sh ├── runner_samplenet.sh ├── runner_samplenet_progressive.sh ├── models │ ├── pointnet_cls.py │ ├── transform_nets.py │ ├── pointnet_cls_basic.py │ └── samplenet_model.py ├── provider.py └── README.md ├── doc └── teaser.png ├── registration ├── src │ ├── chamfer_distance │ │ ├── __init__.py │ │ ├── chamfer_distance.py │ │ └── chamfer_distance.cu │ ├── __init__.py │ ├── fps.py │ ├── random_sampling.py │ ├── sputils.py │ └── pctransforms.py ├── runner_samplenet.sh ├── Dockerfile ├── data │ ├── create_dataset_torch.py │ └── modelnet_loader_torch.py ├── models │ └── pcrnet.py └── README.md ├── .gitignore └── README.md /reconstruction/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/autoencoder/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/external/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/sampler/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /classification/grouping/test/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/external/sampling/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/external/python_plyfile/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /reconstruction/sampler/.gitignore: -------------------------------------------------------------------------------- 1 | plots 2 | *.png 3 | -------------------------------------------------------------------------------- /classification/structural_losses/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | *.pyc 4 | -------------------------------------------------------------------------------- /reconstruction/external/sampling/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | *.pyc 4 | -------------------------------------------------------------------------------- /doc/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/itailang/SampleNet/HEAD/doc/teaser.png -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | *.pyc 4 | -------------------------------------------------------------------------------- /registration/src/chamfer_distance/__init__.py: -------------------------------------------------------------------------------- 1 | from .chamfer_distance import ChamferDistance 2 | -------------------------------------------------------------------------------- /reconstruction/external/python_plyfile/.gitignore: -------------------------------------------------------------------------------- 1 | *~ 2 | *.pyc 3 | *.swp 4 | *.egg-info 5 | plyfile-venv/ 6 | build/ 7 | dist/ 8 | .tox 9 | .cache 10 | -------------------------------------------------------------------------------- /reconstruction/.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | .project 3 | .ipynb_checkpoints 4 | .DS_Store 5 | .pydevproject 6 | *.pyc 7 | *.nfs* 8 | data/* 9 | results/* 10 | log/* 11 | external/structural_losses/*.o 12 | external/structural_losses/*.so 13 | utils/*.so 14 | -------------------------------------------------------------------------------- /classification/grouping/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | 4 | try: 5 | from .tf_grouping import query_ball_point, group_point 6 | except: 7 | print("TF grouping ops (query_ball_point, group_point) were not loaded.") 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .ropeproject 2 | .idea 3 | .vscode 4 | .project 5 | .ipynb_checkpoints 6 | .DS_Store 7 | .pydevproject 8 | *.pyc 9 | *.nfs* 10 | 11 | **/data/*_hdf5_2048* 12 | **/data/ModelNet40/ 13 | **/data/shape_net_core_uniform_samples_2048/ 14 | log/ 15 | dump/ 16 | *.o 17 | *.so 18 | -------------------------------------------------------------------------------- /classification/grouping/.gitignore: -------------------------------------------------------------------------------- 1 | a.out 2 | query_ball_point 3 | query_ball_point_block 4 | query_ball_point_cuda 5 | query_ball_point_grid 6 | tf_grouping_g.cu.o 7 | tf_grouping_so.so 8 | selection_sort 9 | selection_sort_cuda 10 | selection_sort_const_cuda 11 | *.o 12 | *.so 13 | *.pyc 14 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | 4 | try: 5 | from .tf_grouping import query_ball_point, group_point 6 | except: 7 | print("TF grouping ops (query_ball_point, group_point) were not loaded.") 8 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/.gitignore: -------------------------------------------------------------------------------- 1 | a.out 2 | query_ball_point 3 | query_ball_point_block 4 | query_ball_point_cuda 5 | query_ball_point_grid 6 | tf_grouping_g.cu.o 7 | tf_grouping_so.so 8 | selection_sort 9 | selection_sort_cuda 10 | selection_sort_const_cuda 11 | *.o 12 | *.so 13 | *.pyc 14 | -------------------------------------------------------------------------------- /classification/compile_ops.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | cd ./grouping/ 4 | 5 | # compile grouping op 6 | sh tf_grouping_compile.sh 7 | 8 | cd ../structural_losses/ 9 | 10 | # compile nndistance op 11 | sh tf_nndistance_compile.sh 12 | 13 | # compile approxmatch op 14 | sh tf_approxmatch_compile.sh 15 | -------------------------------------------------------------------------------- /registration/src/__init__.py: -------------------------------------------------------------------------------- 1 | from . import quaternion, sputils, pctransforms, qdataset 2 | from .chamfer_distance.chamfer_distance import ChamferDistance 3 | from .soft_projection import SoftProjection 4 | from .samplenet import SampleNet 5 | from .fps import FPSSampler 6 | from .random_sampling import RandomSampler 7 | -------------------------------------------------------------------------------- /classification/structural_losses/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | 4 | try: 5 | from .tf_nndistance import nn_distance 6 | from .tf_approxmatch import approx_match, match_cost 7 | except: 8 | print("External Losses (Chamfer-EMD) were not loaded.") 9 | -------------------------------------------------------------------------------- /classification/grouping/test/compile.sh: -------------------------------------------------------------------------------- 1 | g++ query_ball_point.cpp -o query_ball_point 2 | nvcc query_ball_point.cu -o query_ball_point_cuda 3 | nvcc query_ball_point_block.cu -o query_ball_point_block 4 | nvcc query_ball_point_grid.cu -o query_ball_point_grid 5 | g++ -Wall selection_sort.cpp -o selection_sort 6 | nvcc selection_sort.cu -o selection_sort_cuda 7 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/compile.sh: -------------------------------------------------------------------------------- 1 | g++ query_ball_point.cpp -o query_ball_point 2 | nvcc query_ball_point.cu -o query_ball_point_cuda 3 | nvcc query_ball_point_block.cu -o query_ball_point_block 4 | nvcc query_ball_point_grid.cu -o query_ball_point_grid 5 | g++ -Wall selection_sort.cpp -o selection_sort 6 | nvcc selection_sort.cu -o selection_sort_cuda 7 | -------------------------------------------------------------------------------- /reconstruction/compile_ops.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | cd ./external/grouping/ 4 | 5 | # compile grouping op 6 | sh tf_grouping_compile.sh 7 | 8 | cd ../sampling/ 9 | 10 | # compile sampling op 11 | sh tf_sampling_compile.sh 12 | 13 | cd ../structural_losses/ 14 | 15 | # compile nndistance op 16 | sh tf_nndistance_compile.sh 17 | 18 | # compile approxmatch op 19 | sh tf_approxmatch_compile.sh 20 | -------------------------------------------------------------------------------- /reconstruction/download_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # For each Mesh model of Shape-Net-Core download 1 point-cloud with 2048 points 3 | # sampled uniformly at random (around 1.4GB). 4 | wget https://www.dropbox.com/s/vmsdrae6x5xws1v/shape_net_core_uniform_samples_2048.zip?dl=0 5 | mv shape_net_core_uniform_samples_2048.zip\?dl\=0 shape_net_core_uniform_samples_2048.zip 6 | unzip shape_net_core_uniform_samples_2048.zip 7 | rm shape_net_core_uniform_samples_2048.zip 8 | mkdir -p data 9 | mv shape_net_core_uniform_samples_2048 data 10 | -------------------------------------------------------------------------------- /registration/runner_samplenet.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Train PCR-Net 4 | python main.py \ 5 | -o log/baseline/PCRNet1024 \ 6 | --datafolder car_hdf5_2048 \ 7 | --sampler none \ 8 | --train-pcrnet \ 9 | --epochs 500 \ 10 | 11 | wait 12 | 13 | # Train SampleNet 14 | python main.py \ 15 | -o log/SAMPLENET64 \ 16 | --datafolder car_hdf5_2048 \ 17 | --transfer-from log/baseline/PCRNet1024_model_best.pth \ 18 | --sampler samplenet \ 19 | --train-samplenet \ 20 | --num-out-points 64 \ 21 | --epochs 400 \ 22 | 23 | wait 24 | -------------------------------------------------------------------------------- /classification/grouping/tf_grouping_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc tf_grouping_g.cu -o tf_grouping_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/tf_grouping_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc tf_grouping_g.cu -o tf_grouping_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /reconstruction/external/sampling/tf_sampling_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /reconstruction/runner_samplenet.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # train Autoencoder model 4 | python autoencoder/train_ae.py --train_folder log/autoencoder 5 | wait 6 | 7 | # evaluate Autoencoder model 8 | python autoencoder/evaluate_ae.py --train_folder log/autoencoder 9 | wait 10 | 11 | # train SampleNet, use Autoencoder model as the task network 12 | python sampler/train_samplenet.py --ae_folder log/autoencoder --n_sample_points 64 --train_folder log/SampleNet64 13 | wait 14 | 15 | # evaluate SampleNet 16 | python sampler/evaluate_samplenet.py --train_folder log/SampleNet64 17 | 18 | # see the results in log/SampleNet64/eval/eval_stats_test_set_multi_0064.txt 19 | -------------------------------------------------------------------------------- /classification/structural_losses/tf_approxmatch_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc tf_approxmatch_g.cu -o tf_approxmatch_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_approxmatch.cpp tf_approxmatch_g.cu.o -o tf_approxmatch_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/tf_approxmatch_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc tf_approxmatch_g.cu -o tf_approxmatch_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_approxmatch.cpp tf_approxmatch_g.cu.o -o tf_approxmatch_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /registration/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM pytorch/pytorch:1.4-cuda10.1-cudnn7-devel 2 | RUN apt update 3 | RUN apt install -y git less vim unzip 4 | 5 | # install python packages 6 | RUN pip install kornia 7 | RUN conda install -y opencv 8 | RUN pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl 9 | 10 | # Download Pointnet2_PyTorch 11 | WORKDIR /root 12 | RUN git clone https://github.com/erikwijmans/Pointnet2_PyTorch.git 13 | WORKDIR /workspace 14 | 15 | # environment encoding variables 16 | # https://stackoverflow.com/questions/55646024/writing-accented-characters-from-user-input-to-a-text-file-python-3-7 17 | ENV LANG=C.UTF-8 18 | ENV LC_ALL=C.UTF-8 19 | -------------------------------------------------------------------------------- /classification/structural_losses/tf_nndistance_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -c -o tf_nndistance_g.cu.o tf_nndistance_g.cu -I $TF_INC -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC -O2 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_nndistance.cpp tf_nndistance_g.cu.o -o tf_nndistance_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/tf_nndistance_compile.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') 4 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | 6 | /usr/local/cuda-10.0/bin/nvcc -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -c -o tf_nndistance_g.cu.o tf_nndistance_g.cu -I $TF_INC -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC -O2 7 | 8 | # TF1.13 9 | g++ -std=c++11 tf_nndistance.cpp tf_nndistance_g.cu.o -o tf_nndistance_so.so -shared -fPIC -I $TF_INC -I /usr/local/cuda-10.0/include -I $TF_INC/external/nsync/public -lcudart -L /usr/local/cuda-10.0/lib64/ -L$TF_LIB -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 10 | -------------------------------------------------------------------------------- /classification/runner_samplenet.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # train PointNet classifier 4 | python train_classifier.py --model pointnet_cls --log_dir log/PointNet1024 5 | wait 6 | 7 | # train SampleNet, use PointNet classifier as the task network 8 | python train_samplenet.py --classifier_model pointnet_cls --classifier_model_path log/PointNet1024/model.ckpt \ 9 | --num_out_points 32 --log_dir log/SampleNet32 10 | wait 11 | 12 | # infer SampleNet and evaluate PointNet classifier with sampled points of SampleNet 13 | python evaluate_samplenet.py --sampler_model_path log/SampleNet32/model.ckpt \ 14 | --num_out_points 32 --dump_dir log/SampleNet32/eval 15 | 16 | # see the results in log/SampleNet32/eval/log_evaluate.txt 17 | -------------------------------------------------------------------------------- /reconstruction/runner_samplenet_progressive.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # train Autoencoder model 4 | python autoencoder/train_ae.py --train_folder log/autoencoder 5 | wait 6 | 7 | # evaluate Autoencoder model 8 | python autoencoder/evaluate_ae.py --train_folder log/autoencoder 9 | wait 10 | 11 | # train SampleNetProgressive, use Autoencoder model as the task network 12 | python sampler/train_samplenet_progressive.py --ae_folder log/autoencoder --n_sample_points 64 --train_folder log/SampleNetProgressive 13 | wait 14 | 15 | # evaluate SampleNetProgressive 16 | python sampler/evaluate_samplenet_progressive.py --n_sample_points 64 --train_folder log/SampleNetProgressive 17 | 18 | # see the results in log/SampleNetProgressive/eval/eval_stats_test_set_multi_0064.txt 19 | -------------------------------------------------------------------------------- /classification/runner_samplenet_progressive.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # train PointNet (vanilla) classifier 4 | python train_classifier.py --model pointnet_cls_basic --log_dir log/PointNetVanilla1024 5 | wait 6 | 7 | # train SampleNetProgressive, use PointNet (vanilla) classifier as the task network 8 | python train_samplenet_progressive.py --classifier_model pointnet_cls_basic --classifier_model_path log/PointNetVanilla1024/model.ckpt \ 9 | --log_dir log/SampleNetProgressive 10 | wait 11 | 12 | # infer SampleNetProgressive and save the ordered point clouds to .h5 files 13 | python infer_samplenet_progressive.py --sampler_model_path log/SampleNetProgressive/model.ckpt \ 14 | --dump_dir log/SampleNetProgressive 15 | wait 16 | 17 | # evaluate PointNet (vanilla) classifier with sampled points of SampleNetProgressive 18 | python evaluate_from_files.py --classifier_model pointnet_cls_basic --classifier_model_path log/PointNetVanilla1024/model.ckpt \ 19 | --data_path log/SampleNetProgressive/sampled --dump_dir log/SampleNetProgressive/eval/sampled 20 | 21 | # see the results in log/SampleNetProgressive/eval/sampled/log_evaluate.txt 22 | -------------------------------------------------------------------------------- /reconstruction/src/samplers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on September 7th, 2018 3 | 4 | @author: itailang 5 | """ 6 | 7 | import tensorflow as tf 8 | import numpy as np 9 | 10 | from .encoders_decoders import encoder_with_convs_and_symmetry, decoder_with_fc_only 11 | 12 | 13 | def sampler_with_convs_and_symmetry_and_fc( 14 | in_signal, pc_dim_out, non_linearity=tf.nn.relu 15 | ): 16 | """A sampling network that generate sample points from an input point cloud.""" 17 | 18 | n_points, dummy = pc_dim_out 19 | if dummy != 3: 20 | raise ValueError() 21 | 22 | encoder_args = { 23 | "n_filters": [64, 128, 128, 256, 128], 24 | "filter_sizes": [1], 25 | "strides": [1], 26 | "non_linearity": non_linearity, 27 | "b_norm": True, 28 | "verbose": True, 29 | } 30 | layer = encoder_with_convs_and_symmetry(in_signal, **encoder_args) 31 | 32 | decoder_args = { 33 | "layer_sizes": [256, 256, np.prod([n_points, 3])], 34 | "b_norm": False, 35 | "b_norm_finish": False, 36 | "verbose": True, 37 | } 38 | out_signal = decoder_with_fc_only(layer, **decoder_args) 39 | 40 | out_signal = tf.reshape(out_signal, [-1, n_points, 3]) 41 | return out_signal 42 | -------------------------------------------------------------------------------- /classification/grouping/tf_grouping_op_test.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | import tensorflow as tf 4 | import numpy as np 5 | from .tf_grouping import query_ball_point, group_point 6 | 7 | 8 | class GroupPointTest(tf.test.TestCase): 9 | def test(self): 10 | pass 11 | 12 | def test_grad(self): 13 | with tf.device("/gpu:0"): 14 | points = tf.constant(np.random.random((1, 128, 16)).astype("float32")) 15 | print(points) 16 | xyz1 = tf.constant(np.random.random((1, 128, 3)).astype("float32")) 17 | xyz2 = tf.constant(np.random.random((1, 8, 3)).astype("float32")) 18 | radius = 0.3 19 | nsample = 32 20 | idx, pts_cnt = query_ball_point(radius, nsample, xyz1, xyz2) 21 | grouped_points = group_point(points, idx) 22 | print(grouped_points) 23 | 24 | with self.test_session(): 25 | print("---- Going to compute gradient error") 26 | err = tf.test.compute_gradient_error( 27 | points, (1, 128, 16), grouped_points, (1, 8, 32, 16) 28 | ) 29 | print(err) 30 | self.assertLess(err, 1e-4) 31 | 32 | 33 | if __name__ == "__main__": 34 | tf.test.main() 35 | -------------------------------------------------------------------------------- /reconstruction/src/tf_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on November 26, 2017 3 | 4 | @author: optas 5 | """ 6 | 7 | import tensorflow as tf 8 | import numpy as np 9 | 10 | 11 | def expand_scope_by_name(scope, name): 12 | """ expand tf scope by given name. 13 | """ 14 | 15 | if isinstance(scope, str): 16 | scope += "/" + name 17 | return scope 18 | 19 | if scope is not None: 20 | return scope.name + "/" + name 21 | else: 22 | return scope 23 | 24 | 25 | def replicate_parameter_for_all_layers(parameter, n_layers): 26 | if parameter is not None and len(parameter) != n_layers: 27 | if len(parameter) != 1: 28 | raise ValueError() 29 | parameter = np.array(parameter) 30 | parameter = parameter.repeat(n_layers).tolist() 31 | return parameter 32 | 33 | 34 | def reset_tf_graph(): 35 | """ Reset's all variables of default-tf graph. Useful for jupyter. 36 | """ 37 | if "sess" in globals() and sess: 38 | sess.close() 39 | tf.reset_default_graph() 40 | 41 | 42 | def leaky_relu(alpha): 43 | if not (alpha < 1 and alpha > 0): 44 | raise ValueError() 45 | 46 | return lambda x: tf.maximum(alpha * x, x) 47 | 48 | 49 | def safe_log(x, eps=1e-12): 50 | return tf.log(tf.maximum(x, eps)) 51 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/tf_grouping_op_test.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | import tensorflow as tf 4 | import numpy as np 5 | from .tf_grouping import query_ball_point, group_point 6 | 7 | 8 | class GroupPointTest(tf.test.TestCase): 9 | def test(self): 10 | pass 11 | 12 | def test_grad(self): 13 | with tf.device("/gpu:0"): 14 | points = tf.constant(np.random.random((1, 128, 16)).astype("float32")) 15 | print(points) 16 | xyz1 = tf.constant(np.random.random((1, 128, 3)).astype("float32")) 17 | xyz2 = tf.constant(np.random.random((1, 8, 3)).astype("float32")) 18 | radius = 0.3 19 | nsample = 32 20 | idx, pts_cnt = query_ball_point(radius, nsample, xyz1, xyz2) 21 | grouped_points = group_point(points, idx) 22 | print(grouped_points) 23 | 24 | with self.test_session(): 25 | print("---- Going to compute gradient error") 26 | err = tf.test.compute_gradient_error( 27 | points, (1, 128, 16), grouped_points, (1, 8, 32, 16) 28 | ) 29 | print(err) 30 | self.assertLess(err, 1e-4) 31 | 32 | 33 | if __name__ == "__main__": 34 | tf.test.main() 35 | -------------------------------------------------------------------------------- /reconstruction/src/ae_templates.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on September 2, 2017 3 | 4 | @author: optas 5 | """ 6 | import numpy as np 7 | 8 | from .encoders_decoders import encoder_with_convs_and_symmetry, decoder_with_fc_only 9 | 10 | 11 | def mlp_architecture_ala_iclr_18( 12 | n_pc_points, bneck_size, bneck_post_mlp=False, check_n_pc_points=True 13 | ): 14 | """ Single class experiments. 15 | """ 16 | if check_n_pc_points and n_pc_points != 2048: 17 | raise ValueError() 18 | 19 | encoder = encoder_with_convs_and_symmetry 20 | decoder = decoder_with_fc_only 21 | 22 | n_input = [n_pc_points, 3] 23 | 24 | encoder_args = { 25 | "n_filters": [64, 128, 128, 256, bneck_size], 26 | "filter_sizes": [1], 27 | "strides": [1], 28 | "b_norm": True, 29 | "verbose": True, 30 | } 31 | 32 | decoder_args = { 33 | "layer_sizes": [256, 256, np.prod(n_input)], 34 | "b_norm": False, 35 | "b_norm_finish": False, 36 | "verbose": True, 37 | } 38 | 39 | if bneck_post_mlp: 40 | encoder_args["n_filters"].pop() 41 | decoder_args["layer_sizes"][0] = bneck_size 42 | 43 | return encoder, decoder, encoder_args, decoder_args 44 | 45 | 46 | def default_train_params(): 47 | params = { 48 | "batch_size": 50, 49 | "training_epochs": 500, 50 | "denoising": False, 51 | "learning_rate": 0.0005, 52 | "z_rotate": False, 53 | "saver_step": 50, 54 | "loss_display_step": 1, 55 | } 56 | return params 57 | -------------------------------------------------------------------------------- /reconstruction/src/neural_net.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on August 28, 2017 3 | 4 | @author: optas 5 | """ 6 | from __future__ import print_function 7 | 8 | from builtins import str 9 | from builtins import object 10 | import os.path as osp 11 | import tensorflow as tf 12 | import warnings 13 | 14 | MODEL_SAVER_ID = "models.ckpt" 15 | 16 | 17 | class Neural_Net(object): 18 | def __init__(self, name, graph): 19 | if graph is None: 20 | graph = tf.get_default_graph() 21 | 22 | self.graph = graph 23 | self.name = name 24 | 25 | with tf.variable_scope(name): 26 | with tf.device("/cpu:0"): 27 | self.epoch = tf.get_variable( 28 | "epoch", [], initializer=tf.constant_initializer(0), trainable=False 29 | ) 30 | self.increment_epoch = self.epoch.assign_add(tf.constant(1.0)) 31 | 32 | self.no_op = tf.no_op() 33 | 34 | def is_training(self): 35 | is_training_op = self.graph.get_collection("is_training") 36 | return self.sess.run(is_training_op)[0] 37 | 38 | def restore_model(self, model_path, epoch, verbose=False): 39 | """Restore all the variables of a saved model. 40 | """ 41 | self.saver.restore( 42 | self.sess, osp.join(model_path, MODEL_SAVER_ID + "-" + str(int(epoch))) 43 | ) 44 | 45 | if self.epoch.eval(session=self.sess) != epoch: 46 | warnings.warn("Loaded model's epoch doesn't match the requested one.") 47 | else: 48 | if verbose: 49 | print("Model restored in epoch {0}.".format(epoch)) 50 | -------------------------------------------------------------------------------- /registration/src/fps.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import warnings 3 | 4 | from pointnet2.utils.pointnet2_utils import furthest_point_sample as fps 5 | from pointnet2.utils.pointnet2_utils import gather_operation as gather 6 | 7 | 8 | class FPSSampler(torch.nn.Module): 9 | def __init__(self, num_out_points, permute, input_shape="bcn", output_shape="bcn"): 10 | super().__init__() 11 | self.num_out_points = num_out_points 12 | self.permute = permute 13 | self.name = "fps" 14 | 15 | # input / output shapes 16 | if input_shape not in ["bcn", "bnc"]: 17 | raise ValueError( 18 | "allowed shape are 'bcn' (batch * channels * num_in_points), 'bnc' " 19 | ) 20 | if output_shape not in ["bcn", "bnc"]: 21 | raise ValueError( 22 | "allowed shape are 'bcn' (batch * channels * num_in_points), 'bnc' " 23 | ) 24 | if input_shape != output_shape: 25 | warnings.warn("FPS: input_shape is different to output_shape.") 26 | self.input_shape = input_shape 27 | self.output_shape = output_shape 28 | 29 | def forward(self, x: torch.Tensor): 30 | # x shape should be B x 3 x N 31 | if self.permute: 32 | _, N, _ = x.shape 33 | x = x[:, torch.randperm(N), :] 34 | 35 | idx = fps(x, self.num_out_points) 36 | 37 | if self.input_shape == "bnc": 38 | x = x.permute(0, 2, 1).contiguous() 39 | y = gather(x, idx) 40 | if self.output_shape == "bnc": 41 | y = y.permute(0, 2, 1).contiguous() 42 | 43 | return y 44 | -------------------------------------------------------------------------------- /registration/src/random_sampling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import warnings 3 | 4 | from pointnet2.utils.pointnet2_utils import gather_operation as gather 5 | 6 | 7 | class RandomSampler(torch.nn.Module): 8 | def __init__(self, num_out_points, input_shape="bcn", output_shape="bcn"): 9 | super().__init__() 10 | self.num_out_points = num_out_points 11 | self.name = "random" 12 | 13 | # input / output shapes 14 | if input_shape not in ["bcn", "bnc"]: 15 | raise ValueError( 16 | "allowed shape are 'bcn' (batch * channels * num_in_points), 'bnc' " 17 | ) 18 | if output_shape not in ["bcn", "bnc"]: 19 | raise ValueError( 20 | "allowed shape are 'bcn' (batch * channels * num_in_points), 'bnc' " 21 | ) 22 | if input_shape != output_shape: 23 | warnings.warn("RandomSampler: input_shape is different to output_shape.") 24 | self.input_shape = input_shape 25 | self.output_shape = output_shape 26 | 27 | def forward(self, x: torch.Tensor): 28 | if self.input_shape == "bnc": 29 | x = x.permute(0, 2, 1).contiguous() 30 | 31 | B, _, N = x.shape 32 | idx = torch.zeros(B, self.num_out_points, dtype=torch.int32, device=x.device) 33 | for i in range(B): 34 | rand_perm = torch.randperm( 35 | N, 36 | dtype=torch.int32, 37 | device=x.device 38 | ) 39 | idx[i] = rand_perm[:self.num_out_points] 40 | 41 | 42 | y = gather(x, idx) 43 | if self.output_shape == "bnc": 44 | y = y.permute(0, 2, 1).contiguous() 45 | 46 | return y 47 | -------------------------------------------------------------------------------- /registration/src/chamfer_distance/chamfer_distance.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.cpp_extension import load 3 | import os 4 | 5 | script_dir = os.path.dirname(__file__) 6 | sources = [ 7 | os.path.join(script_dir, "chamfer_distance.cpp"), 8 | os.path.join(script_dir, "chamfer_distance.cu"), 9 | ] 10 | 11 | cd = load(name="cd", sources=sources) 12 | 13 | 14 | class ChamferDistanceFunction(torch.autograd.Function): 15 | @staticmethod 16 | def forward(ctx, xyz1, xyz2): 17 | batchsize, n, _ = xyz1.size() 18 | _, m, _ = xyz2.size() 19 | xyz1 = xyz1.contiguous() 20 | xyz2 = xyz2.contiguous() 21 | dist1 = torch.zeros(batchsize, n) 22 | dist2 = torch.zeros(batchsize, m) 23 | 24 | idx1 = torch.zeros(batchsize, n, dtype=torch.int) 25 | idx2 = torch.zeros(batchsize, m, dtype=torch.int) 26 | 27 | if not xyz1.is_cuda: 28 | cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 29 | else: 30 | dist1 = dist1.cuda() 31 | dist2 = dist2.cuda() 32 | idx1 = idx1.cuda() 33 | idx2 = idx2.cuda() 34 | cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2) 35 | 36 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 37 | 38 | return dist1, dist2 39 | 40 | @staticmethod 41 | def backward(ctx, graddist1, graddist2): 42 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 43 | 44 | graddist1 = graddist1.contiguous() 45 | graddist2 = graddist2.contiguous() 46 | 47 | gradxyz1 = torch.zeros(xyz1.size()) 48 | gradxyz2 = torch.zeros(xyz2.size()) 49 | 50 | if not graddist1.is_cuda: 51 | cd.backward( 52 | xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 53 | ) 54 | else: 55 | gradxyz1 = gradxyz1.cuda() 56 | gradxyz2 = gradxyz2.cuda() 57 | cd.backward_cuda( 58 | xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 59 | ) 60 | 61 | return gradxyz1, gradxyz2 62 | 63 | 64 | class ChamferDistance(torch.nn.Module): 65 | def forward(self, xyz1, xyz2): 66 | return ChamferDistanceFunction.apply(xyz1, xyz2) 67 | -------------------------------------------------------------------------------- /classification/grouping/test/selection_sort.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,k), val (b,m,k) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, float *dist, int *idx, float *val) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | idx+=m*k*batch_index; 24 | val+=m*k*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | float *p_dist; 30 | for (int j=index;j>>(b,n,m,k,dist,idx,val); 68 | cudaDeviceSynchronize(); 69 | printf("selection sort cpu time %f\n",get_time()-t0); 70 | 71 | return 0; 72 | } 73 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/selection_sort.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,k), val (b,m,k) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, float *dist, int *idx, float *val) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | idx+=m*k*batch_index; 24 | val+=m*k*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | float *p_dist; 30 | for (int j=index;j>>(b,n,m,k,dist,idx,val); 68 | cudaDeviceSynchronize(); 69 | printf("selection sort cpu time %f\n",get_time()-t0); 70 | 71 | return 0; 72 | } 73 | -------------------------------------------------------------------------------- /registration/data/create_dataset_torch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import json 3 | import h5py 4 | import os 5 | 6 | from modelnet_loader_torch import ModelNetCls 7 | 8 | CATEGORY = "car" 9 | FOLDER40 = "modelnet40_ply_hdf5_2048" 10 | CATFILE40 = "data/modelnet40_ply_hdf5_2048/shape_names.txt" 11 | # CATFILE10 = 'data/modelnet10_hdf5_2048/shape_names_10.txt' 12 | NUMPOINTS = 2048 13 | BATCHSIZE = 32 14 | 15 | 16 | def save_h5(h5_filename, data, label, data_dtype="uint8", label_dtype="uint8"): 17 | h5_fout = h5py.File(h5_filename, "w") 18 | h5_fout.create_dataset( 19 | "data", data=data, compression="gzip", compression_opts=4, dtype=data_dtype 20 | ) 21 | h5_fout.create_dataset( 22 | "label", data=label, compression="gzip", compression_opts=1, dtype=label_dtype 23 | ) 24 | h5_fout.close() 25 | 26 | 27 | # Assert dataset exists. 28 | ModelNetCls( 29 | NUMPOINTS, 30 | transforms=None, 31 | train=True, 32 | download=True, 33 | folder=FOLDER40, 34 | include_shapes=False, 35 | ) 36 | 37 | all_categories = [line.rstrip("\n") for line in open(CATFILE40)] 38 | # allowed_catagories = [line.rstrip('\n') for line in open(CATFILE10)] 39 | allowed_catagories = [CATEGORY] 40 | 41 | # load the dataset 42 | for T, b in zip(("train", "test"), (True, False)): 43 | dataset = ModelNetCls( 44 | NUMPOINTS, 45 | transforms=None, 46 | train=b, 47 | download=False, 48 | folder=FOLDER40, 49 | include_shapes=True, 50 | ) 51 | 52 | # filter correct classes 53 | newset = [x for x in dataset if all_categories[x[1]] in allowed_catagories] 54 | loader = torch.utils.data.DataLoader(newset, batch_size=BATCHSIZE, shuffle=True) 55 | 56 | try: 57 | flist = open(f"data/{CATEGORY}_hdf5_2048/{T}_files.txt", "w") 58 | except FileNotFoundError: 59 | os.mkdir(f"data/{CATEGORY}_hdf5_2048/") 60 | flist = open(f"data/{CATEGORY}_hdf5_2048/{T}_files.txt", "w") 61 | 62 | for n, databatch in enumerate(loader): 63 | h5name = f"data/{CATEGORY}_hdf5_2048/{T}{n}.h5" 64 | jname = f"data/{CATEGORY}_hdf5_2048/ply_data_{T}_{n}_id2file.json" 65 | 66 | data, label, model = databatch 67 | 68 | save_h5(h5name, data, label, data_dtype="float32", label_dtype="uint8") 69 | with open(jname, "w") as write_file: 70 | json.dump(model, write_file) 71 | flist.write(f"{h5name}\n") 72 | print(f"saved {h5name}") 73 | 74 | flist.close() 75 | -------------------------------------------------------------------------------- /registration/src/sputils.py: -------------------------------------------------------------------------------- 1 | """Utility functions for SampleNet that can be shared between PyTorch and Tensorflow implementations""" 2 | 3 | import numpy as np 4 | import argparse 5 | 6 | 7 | def _calc_distances(p0, points): 8 | return ((p0 - points) ** 2).sum(axis=1) 9 | 10 | 11 | def _fps_from_given_pc(pts, k, given_pc): 12 | farthest_pts = np.zeros((k, 3)) 13 | t = np.size(given_pc) // 3 14 | farthest_pts[0:t] = given_pc 15 | 16 | distances = _calc_distances(farthest_pts[0], pts) 17 | for i in range(1, t): 18 | distances = np.minimum(distances, _calc_distances(farthest_pts[i], pts)) 19 | 20 | for i in range(t, k): 21 | farthest_pts[i] = pts[np.argmax(distances)] 22 | distances = np.minimum(distances, _calc_distances(farthest_pts[i], pts)) 23 | return farthest_pts 24 | 25 | 26 | def _unique(arr): 27 | _, idx = np.unique(arr, return_index=True) 28 | return arr[np.sort(idx)] 29 | 30 | 31 | def nn_matching(full_pc, idx, k, complete_fps=True): 32 | batch_size = np.size(full_pc, 0) 33 | out_pc = np.zeros((full_pc.shape[0], k, 3)) 34 | for ii in range(0, batch_size): 35 | best_idx = idx[ii] 36 | if complete_fps: 37 | best_idx = _unique(best_idx) 38 | out_pc[ii] = _fps_from_given_pc(full_pc[ii], k, full_pc[ii][best_idx]) 39 | else: 40 | out_pc[ii] = full_pc[ii][best_idx] 41 | return out_pc[:, 0:k, :] 42 | 43 | 44 | # fmt: off 45 | def get_parser(): 46 | parser = argparse.ArgumentParser("SampleNet: Differentiable Point Cloud Sampling") 47 | 48 | parser.add_argument("--skip-projection", action="store_true", help="Do not project points in training") 49 | 50 | parser.add_argument("-in", "--num-in-points", type=int, default=1024, help="Number of input Points [default: 1024]") 51 | parser.add_argument("-out", "--num-out-points", type=int, default=64, help="Number of output points [2, 1024] [default: 64]") 52 | parser.add_argument("--bottleneck-size", type=int, default=128, help="bottleneck size [default: 128]") 53 | parser.add_argument("--alpha", type=float, default=0.01, help="Simplification regularization loss weight [default: 0.01]") 54 | parser.add_argument("--gamma", type=float, default=1, help="Lb constant regularization loss weight [default: 1]") 55 | parser.add_argument("--delta", type=float, default=0, help="Lb linear regularization loss weight [default: 0]") 56 | 57 | # projection arguments 58 | parser.add_argument("-gs", "--projection-group-size", type=int, default=8, help='Neighborhood size in Soft Projection [default: 8]') 59 | parser.add_argument("--lmbda", type=float, default=0.01, help="Projection regularization loss weight [default: 0.01]") 60 | 61 | return parser 62 | # fmt: on 63 | -------------------------------------------------------------------------------- /classification/grouping/test/selection_sort_const.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), dist_out (b,m,n) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | outi+=m*n*batch_index; 24 | out+=m*n*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | // copy from dist to dist_out 30 | for (int j=index;j>>(b,n,m,k,dist,idx,dist_out); 84 | cudaDeviceSynchronize(); 85 | printf("selection sort cpu time %f\n",get_time()-t0); 86 | 87 | //for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), dist_out (b,m,n) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | outi+=m*n*batch_index; 24 | out+=m*n*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | // copy from dist to dist_out 30 | for (int j=index;j>>(b,n,m,k,dist,idx,dist_out); 84 | cudaDeviceSynchronize(); 85 | printf("selection sort cpu time %f\n",get_time()-t0); 86 | 87 | //for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), val (b,m,n) 20 | void selection_sort_cpu(int b, int n, int m, int k, const float *dist, int *idx, float *val) { 21 | float *p_dist; 22 | float tmp; 23 | int tmpi; 24 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), val (b,m,n) 20 | void selection_sort_cpu(int b, int n, int m, int k, const float *dist, int *idx, float *val) { 21 | float *p_dist; 22 | float tmp; 23 | int tmpi; 24 | for (int i=0;iDocker Engine and Nvidia-Docker. 3 | At this moment, the container can only run on a CUDA (_linux/amd64_) enabled machine due to specific compiled ops from Pointnet2_PyTorch. 4 | 5 | ### Pull and run the Docker container 6 | ```bash 7 | docker pull asafmanor/pytorch:samplenetreg_torch1.4 8 | docker run --runtime nvidia -v $(pwd):/workspace/ -it --name samplenetreg asafmanor/pytorch:samplenetreg_torch1.4 9 | ``` 10 | 11 | ### Alternatively, build your own Docker image 12 | #### On the host machine 13 | ```bash 14 | docker build -t samplenetreg_torch1.4_image . 15 | docker run --runtime nvidia -v $(pwd):/workspace/ -it --name samplenetreg samplenetreg_torch1.4_image 16 | ``` 17 | #### Inside the Docker container 18 | ```bash 19 | cd /root/Pointnet2_PyTorch 20 | git checkout 5ff4382f56a8cbed2b5edd3572f97436271aba89 21 | pip install -r requirements.txt 22 | pip install -e . 23 | cd /workspace 24 | ``` 25 | 26 | ## Usage 27 | ### Data preparation 28 | Create the 'car' dataset (ModelNet40 data will automatically be downloaded to `data/modelnet40_ply_hdf5_2048` if needed) and log directories: 29 | ```bash 30 | mkdir log 31 | mkdir log/baseline 32 | python data/create_dataset_torch.py 33 | ``` 34 | Point clouds of ModelNet40 models in HDF5 files (provided by Qi et al.) will be automatically downloaded (416MB) to the data folder. Each point cloud contains 2048 points uniformly sampled from a shape surface. Each cloud is zero-mean and normalized into an unit sphere. There are also text files in `data/modelnet40_ply_hdf5_2048` specifying the ids of shapes in h5 files. 35 | 36 | ### Training and evaluating 37 | For a quick start please use: 38 | ```bash 39 | sh runner_samplenet.sh 40 | ``` 41 | 42 | ### Train *PCRNet* (supervised) registration network 43 | To train a *PCRNet* model to register point clouds, use: 44 | ```bash 45 | python main.py -o log/baseline/PCRNet1024 --datafolder car_hdf5_2048 --sampler none --train-pcrnet --epochs 500 46 | ``` 47 | 48 | ### Train SampleNet 49 | To train SampleNet (with sample size 64 in this example), using an existing PCRNet as the task network, use: 50 | ```bash 51 | python main.py -o log/SAMPLENET64 --datafolder car_hdf5_2048 --transfer-from log/baseline/PCRNet1024_model_best.pth --sampler samplenet --train-samplenet --num-out-points 64 52 | ``` 53 | 54 | ### Evaluate SampleNet 55 | To evaluate PCRNet with SampleNet's sampled points (with sample size 64 in this example), use: 56 | ```bash 57 | python main.py -o log/SAMPLENET64 --datafolder car_hdf5_2048 --pretrained log/SAMPLENET64_model_best.pth --sampler samplenet --num-out-points 64 --test 58 | ``` 59 | 60 | Additional options for training and evaluating can be found using `python main.py --help`. 61 | 62 | ## Acknowledgment 63 | This code builds upon the code provided in PointNetLK, Pointnet2_PyTorch and KNN_CUDA. We thank the authors for sharing their code. 64 | 65 | -------------------------------------------------------------------------------- /classification/structural_losses/tf_nndistance.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from builtins import range 3 | import tensorflow as tf 4 | from tensorflow.python.framework import ops 5 | import os.path as osp 6 | 7 | base_dir = osp.dirname(osp.abspath(__file__)) 8 | 9 | nn_distance_module = tf.load_op_library(osp.join(base_dir, "tf_nndistance_so.so")) 10 | 11 | 12 | def nn_distance(xyz1, xyz2): 13 | """ 14 | Computes the distance of nearest neighbors for a pair of point clouds 15 | input: xyz1: (batch_size,#points_1,3) the first point cloud 16 | input: xyz2: (batch_size,#points_2,3) the second point cloud 17 | output: dist1: (batch_size,#point_1) distance from first to second 18 | output: idx1: (batch_size,#point_1) nearest neighbor from first to second 19 | output: dist2: (batch_size,#point_2) distance from second to first 20 | output: idx2: (batch_size,#point_2) nearest neighbor from second to first 21 | """ 22 | 23 | return nn_distance_module.nn_distance(xyz1, xyz2) 24 | 25 | 26 | # @tf.RegisterShape('NnDistance') 27 | @ops.RegisterShape("NnDistance") 28 | def _nn_distance_shape(op): 29 | shape1 = op.inputs[0].get_shape().with_rank(3) 30 | shape2 = op.inputs[1].get_shape().with_rank(3) 31 | return [ 32 | tf.TensorShape([shape1.dims[0], shape1.dims[1]]), 33 | tf.TensorShape([shape1.dims[0], shape1.dims[1]]), 34 | tf.TensorShape([shape2.dims[0], shape2.dims[1]]), 35 | tf.TensorShape([shape2.dims[0], shape2.dims[1]]), 36 | ] 37 | 38 | 39 | @ops.RegisterGradient("NnDistance") 40 | def _nn_distance_grad(op, grad_dist1, grad_idx1, grad_dist2, grad_idx2): 41 | xyz1 = op.inputs[0] 42 | xyz2 = op.inputs[1] 43 | idx1 = op.outputs[1] 44 | idx2 = op.outputs[3] 45 | return nn_distance_module.nn_distance_grad( 46 | xyz1, xyz2, grad_dist1, idx1, grad_dist2, idx2 47 | ) 48 | 49 | 50 | if __name__ == "__main__": 51 | import numpy as np 52 | import random 53 | import time 54 | from tensorflow.python.kernel_tests.gradient_checker import compute_gradient 55 | 56 | random.seed(100) 57 | np.random.seed(100) 58 | with tf.Session("") as sess: 59 | xyz1 = np.random.randn(32, 16384, 3).astype("float32") 60 | xyz2 = np.random.randn(32, 1024, 3).astype("float32") 61 | with tf.device("/gpu:0"): 62 | inp1 = tf.Variable(xyz1) 63 | inp2 = tf.constant(xyz2) 64 | reta, retb, retc, retd = nn_distance(inp1, inp2) 65 | loss = tf.reduce_sum(reta) + tf.reduce_sum(retc) 66 | train = tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss) 67 | sess.run(tf.initialize_all_variables()) 68 | t0 = time.time() 69 | t1 = t0 70 | best = 1e100 71 | for i in range(100): 72 | trainloss, _ = sess.run([loss, train]) 73 | newt = time.time() 74 | best = min(best, newt - t1) 75 | print((i, trainloss, (newt - t0) / (i + 1), best)) 76 | t1 = newt 77 | # print(sess.run([inp1,retb,inp2,retd])) 78 | # grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2]) 79 | # for i,j in grads: 80 | # print(i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j))) 81 | # for i in xrange(10): 82 | # t0=time.time() 83 | # a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2}) 84 | # print('time',time.time()-t0) 85 | # print(a.shape,b.shape,c.shape,d.shape) 86 | # print(a.dtype,b.dtype,c.dtype,d.dtype) 87 | # samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32') 88 | # dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 89 | # idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 90 | # print(np.abs(dist1-a[:,samples]).max()) 91 | # print(np.abs(idx1-b[:,samples]).max()) 92 | # dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 93 | # idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 94 | # print(np.abs(dist2-c[:,samples]).max()) 95 | # print(np.abs(idx2-d[:,samples]).max()) 96 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/tf_nndistance.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from builtins import range 3 | import tensorflow as tf 4 | from tensorflow.python.framework import ops 5 | import os.path as osp 6 | 7 | base_dir = osp.dirname(osp.abspath(__file__)) 8 | 9 | nn_distance_module = tf.load_op_library(osp.join(base_dir, "tf_nndistance_so.so")) 10 | 11 | 12 | def nn_distance(xyz1, xyz2): 13 | """ 14 | Computes the distance of nearest neighbors for a pair of point clouds 15 | input: xyz1: (batch_size,#points_1,3) the first point cloud 16 | input: xyz2: (batch_size,#points_2,3) the second point cloud 17 | output: dist1: (batch_size,#point_1) distance from first to second 18 | output: idx1: (batch_size,#point_1) nearest neighbor from first to second 19 | output: dist2: (batch_size,#point_2) distance from second to first 20 | output: idx2: (batch_size,#point_2) nearest neighbor from second to first 21 | """ 22 | 23 | return nn_distance_module.nn_distance(xyz1, xyz2) 24 | 25 | 26 | # @tf.RegisterShape('NnDistance') 27 | @ops.RegisterShape("NnDistance") 28 | def _nn_distance_shape(op): 29 | shape1 = op.inputs[0].get_shape().with_rank(3) 30 | shape2 = op.inputs[1].get_shape().with_rank(3) 31 | return [ 32 | tf.TensorShape([shape1.dims[0], shape1.dims[1]]), 33 | tf.TensorShape([shape1.dims[0], shape1.dims[1]]), 34 | tf.TensorShape([shape2.dims[0], shape2.dims[1]]), 35 | tf.TensorShape([shape2.dims[0], shape2.dims[1]]), 36 | ] 37 | 38 | 39 | @ops.RegisterGradient("NnDistance") 40 | def _nn_distance_grad(op, grad_dist1, grad_idx1, grad_dist2, grad_idx2): 41 | xyz1 = op.inputs[0] 42 | xyz2 = op.inputs[1] 43 | idx1 = op.outputs[1] 44 | idx2 = op.outputs[3] 45 | return nn_distance_module.nn_distance_grad( 46 | xyz1, xyz2, grad_dist1, idx1, grad_dist2, idx2 47 | ) 48 | 49 | 50 | if __name__ == "__main__": 51 | import numpy as np 52 | import random 53 | import time 54 | from tensorflow.python.kernel_tests.gradient_checker import compute_gradient 55 | 56 | random.seed(100) 57 | np.random.seed(100) 58 | with tf.Session("") as sess: 59 | xyz1 = np.random.randn(32, 16384, 3).astype("float32") 60 | xyz2 = np.random.randn(32, 1024, 3).astype("float32") 61 | with tf.device("/gpu:0"): 62 | inp1 = tf.Variable(xyz1) 63 | inp2 = tf.constant(xyz2) 64 | reta, retb, retc, retd = nn_distance(inp1, inp2) 65 | loss = tf.reduce_sum(reta) + tf.reduce_sum(retc) 66 | train = tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss) 67 | sess.run(tf.initialize_all_variables()) 68 | t0 = time.time() 69 | t1 = t0 70 | best = 1e100 71 | for i in range(100): 72 | trainloss, _ = sess.run([loss, train]) 73 | newt = time.time() 74 | best = min(best, newt - t1) 75 | print((i, trainloss, (newt - t0) / (i + 1), best)) 76 | t1 = newt 77 | # print(sess.run([inp1,retb,inp2,retd])) 78 | # grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2]) 79 | # for i,j in grads: 80 | # print(i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j))) 81 | # for i in xrange(10): 82 | # t0=time.time() 83 | # a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2}) 84 | # print('time',time.time()-t0) 85 | # print(a.shape,b.shape,c.shape,d.shape) 86 | # print(a.dtype,b.dtype,c.dtype,d.dtype) 87 | # samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32') 88 | # dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 89 | # idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 90 | # print(np.abs(dist1-a[:,samples]).max()) 91 | # print(np.abs(idx1-b[:,samples]).max()) 92 | # dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 93 | # idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 94 | # print(np.abs(dist2-c[:,samples]).max()) 95 | # print(np.abs(idx2-d[:,samples]).max()) 96 | -------------------------------------------------------------------------------- /classification/grouping/test/query_ball_point.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | void query_ball_point_cpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | void query_ball_point_cpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int index = threadIdx.x; 21 | xyz1 += n*3*index; 22 | xyz2 += m*3*index; 23 | idx += m*nsample*index; 24 | 25 | for (int j=0;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,b>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,b>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/query_ball_point_block.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int index = threadIdx.x; 21 | xyz1 += n*3*index; 22 | xyz2 += m*3*index; 23 | idx += m*nsample*index; 24 | 25 | for (int j=0;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,b>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,b>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /classification/grouping/tf_grouping.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from builtins import range 3 | import tensorflow as tf 4 | from tensorflow.python.framework import ops 5 | import sys 6 | import os 7 | 8 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 9 | sys.path.append(BASE_DIR) 10 | grouping_module = tf.load_op_library(os.path.join(BASE_DIR, "tf_grouping_so.so")) 11 | 12 | 13 | def query_ball_point(radius, nsample, xyz1, xyz2): 14 | """ 15 | Input: 16 | radius: float32, ball search radius 17 | nsample: int32, number of points selected in each ball region 18 | xyz1: (batch_size, ndataset, 3) float32 array, input points 19 | xyz2: (batch_size, npoint, 3) float32 array, query points 20 | Output: 21 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 22 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 23 | """ 24 | # return grouping_module.query_ball_point(radius, nsample, xyz1, xyz2) 25 | return grouping_module.query_ball_point(xyz1, xyz2, radius, nsample) 26 | 27 | 28 | ops.NoGradient("QueryBallPoint") 29 | 30 | 31 | def select_top_k(k, dist): 32 | """ 33 | Input: 34 | k: int32, number of k SMALLEST elements selected 35 | dist: (b,m,n) float32 array, distance matrix, m query points, n dataset points 36 | Output: 37 | idx: (b,m,n) int32 array, first k in n are indices to the top k 38 | dist_out: (b,m,n) float32 array, first k in n are the top k 39 | """ 40 | return grouping_module.selection_sort(dist, k) 41 | 42 | 43 | ops.NoGradient("SelectionSort") 44 | 45 | 46 | def group_point(points, idx): 47 | """ 48 | Input: 49 | points: (batch_size, ndataset, channel) float32 array, points to sample from 50 | idx: (batch_size, npoint, nsample) int32 array, indices to points 51 | Output: 52 | out: (batch_size, npoint, nsample, channel) float32 array, values sampled from points 53 | """ 54 | return grouping_module.group_point(points, idx) 55 | 56 | 57 | @tf.RegisterGradient("GroupPoint") 58 | def _group_point_grad(op, grad_out): 59 | points = op.inputs[0] 60 | idx = op.inputs[1] 61 | return [grouping_module.group_point_grad(points, idx, grad_out), None] 62 | 63 | 64 | def knn_point(k, xyz1, xyz2): 65 | """ 66 | Input: 67 | k: int32, number of k in k-nn search 68 | xyz1: (batch_size, ndataset, c) float32 array, input points 69 | xyz2: (batch_size, npoint, c) float32 array, query points 70 | Output: 71 | val: (batch_size, npoint, k) float32 array, L2 distances 72 | idx: (batch_size, npoint, k) int32 array, indices to input points 73 | """ 74 | b = xyz1.get_shape()[0].value 75 | n = xyz1.get_shape()[1].value 76 | c = xyz1.get_shape()[2].value 77 | m = xyz2.get_shape()[1].value 78 | print((b, n, c, m)) 79 | print((xyz1, (b, 1, n, c))) 80 | # xyz1 = tf.tile(tf.reshape(xyz1, (b,1,n,c)), [1,m,1,1]) # this does not work when the batch size is unknown (None) 81 | # xyz2 = tf.tile(tf.reshape(xyz2, (b,m,1,c)), [1,1,n,1]) # this does not work when the batch size is unknown (None) 82 | xyz1 = tf.tile(tf.expand_dims(xyz1, axis=1), [1, m, 1, 1]) 83 | xyz2 = tf.tile(tf.expand_dims(xyz2, axis=2), [1, 1, n, 1]) 84 | dist = tf.reduce_sum((xyz1 - xyz2) ** 2, -1) 85 | print((dist, k)) 86 | outi, out = select_top_k(k, dist) 87 | idx = tf.slice(outi, [0, 0, 0], [-1, -1, k]) 88 | val = tf.slice(out, [0, 0, 0], [-1, -1, k]) 89 | print((idx, val)) 90 | # val, idx = tf.nn.top_k(-dist, k=k) # ONLY SUPPORT CPU 91 | return val, idx 92 | 93 | 94 | if __name__ == "__main__": 95 | knn = True 96 | import numpy as np 97 | import time 98 | 99 | np.random.seed(100) 100 | pts = np.random.random((32, 512, 64)).astype("float32") 101 | tmp1 = np.random.random((32, 512, 3)).astype("float32") 102 | tmp2 = np.random.random((32, 128, 3)).astype("float32") 103 | with tf.device("/gpu:0"): 104 | points = tf.constant(pts) 105 | xyz1 = tf.constant(tmp1) 106 | xyz2 = tf.constant(tmp2) 107 | radius = 0.1 108 | nsample = 64 109 | if knn: 110 | _, idx = knn_point(nsample, xyz1, xyz2) 111 | grouped_points = group_point(points, idx) 112 | else: 113 | idx, _ = query_ball_point(radius, nsample, xyz1, xyz2) 114 | grouped_points = group_point(points, idx) 115 | # grouped_points_grad = tf.ones_like(grouped_points) 116 | # points_grad = tf.gradients(grouped_points, points, grouped_points_grad) 117 | with tf.Session("") as sess: 118 | now = time.time() 119 | for _ in range(100): 120 | ret = sess.run(grouped_points) 121 | print(time.time() - now) 122 | print((ret.shape, ret.dtype)) 123 | print(ret) 124 | -------------------------------------------------------------------------------- /classification/grouping/test/query_ball_point.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,1>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/tf_grouping.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from builtins import range 3 | import tensorflow as tf 4 | from tensorflow.python.framework import ops 5 | import sys 6 | import os 7 | 8 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 9 | sys.path.append(BASE_DIR) 10 | grouping_module = tf.load_op_library(os.path.join(BASE_DIR, "tf_grouping_so.so")) 11 | 12 | 13 | def query_ball_point(radius, nsample, xyz1, xyz2): 14 | """ 15 | Input: 16 | radius: float32, ball search radius 17 | nsample: int32, number of points selected in each ball region 18 | xyz1: (batch_size, ndataset, 3) float32 array, input points 19 | xyz2: (batch_size, npoint, 3) float32 array, query points 20 | Output: 21 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 22 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 23 | """ 24 | # return grouping_module.query_ball_point(radius, nsample, xyz1, xyz2) 25 | return grouping_module.query_ball_point(xyz1, xyz2, radius, nsample) 26 | 27 | 28 | ops.NoGradient("QueryBallPoint") 29 | 30 | 31 | def select_top_k(k, dist): 32 | """ 33 | Input: 34 | k: int32, number of k SMALLEST elements selected 35 | dist: (b,m,n) float32 array, distance matrix, m query points, n dataset points 36 | Output: 37 | idx: (b,m,n) int32 array, first k in n are indices to the top k 38 | dist_out: (b,m,n) float32 array, first k in n are the top k 39 | """ 40 | return grouping_module.selection_sort(dist, k) 41 | 42 | 43 | ops.NoGradient("SelectionSort") 44 | 45 | 46 | def group_point(points, idx): 47 | """ 48 | Input: 49 | points: (batch_size, ndataset, channel) float32 array, points to sample from 50 | idx: (batch_size, npoint, nsample) int32 array, indices to points 51 | Output: 52 | out: (batch_size, npoint, nsample, channel) float32 array, values sampled from points 53 | """ 54 | return grouping_module.group_point(points, idx) 55 | 56 | 57 | @tf.RegisterGradient("GroupPoint") 58 | def _group_point_grad(op, grad_out): 59 | points = op.inputs[0] 60 | idx = op.inputs[1] 61 | return [grouping_module.group_point_grad(points, idx, grad_out), None] 62 | 63 | 64 | def knn_point(k, xyz1, xyz2): 65 | """ 66 | Input: 67 | k: int32, number of k in k-nn search 68 | xyz1: (batch_size, ndataset, c) float32 array, input points 69 | xyz2: (batch_size, npoint, c) float32 array, query points 70 | Output: 71 | val: (batch_size, npoint, k) float32 array, L2 distances 72 | idx: (batch_size, npoint, k) int32 array, indices to input points 73 | """ 74 | b = xyz1.get_shape()[0].value 75 | n = xyz1.get_shape()[1].value 76 | c = xyz1.get_shape()[2].value 77 | m = xyz2.get_shape()[1].value 78 | print((b, n, c, m)) 79 | print((xyz1, (b, 1, n, c))) 80 | # xyz1 = tf.tile(tf.reshape(xyz1, (b,1,n,c)), [1,m,1,1]) # this does not work when the batch size is unknown (None) 81 | # xyz2 = tf.tile(tf.reshape(xyz2, (b,m,1,c)), [1,1,n,1]) # this does not work when the batch size is unknown (None) 82 | xyz1 = tf.tile(tf.expand_dims(xyz1, axis=1), [1, m, 1, 1]) 83 | xyz2 = tf.tile(tf.expand_dims(xyz2, axis=2), [1, 1, n, 1]) 84 | dist = tf.reduce_sum((xyz1 - xyz2) ** 2, -1) 85 | print((dist, k)) 86 | outi, out = select_top_k(k, dist) 87 | idx = tf.slice(outi, [0, 0, 0], [-1, -1, k]) 88 | val = tf.slice(out, [0, 0, 0], [-1, -1, k]) 89 | print((idx, val)) 90 | # val, idx = tf.nn.top_k(-dist, k=k) # ONLY SUPPORT CPU 91 | return val, idx 92 | 93 | 94 | if __name__ == "__main__": 95 | knn = True 96 | import numpy as np 97 | import time 98 | 99 | np.random.seed(100) 100 | pts = np.random.random((32, 512, 64)).astype("float32") 101 | tmp1 = np.random.random((32, 512, 3)).astype("float32") 102 | tmp2 = np.random.random((32, 128, 3)).astype("float32") 103 | with tf.device("/gpu:0"): 104 | points = tf.constant(pts) 105 | xyz1 = tf.constant(tmp1) 106 | xyz2 = tf.constant(tmp2) 107 | radius = 0.1 108 | nsample = 64 109 | if knn: 110 | _, idx = knn_point(nsample, xyz1, xyz2) 111 | grouped_points = group_point(points, idx) 112 | else: 113 | idx, _ = query_ball_point(radius, nsample, xyz1, xyz2) 114 | grouped_points = group_point(points, idx) 115 | # grouped_points_grad = tf.ones_like(grouped_points) 116 | # points_grad = tf.gradients(grouped_points, points, grouped_points_grad) 117 | with tf.Session("") as sess: 118 | now = time.time() 119 | for _ in range(100): 120 | ret = sess.run(grouped_points) 121 | print(time.time() - now) 122 | print((ret.shape, ret.dtype)) 123 | print(ret) 124 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/query_ball_point.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,1>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /classification/provider.py: -------------------------------------------------------------------------------- 1 | from builtins import range 2 | import os 3 | import sys 4 | import numpy as np 5 | import h5py 6 | 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | sys.path.append(BASE_DIR) 9 | 10 | # Download dataset for point cloud classification 11 | DATA_DIR = os.path.join(BASE_DIR, "data") 12 | if not os.path.exists(DATA_DIR): 13 | os.mkdir(DATA_DIR) 14 | if not os.path.exists(os.path.join(DATA_DIR, "modelnet40_ply_hdf5_2048")): 15 | www = "https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip" 16 | zipfile = os.path.basename(www) 17 | os.system("wget %s; unzip %s" % (www, zipfile)) 18 | os.system("mv %s %s" % (zipfile[:-4], DATA_DIR)) 19 | os.system("rm %s" % (zipfile)) 20 | 21 | 22 | def shuffle_data(data, labels): 23 | """ Shuffle data and labels. 24 | Input: 25 | data: B,N,... numpy array 26 | label: B,... numpy array 27 | Return: 28 | shuffled data, label and shuffle indices 29 | """ 30 | idx = np.arange(len(labels)) 31 | np.random.shuffle(idx) 32 | return data[idx, ...], labels[idx], idx 33 | 34 | 35 | def rotate_point_cloud(batch_data): 36 | """ Randomly rotate the point clouds to augument the dataset 37 | rotation is per shape based along up direction 38 | Input: 39 | BxNx3 array, original batch of point clouds 40 | Return: 41 | BxNx3 array, rotated batch of point clouds 42 | """ 43 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 44 | for k in range(batch_data.shape[0]): 45 | rotation_angle = np.random.uniform() * 2 * np.pi 46 | cosval = np.cos(rotation_angle) 47 | sinval = np.sin(rotation_angle) 48 | rotation_matrix = np.array( 49 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 50 | ) 51 | shape_pc = batch_data[k, ...] 52 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 53 | return rotated_data 54 | 55 | 56 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 57 | """ Rotate the point cloud along up direction with certain angle. 58 | Input: 59 | BxNx3 array, original batch of point clouds 60 | Return: 61 | BxNx3 array, rotated batch of point clouds 62 | """ 63 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 64 | for k in range(batch_data.shape[0]): 65 | # rotation_angle = np.random.uniform() * 2 * np.pi 66 | cosval = np.cos(rotation_angle) 67 | sinval = np.sin(rotation_angle) 68 | rotation_matrix = np.array( 69 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 70 | ) 71 | shape_pc = batch_data[k, ...] 72 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 73 | return rotated_data 74 | 75 | 76 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05): 77 | """ Randomly jitter points. jittering is per point. 78 | Input: 79 | BxNx3 array, original batch of point clouds 80 | Return: 81 | BxNx3 array, jittered batch of point clouds 82 | """ 83 | B, N, C = batch_data.shape 84 | assert clip > 0 85 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1 * clip, clip) 86 | jittered_data += batch_data 87 | return jittered_data 88 | 89 | 90 | def noisey_point_cloud(batch_data, ratio=0.1): 91 | """ Randomly replace points with noise. 92 | Input: 93 | BxNx3 array, original batch of point clouds 94 | Return: 95 | BxNx3 array, jittered batch of point clouds 96 | """ 97 | assert ratio < 1 and ratio >= 0 98 | B, N, C = batch_data.shape 99 | noise = np.random.rand(B, N, C) * 2 - 1 100 | 101 | rand_idx = np.random.permutation(list(range(0, N))) 102 | rand_idx = rand_idx[: int(N * ratio)] 103 | 104 | noisey_data = batch_data 105 | noisey_data[:, rand_idx, :] = noise[:, rand_idx, :] 106 | 107 | # noiseThresh = np.random.rand(B, N) 108 | # noiseThresh=(noiseThresh > ratio).choose(1,0) 109 | # for i in range(C): 110 | # noise[:, :, i] = noise[:, :, i] * noiseThresh 111 | # 112 | # noisey_data =(noise > 0).choose(batch_data,noise) 113 | return noisey_data 114 | 115 | 116 | def getDataFiles(list_filename): 117 | return [line.rstrip() for line in open(list_filename)] 118 | 119 | 120 | def load_h5(h5_filename): 121 | f = h5py.File(h5_filename) 122 | data = f["data"][:] 123 | label = f["label"][:] 124 | return (data, label) 125 | 126 | 127 | def loadDataFile(filename): 128 | return load_h5(filename) 129 | 130 | 131 | def load_h5_data_label_seg(h5_filename): 132 | f = h5py.File(h5_filename) 133 | data = f["data"][:] 134 | label = f["label"][:] 135 | seg = f["pid"][:] 136 | return (data, label, seg) 137 | 138 | 139 | def loadDataFile_with_seg(filename): 140 | return load_h5_data_label_seg(filename) 141 | -------------------------------------------------------------------------------- /classification/models/transform_nets.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import sys 4 | import os 5 | 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | sys.path.append(BASE_DIR) 8 | sys.path.append(os.path.join(BASE_DIR, "../utils")) 9 | import tf_util 10 | 11 | 12 | def input_transform_net(point_cloud, is_training, bn_decay=None, K=3): 13 | """ Input (XYZ) Transform Net, input is BxNx3 gray image 14 | Return: 15 | Transformation matrix of size 3xK """ 16 | batch_size = point_cloud.get_shape()[0].value 17 | num_point = point_cloud.get_shape()[1].value 18 | 19 | input_image = tf.expand_dims(point_cloud, -1) 20 | net = tf_util.conv2d( 21 | input_image, 22 | 64, 23 | [1, 3], 24 | padding="VALID", 25 | stride=[1, 1], 26 | bn=True, 27 | is_training=is_training, 28 | scope="tconv1", 29 | bn_decay=bn_decay, 30 | ) 31 | net = tf_util.conv2d( 32 | net, 33 | 128, 34 | [1, 1], 35 | padding="VALID", 36 | stride=[1, 1], 37 | bn=True, 38 | is_training=is_training, 39 | scope="tconv2", 40 | bn_decay=bn_decay, 41 | ) 42 | net = tf_util.conv2d( 43 | net, 44 | 1024, 45 | [1, 1], 46 | padding="VALID", 47 | stride=[1, 1], 48 | bn=True, 49 | is_training=is_training, 50 | scope="tconv3", 51 | bn_decay=bn_decay, 52 | ) 53 | net = tf_util.max_pool2d(net, [num_point, 1], padding="VALID", scope="tmaxpool") 54 | 55 | net = tf.reshape(net, [batch_size, -1]) 56 | net = tf_util.fully_connected( 57 | net, 512, bn=True, is_training=is_training, scope="tfc1", bn_decay=bn_decay 58 | ) 59 | net = tf_util.fully_connected( 60 | net, 256, bn=True, is_training=is_training, scope="tfc2", bn_decay=bn_decay 61 | ) 62 | 63 | with tf.variable_scope("transform_XYZ") as sc: 64 | assert K == 3 65 | weights = tf.get_variable( 66 | "weights", 67 | [256, 3 * K], 68 | initializer=tf.constant_initializer(0.0), 69 | dtype=tf.float32, 70 | ) 71 | biases = tf.get_variable( 72 | "biases", 73 | [3 * K], 74 | initializer=tf.constant_initializer(0.0), 75 | dtype=tf.float32, 76 | ) 77 | biases += tf.constant([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=tf.float32) 78 | transform = tf.matmul(net, weights) 79 | transform = tf.nn.bias_add(transform, biases) 80 | 81 | transform = tf.reshape(transform, [batch_size, 3, K]) 82 | return transform 83 | 84 | 85 | def feature_transform_net(inputs, is_training, bn_decay=None, K=64): 86 | """ Feature Transform Net, input is BxNx1xK 87 | Return: 88 | Transformation matrix of size KxK """ 89 | batch_size = inputs.get_shape()[0].value 90 | num_point = inputs.get_shape()[1].value 91 | 92 | net = tf_util.conv2d( 93 | inputs, 94 | 64, 95 | [1, 1], 96 | padding="VALID", 97 | stride=[1, 1], 98 | bn=True, 99 | is_training=is_training, 100 | scope="tconv1", 101 | bn_decay=bn_decay, 102 | ) 103 | net = tf_util.conv2d( 104 | net, 105 | 128, 106 | [1, 1], 107 | padding="VALID", 108 | stride=[1, 1], 109 | bn=True, 110 | is_training=is_training, 111 | scope="tconv2", 112 | bn_decay=bn_decay, 113 | ) 114 | net = tf_util.conv2d( 115 | net, 116 | 1024, 117 | [1, 1], 118 | padding="VALID", 119 | stride=[1, 1], 120 | bn=True, 121 | is_training=is_training, 122 | scope="tconv3", 123 | bn_decay=bn_decay, 124 | ) 125 | net = tf_util.max_pool2d(net, [num_point, 1], padding="VALID", scope="tmaxpool") 126 | 127 | net = tf.reshape(net, [batch_size, -1]) 128 | net = tf_util.fully_connected( 129 | net, 512, bn=True, is_training=is_training, scope="tfc1", bn_decay=bn_decay 130 | ) 131 | net = tf_util.fully_connected( 132 | net, 256, bn=True, is_training=is_training, scope="tfc2", bn_decay=bn_decay 133 | ) 134 | 135 | with tf.variable_scope("transform_feat") as sc: 136 | weights = tf.get_variable( 137 | "weights", 138 | [256, K * K], 139 | initializer=tf.constant_initializer(0.0), 140 | dtype=tf.float32, 141 | ) 142 | biases = tf.get_variable( 143 | "biases", 144 | [K * K], 145 | initializer=tf.constant_initializer(0.0), 146 | dtype=tf.float32, 147 | ) 148 | biases += tf.constant(np.eye(K).flatten(), dtype=tf.float32) 149 | transform = tf.matmul(net, weights) 150 | transform = tf.nn.bias_add(transform, biases) 151 | 152 | transform = tf.reshape(transform, [batch_size, K, K]) 153 | return transform 154 | -------------------------------------------------------------------------------- /classification/README.md: -------------------------------------------------------------------------------- 1 | ## Installation 2 | 3 | The code has been tested with Python 3.6.9, TensorFlow 1.13.2, CUDA 10.0 and cuDNN 7.6.2 on Ubuntu 16.04. 4 | 5 | Install TensorFlow. You can also use a TensorFlow Docker image. A Docker image that meets the TensorFlow, CUDA and cuDNN version that we used is **tensorflow/tensorflow:1.13.2-gpu-py3**. 6 | 7 | Install h5py for Python: 8 | ```bash 9 | sudo apt-get install libhdf5-dev 10 | sudo pip install h5py 11 | ``` 12 | 13 | In order to download the dataset for training and evaluation, wget package is required. To install wget: 14 | ```bash 15 | sudo apt-get update 16 | sudo apt-get install wget 17 | ``` 18 | 19 | ### Compilation of TensorFlow ops 20 | 21 | Compile TensorFlow ops: nearest neighbor grouping, implemented by [Qi et al.](https://github.com/charlesq34/pointnet2); structural losses, implemented by [Fan et al.](https://github.com/fanhqme/PointSetGeneration) The ops are located under `classification` at `grouping` and `structural losses` folders, respectively. If needed, use a text editor and modify the corresponding `sh` file of each op to point to your `nvcc` path. Then, use: 22 | ```bash 23 | cd classification/ 24 | sh compile_ops.sh 25 | ``` 26 | 27 | An `o` and `so` files should be created in the corresponding folder of each op. 28 | 29 | ## Usage 30 | For a quick start please use: 31 | ```bash 32 | sh runner_samplenet.sh 33 | ``` 34 | or: 35 | ```bash 36 | sh runner_samplenet_progressive.sh 37 | ``` 38 | 39 | These scripts train a classifier model with complete point clouds, use it to train a sampler (SampleNet or SampleNetProgressive), and then evaluate sampler by running its sampled points through the classifier model. In the following sections, we explain how to run each part of this pipeline separately. 40 | 41 | ### Data Set 42 | 43 | Point clouds of ModelNet40 models in `h5` files (provided by Qi et al.) will be automatically downloaded (416MB) on the first training of a classifier model. Each point cloud contains 2048 points uniformly sampled from a shape surface. Each cloud is zero-mean and normalized into an unit sphere. There are also `json` files specifying the IDs of shapes in the `h5` files. The data will be downloaded to the folder `classification/data/modelnet40_ply_hdf5_2048`. 44 | 45 | ### Classifier 46 | 47 | To train a PointNet model to classify point clouds, use: 48 | ```bash 49 | python train_classifier.py --model pointnet_cls --log_dir log/PointNet1024 50 | ``` 51 | 52 | To train the vanilla version of PointNet, use: 53 | ```bash 54 | python train_classifier.py --model pointnet_cls_basic --log_dir log/PointNetVanilla1024 55 | ``` 56 | 57 | ### SampleNet 58 | To train SampleNet (with sample size 32 in this example), using an existing classifier (PointNet in this example) as the task network (provided in `classifier_model_path` argument), use: 59 | ```bash 60 | python train_samplenet.py --classifier_model pointnet_cls --classifier_model_path log/PointNet1024/model.ckpt --num_out_points 32 --log_dir log/SampleNet32 61 | ``` 62 | 63 | To evaluate classification with SampleNet's sampled points (with sample size 32 in this example), use: 64 | ```bash 65 | python evaluate_samplenet.py --sampler_model_path log/SampleNet32/model.ckpt --num_out_points 32 --dump_dir log/SampleNet32/eval 66 | ``` 67 | 68 | This evaluation script computes classification accuracy results and saves them to the `dump_dir`. 69 | 70 | ### SampleNetProgressive 71 | To train SampleNetProgressive, using an existing classifier (PointNet vanilla in this example) as the task network (provided in `classifier_model_path` argument), use: 72 | ```bash 73 | python train_samplenet_progressive.py --classifier_model pointnet_cls_basic --classifier_model_path log/PointNetVanilla1024/model.ckpt --log_dir log/SampleNetProgressive 74 | ``` 75 | 76 | Evaluation of SampleNetProgressive is done in two steps. First, infer SampleNetProgressive and save the ordered point clouds to `h5` files: 77 | ```bash 78 | python infer_samplenet_progressive.py --sampler_model_path log/SampleNetProgressive/model.ckpt --dump_dir log/SampleNetProgressive 79 | ``` 80 | 81 | Then, evaluate the classifier using SampleNetProgressive's sampled points: 82 | ```bash 83 | python evaluate_from_files.py --classifier_model pointnet_cls_basic --classifier_model_path log/PointNetVanilla1024/model.ckpt --data_path log/SampleNetProgressive/sampled --dump_dir log/SampleNetProgressive/eval/sampled 84 | ``` 85 | 86 | This evaluation script computes classification accuracy results for different sample sizes and saves them to the `dump_dir`. 87 | 88 | ## Acknowledgment 89 | Our code builds upon the code provided by Qi et al. and Dovrat et al. We thank the authors for sharing their code. 90 | -------------------------------------------------------------------------------- /reconstruction/autoencoder/train_ae.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on September 5th, 2018 3 | 4 | @author: itailang 5 | """ 6 | from __future__ import print_function 7 | 8 | # import system modules 9 | from builtins import str 10 | from builtins import range 11 | import os.path as osp 12 | import sys 13 | import argparse 14 | 15 | # add paths 16 | parent_dir = osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))) 17 | if parent_dir not in sys.path: 18 | sys.path.append(parent_dir) 19 | 20 | # import modules 21 | from reconstruction.src.ae_templates import ( 22 | mlp_architecture_ala_iclr_18, 23 | default_train_params, 24 | ) 25 | from reconstruction.src.autoencoder import Configuration as Conf 26 | from reconstruction.src.pointnet_ae import PointNetAE 27 | 28 | from reconstruction.src.in_out import ( 29 | snc_category_to_synth_id, 30 | create_dir, 31 | PointCloudDataSet, 32 | load_and_split_all_point_clouds_under_folder, 33 | ) 34 | 35 | from reconstruction.src.tf_utils import reset_tf_graph 36 | 37 | # command line arguments 38 | # fmt: off 39 | parser = argparse.ArgumentParser() 40 | parser.add_argument('--use_fps', type=int, default=0, help='1: FPS sampling before autoencoder, 0: do not sample [default: 0]') 41 | parser.add_argument('--n_sample_points', type=int, default=2048, help='Number of sample points of the input [default: 2048]') 42 | parser.add_argument('--object_class', type=str, default='multi', help='Single class name (for example: chair) or multi [default: multi]') 43 | parser.add_argument('--train_folder', type=str, default='log/autoencoder', help='Folder for saving data form the training [default: log/autoencoder]') 44 | parser.add_argument('--training_epochs', type=int, default=500, help='Number of training epochs [default: 500]') 45 | flags = parser.parse_args() 46 | # fmt: on 47 | 48 | print(("Train flags:", flags)) 49 | 50 | # define basic parameters 51 | project_dir = osp.dirname(osp.dirname(osp.abspath(__file__))) 52 | top_in_dir = osp.join( 53 | project_dir, "data", "shape_net_core_uniform_samples_2048" 54 | ) # top-dir of where point-clouds are stored. 55 | top_out_dir = osp.join(project_dir) # use to save Neural-Net check-points etc. 56 | 57 | if flags.object_class == "multi": 58 | class_name = ["chair", "table", "car", "airplane"] 59 | else: 60 | class_name = [str(flags.object_class)] 61 | 62 | experiment_name = "autoencoder" 63 | n_pc_points = 2048 # Number of points per model 64 | bneck_size = 128 # Bottleneck-AE size 65 | ae_loss = "chamfer" # Loss to optimize 66 | 67 | # load Point-Clouds 68 | syn_id = snc_category_to_synth_id()[class_name[0]] 69 | class_dir = osp.join(top_in_dir, syn_id) 70 | pc_data_train, pc_data_val, _ = load_and_split_all_point_clouds_under_folder( 71 | class_dir, n_threads=8, file_ending=".ply", verbose=True 72 | ) 73 | 74 | for i in range(1, len(class_name)): 75 | syn_id = snc_category_to_synth_id()[class_name[i]] 76 | class_dir = osp.join(top_in_dir, syn_id) 77 | ( 78 | pc_data_train_curr, 79 | pc_data_val_curr, 80 | _, 81 | ) = load_and_split_all_point_clouds_under_folder( 82 | class_dir, n_threads=8, file_ending=".ply", verbose=True 83 | ) 84 | pc_data_train.merge(pc_data_train_curr) 85 | pc_data_val.merge(pc_data_val_curr) 86 | 87 | if flags.object_class == "multi": 88 | pc_data_train.shuffle_data(seed=55) 89 | pc_data_val.shuffle_data(seed=55) 90 | 91 | # load default training parameters 92 | train_params = default_train_params() 93 | 94 | encoder, decoder, enc_args, dec_args = mlp_architecture_ala_iclr_18( 95 | n_pc_points, bneck_size 96 | ) 97 | 98 | train_dir = create_dir(osp.join(top_out_dir, flags.train_folder)) 99 | 100 | conf = Conf( 101 | n_input=[n_pc_points, 3], 102 | loss=ae_loss, 103 | training_epochs=train_params["training_epochs"], 104 | batch_size=train_params["batch_size"], 105 | denoising=train_params["denoising"], 106 | learning_rate=train_params["learning_rate"], 107 | train_dir=train_dir, 108 | loss_display_step=train_params["loss_display_step"], 109 | saver_step=train_params["saver_step"], 110 | z_rotate=train_params["z_rotate"], 111 | encoder=encoder, 112 | decoder=decoder, 113 | encoder_args=enc_args, 114 | decoder_args=dec_args, 115 | ) 116 | conf.experiment_name = experiment_name 117 | conf.held_out_step = 5 # how often to evaluate/print(out loss on) 118 | # held_out data (if they are provided in ae.train()). 119 | conf.class_name = class_name 120 | conf.use_fps = flags.use_fps 121 | conf.n_sample_points = flags.n_sample_points 122 | conf.n_samp_out = [2048, 3] 123 | conf.training_epochs = flags.training_epochs 124 | conf.save(osp.join(train_dir, "configuration")) 125 | 126 | # build AE Model 127 | reset_tf_graph() 128 | ae = PointNetAE(conf.experiment_name, conf) 129 | 130 | # train the AE (save output to train_stats.txt) 131 | buf_size = 1 # Make 'training_stats' file to flush each output line regarding training. 132 | fout = open(osp.join(conf.train_dir, "train_stats.txt"), "a", buf_size) 133 | train_stats = ae.train(pc_data_train, conf, log_file=fout, held_out_data=pc_data_val) 134 | fout.close() 135 | -------------------------------------------------------------------------------- /classification/grouping/test/query_ball_point_grid.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int batch_index = blockIdx.x; 21 | xyz1 += n*3*batch_index; 22 | xyz2 += m*3*batch_index; 23 | idx += m*nsample*batch_index; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | 28 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 123 | cudaDeviceSynchronize(); 124 | printf("query_ball_point gpu time %f\n",get_time()-t0); 125 | 126 | t0=get_time(); 127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 128 | cudaDeviceSynchronize(); 129 | printf("grou_point gpu time %f\n",get_time()-t0); 130 | 131 | t0=get_time(); 132 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 133 | cudaDeviceSynchronize(); 134 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 135 | 136 | cudaFree(xyz1); 137 | cudaFree(xyz2); 138 | cudaFree(points); 139 | cudaFree(idx); 140 | cudaFree(out); 141 | cudaFree(grad_out); 142 | cudaFree(grad_points); 143 | return 0; 144 | } 145 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/test/query_ball_point_grid.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int batch_index = blockIdx.x; 21 | xyz1 += n*3*batch_index; 22 | xyz2 += m*3*batch_index; 23 | idx += m*nsample*batch_index; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | 28 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 123 | cudaDeviceSynchronize(); 124 | printf("query_ball_point gpu time %f\n",get_time()-t0); 125 | 126 | t0=get_time(); 127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 128 | cudaDeviceSynchronize(); 129 | printf("grou_point gpu time %f\n",get_time()-t0); 130 | 131 | t0=get_time(); 132 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 133 | cudaDeviceSynchronize(); 134 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 135 | 136 | cudaFree(xyz1); 137 | cudaFree(xyz2); 138 | cudaFree(points); 139 | cudaFree(idx); 140 | cudaFree(out); 141 | cudaFree(grad_out); 142 | cudaFree(grad_points); 143 | return 0; 144 | } 145 | -------------------------------------------------------------------------------- /classification/structural_losses/tf_nndistance_g.cu: -------------------------------------------------------------------------------- 1 | #if GOOGLE_CUDA 2 | #define EIGEN_USE_GPU 3 | #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 4 | 5 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 6 | const int batch=512; 7 | __shared__ float buf[batch*3]; 8 | for (int i=blockIdx.x;ibest){ 120 | result[(i*n+j)]=best; 121 | result_i[(i*n+j)]=best_i; 122 | } 123 | } 124 | __syncthreads(); 125 | } 126 | } 127 | } 128 | void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){ 129 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 130 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 131 | } 132 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 133 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 156 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 157 | } 158 | 159 | #endif 160 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/tf_nndistance_g.cu: -------------------------------------------------------------------------------- 1 | #if GOOGLE_CUDA 2 | #define EIGEN_USE_GPU 3 | #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 4 | 5 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 6 | const int batch=512; 7 | __shared__ float buf[batch*3]; 8 | for (int i=blockIdx.x;ibest){ 120 | result[(i*n+j)]=best; 121 | result_i[(i*n+j)]=best_i; 122 | } 123 | } 124 | __syncthreads(); 125 | } 126 | } 127 | } 128 | void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){ 129 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 130 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 131 | } 132 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 133 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 156 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 157 | } 158 | 159 | #endif 160 | -------------------------------------------------------------------------------- /classification/models/pointnet_cls_basic.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import tensorflow as tf 3 | import numpy as np 4 | import math 5 | import sys 6 | import os 7 | 8 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(os.path.join(BASE_DIR, "../utils")) 11 | import tf_util 12 | 13 | 14 | def placeholder_inputs(batch_size, num_point): 15 | pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 16 | labels_pl = tf.placeholder(tf.int32, shape=(batch_size)) 17 | return pointclouds_pl, labels_pl 18 | 19 | 20 | # def get_latent_vector(point_cloud, is_training, bn_decay=None): 21 | # num_point = point_cloud.get_shape()[1].value 22 | # input_image = tf.expand_dims(point_cloud, -1) 23 | # 24 | # # Point functions (MLP implemented as conv2d) 25 | # net = tf_util.conv2d(input_image, 64, [1, 3], 26 | # padding='VALID', stride=[1, 1], 27 | # bn=True, is_training=is_training, 28 | # scope='conv1', bn_decay=bn_decay) 29 | # net = tf_util.conv2d(net, 64, [1, 1], 30 | # padding='VALID', stride=[1, 1], 31 | # bn=True, is_training=is_training, 32 | # scope='conv2', bn_decay=bn_decay) 33 | # net = tf_util.conv2d(net, 64, [1, 1], 34 | # padding='VALID', stride=[1, 1], 35 | # bn=True, is_training=is_training, 36 | # scope='conv3', bn_decay=bn_decay) 37 | # net = tf_util.conv2d(net, 128, [1, 1], 38 | # padding='VALID', stride=[1, 1], 39 | # bn=True, is_training=is_training, 40 | # scope='conv4', bn_decay=bn_decay) 41 | # net = tf_util.conv2d(net, 1024, [1, 1], 42 | # padding='VALID', stride=[1, 1], 43 | # bn=True, is_training=is_training, 44 | # scope='conv5', bn_decay=bn_decay) 45 | # 46 | # # Symmetric function: max pooling 47 | # net = tf_util.max_pool2d(net, [num_point, 1], 48 | # padding='VALID', scope='maxpool') 49 | # 50 | # latent_vector = tf.squeeze(net, [1, 2]) 51 | # 52 | # return latent_vector 53 | 54 | 55 | def get_model(point_cloud, is_training, bn_decay=None): 56 | """ Classification PointNet, input is BxNx3, output Bx40 """ 57 | batch_size = point_cloud.get_shape()[0].value 58 | num_point = point_cloud.get_shape()[1].value 59 | end_points = {} 60 | input_image = tf.expand_dims(point_cloud, -1) 61 | 62 | # Point functions (MLP implemented as conv2d) 63 | net = tf_util.conv2d( 64 | input_image, 65 | 64, 66 | [1, 3], 67 | padding="VALID", 68 | stride=[1, 1], 69 | bn=True, 70 | is_training=is_training, 71 | scope="conv1", 72 | bn_decay=bn_decay, 73 | ) 74 | net = tf_util.conv2d( 75 | net, 76 | 64, 77 | [1, 1], 78 | padding="VALID", 79 | stride=[1, 1], 80 | bn=True, 81 | is_training=is_training, 82 | scope="conv2", 83 | bn_decay=bn_decay, 84 | ) 85 | net = tf_util.conv2d( 86 | net, 87 | 64, 88 | [1, 1], 89 | padding="VALID", 90 | stride=[1, 1], 91 | bn=True, 92 | is_training=is_training, 93 | scope="conv3", 94 | bn_decay=bn_decay, 95 | ) 96 | net = tf_util.conv2d( 97 | net, 98 | 128, 99 | [1, 1], 100 | padding="VALID", 101 | stride=[1, 1], 102 | bn=True, 103 | is_training=is_training, 104 | scope="conv4", 105 | bn_decay=bn_decay, 106 | ) 107 | net = tf_util.conv2d( 108 | net, 109 | 1024, 110 | [1, 1], 111 | padding="VALID", 112 | stride=[1, 1], 113 | bn=True, 114 | is_training=is_training, 115 | scope="conv5", 116 | bn_decay=bn_decay, 117 | ) 118 | 119 | # Symmetric function: max pooling 120 | end_points["critical_set_idx"] = tf.arg_max(net, 1) 121 | net = tf_util.max_pool2d(net, [num_point, 1], padding="VALID", scope="maxpool") 122 | end_points["GFV"] = net 123 | # latent_vector = tf.squeeze(net, [1, 2]) 124 | 125 | # MLP on global point cloud vector 126 | net = tf.reshape(net, [batch_size, -1]) 127 | net = tf_util.fully_connected( 128 | net, 512, bn=True, is_training=is_training, scope="fc1", bn_decay=bn_decay 129 | ) 130 | net = tf_util.fully_connected( 131 | net, 256, bn=True, is_training=is_training, scope="fc2", bn_decay=bn_decay 132 | ) 133 | net = tf_util.dropout(net, keep_prob=0.7, is_training=is_training, scope="dp1") 134 | net = tf_util.fully_connected(net, 40, activation_fn=None, scope="fc3") 135 | 136 | return net, end_points 137 | 138 | 139 | def get_loss(pred, label, end_points): 140 | """ pred: B*NUM_CLASSES, 141 | label: B, """ 142 | loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=label) 143 | classify_loss = tf.reduce_mean(loss) 144 | tf.summary.scalar("classify loss", classify_loss) 145 | return classify_loss 146 | 147 | 148 | if __name__ == "__main__": 149 | with tf.Graph().as_default(): 150 | inputs = tf.zeros((32, 1024, 3)) 151 | outputs = get_model(inputs, tf.constant(True)) 152 | print(outputs) 153 | -------------------------------------------------------------------------------- /reconstruction/external/sampling/tf_sampling.py: -------------------------------------------------------------------------------- 1 | """ Furthest point sampling 2 | Original author: Haoqiang Fan 3 | Modified by Charles R. Qi 4 | All Rights Reserved. 2017. 5 | """ 6 | from __future__ import print_function 7 | 8 | try: 9 | from future import standard_library 10 | standard_library.install_aliases() 11 | except ModuleNotFoundError: 12 | print("No module named 'future'. For python2 compatibility, please install 'future' module.") 13 | 14 | import tensorflow as tf 15 | from tensorflow.python.framework import ops 16 | import sys 17 | import os 18 | 19 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 20 | sys.path.append(BASE_DIR) 21 | sampling_module = tf.load_op_library(os.path.join(BASE_DIR, "tf_sampling_so.so")) 22 | 23 | 24 | def prob_sample(inp, inpr): 25 | """ 26 | input: 27 | batch_size * ncategory float32 28 | batch_size * npoints float32 29 | returns: 30 | batch_size * npoints int32 31 | """ 32 | return sampling_module.prob_sample(inp, inpr) 33 | 34 | 35 | ops.NoGradient("ProbSample") 36 | # TF1.0 API requires set shape in C++ 37 | # @tf.RegisterShape('ProbSample') 38 | # def _prob_sample_shape(op): 39 | # shape1=op.inputs[0].get_shape().with_rank(2) 40 | # shape2=op.inputs[1].get_shape().with_rank(2) 41 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 42 | def gather_point(inp, idx): 43 | """ 44 | input: 45 | batch_size * ndataset * 3 float32 46 | batch_size * npoints int32 47 | returns: 48 | batch_size * npoints * 3 float32 49 | """ 50 | return sampling_module.gather_point(inp, idx) 51 | 52 | 53 | # @tf.RegisterShape('GatherPoint') 54 | # def _gather_point_shape(op): 55 | # shape1=op.inputs[0].get_shape().with_rank(3) 56 | # shape2=op.inputs[1].get_shape().with_rank(2) 57 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 58 | @tf.RegisterGradient("GatherPoint") 59 | def _gather_point_grad(op, out_g): 60 | inp = op.inputs[0] 61 | idx = op.inputs[1] 62 | return [sampling_module.gather_point_grad(inp, idx, out_g), None] 63 | 64 | 65 | def farthest_point_sample(npoint, inp): 66 | """ 67 | input: 68 | int32 69 | batch_size * ndataset * 3 float32 70 | returns: 71 | batch_size * npoint int32 72 | """ 73 | return sampling_module.farthest_point_sample(inp, npoint) 74 | 75 | 76 | ops.NoGradient("FarthestPointSample") 77 | 78 | 79 | def non_sampled(ndataset, idx): 80 | """ 81 | input: 82 | int32 83 | batch_size * npoint int32 84 | returns: 85 | batch_size * (ndataset-npoint) int32 86 | """ 87 | 88 | idx_shape = idx.get_shape().as_list() 89 | batch_size = idx_shape[ 90 | 0 91 | ] # use static shape instead of dynamic shape (tf.shape(idx)[0]) 92 | npoint = idx_shape[ 93 | 1 94 | ] # # use static shape instead of dynamic shape (tf.shape(idx)[1]) 95 | 96 | x = tf.range(npoint) # [0, 1, ..., npoint-1] 97 | y = tf.range(batch_size) # [0, 1, ..., batch_size-1] 98 | _, iy = tf.meshgrid(x, y) 99 | 100 | sampled_indices = tf.stack([iy, idx], 2) 101 | 102 | ones = tf.ones([batch_size, npoint], tf.int32) 103 | sampled_indicator = tf.scatter_nd(sampled_indices, ones, [batch_size, ndataset]) 104 | 105 | zero = tf.constant(0, dtype=tf.int32) 106 | non_sampled_ij = tf.where(tf.equal(sampled_indicator, zero)) 107 | 108 | non_idx = tf.reshape(non_sampled_ij[:, 1], [batch_size, ndataset - npoint]) 109 | non_idx = tf.cast(non_idx, tf.int32) 110 | return non_idx 111 | 112 | 113 | if __name__ == "__main__": 114 | # test non_sampled function 115 | ndataset = 6 116 | idx = tf.constant([[0, 2, 1], [5, 3, 0]], dtype=tf.int32) 117 | idx_non = non_sampled(ndataset, idx) 118 | 119 | with tf.Session("") as sess: 120 | print("\n" * 2) 121 | print("non-sampled indices:") 122 | print(sess.run(idx_non)) 123 | print("\n" * 2) 124 | 125 | # test farthest_point_sample 126 | import numpy as np 127 | 128 | np.random.seed(100) 129 | triangles = np.random.rand(1, 5, 3, 3).astype("float32") 130 | with tf.device("/gpu:1"): 131 | inp = tf.constant(triangles) 132 | tria = inp[:, :, 0, :] 133 | trib = inp[:, :, 1, :] 134 | tric = inp[:, :, 2, :] 135 | areas = tf.sqrt( 136 | tf.reduce_sum(tf.cross(trib - tria, tric - tria) ** 2, 2) + 1e-9 137 | ) 138 | randomnumbers = tf.random_uniform((1, 8192)) 139 | triids = prob_sample(areas, randomnumbers) 140 | tria_sample = gather_point(tria, triids) 141 | trib_sample = gather_point(trib, triids) 142 | tric_sample = gather_point(tric, triids) 143 | us = tf.random_uniform((1, 8192)) 144 | vs = tf.random_uniform((1, 8192)) 145 | uplusv = 1 - tf.abs(us + vs - 1) 146 | uminusv = us - vs 147 | us = (uplusv + uminusv) * 0.5 148 | vs = (uplusv - uminusv) * 0.5 149 | pt_sample = ( 150 | tria_sample 151 | + (trib_sample - tria_sample) * tf.expand_dims(us, -1) 152 | + (tric_sample - tria_sample) * tf.expand_dims(vs, -1) 153 | ) 154 | print(("pt_sample: ", pt_sample)) 155 | reduced_sample = gather_point(pt_sample, farthest_point_sample(1024, pt_sample)) 156 | print(reduced_sample) 157 | with tf.Session("") as sess: 158 | ret = sess.run(reduced_sample) 159 | print((ret.shape, ret.dtype)) 160 | import pickle as pickle 161 | 162 | pickle.dump(ret, open("1.pkl", "wb"), -1) 163 | -------------------------------------------------------------------------------- /registration/src/pctransforms.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | import torch 9 | import numpy as np 10 | 11 | 12 | def angle_axis(angle, axis): 13 | # type: (float, np.ndarray) -> float 14 | r"""Returns a 4x4 rotation matrix that performs a rotation around axis by angle 15 | 16 | Parameters 17 | ---------- 18 | angle : float 19 | Angle to rotate by 20 | axis: np.ndarray 21 | Axis to rotate about 22 | 23 | Returns 24 | ------- 25 | torch.Tensor 26 | 3x3 rotation matrix 27 | """ 28 | u = axis / np.linalg.norm(axis) 29 | cosval, sinval = np.cos(angle), np.sin(angle) 30 | 31 | # yapf: disable 32 | cross_prod_mat = np.array([[0.0, -u[2], u[1]], 33 | [u[2], 0.0, -u[0]], 34 | [-u[1], u[0], 0.0]]) 35 | 36 | R = torch.from_numpy( 37 | cosval * np.eye(3) 38 | + sinval * cross_prod_mat 39 | + (1.0 - cosval) * np.outer(u, u) 40 | ) 41 | # yapf: enable 42 | return R.float() 43 | 44 | 45 | class PointcloudScale(object): 46 | def __init__(self, lo=0.8, hi=1.25): 47 | self.lo, self.hi = lo, hi 48 | 49 | def __call__(self, points): 50 | scaler = np.random.uniform(self.lo, self.hi) 51 | points[:, 0:3] *= scaler 52 | return points 53 | 54 | 55 | class PointcloudRotate(object): 56 | def __init__(self, axis=np.array([0.0, 1.0, 0.0])): 57 | self.axis = axis 58 | 59 | def __call__(self, points): 60 | rotation_angle = np.random.uniform() * 2 * np.pi 61 | rotation_matrix = angle_axis(rotation_angle, self.axis) 62 | 63 | normals = points.size(1) > 3 64 | if not normals: 65 | return torch.matmul(points, rotation_matrix.t()) 66 | else: 67 | pc_xyz = points[:, 0:3] 68 | pc_normals = points[:, 3:] 69 | points[:, 0:3] = torch.matmul(pc_xyz, rotation_matrix.t()) 70 | points[:, 3:] = torch.matmul(pc_normals, rotation_matrix.t()) 71 | 72 | return points 73 | 74 | 75 | class PointcloudRotatePerturbation(object): 76 | def __init__(self, angle_sigma=0.06, angle_clip=0.18): 77 | self.angle_sigma, self.angle_clip = angle_sigma, angle_clip 78 | 79 | def _get_angles(self): 80 | angles = np.clip( 81 | self.angle_sigma * np.random.randn(3), -self.angle_clip, self.angle_clip 82 | ) 83 | 84 | return angles 85 | 86 | def __call__(self, points): 87 | angles = self._get_angles() 88 | Rx = angle_axis(angles[0], np.array([1.0, 0.0, 0.0])) 89 | Ry = angle_axis(angles[1], np.array([0.0, 1.0, 0.0])) 90 | Rz = angle_axis(angles[2], np.array([0.0, 0.0, 1.0])) 91 | 92 | rotation_matrix = torch.matmul(torch.matmul(Rz, Ry), Rx) 93 | 94 | normals = points.size(1) > 3 95 | if not normals: 96 | return torch.matmul(points, rotation_matrix.t()) 97 | else: 98 | pc_xyz = points[:, 0:3] 99 | pc_normals = points[:, 3:] 100 | points[:, 0:3] = torch.matmul(pc_xyz, rotation_matrix.t()) 101 | points[:, 3:] = torch.matmul(pc_normals, rotation_matrix.t()) 102 | 103 | return points 104 | 105 | 106 | class PointcloudJitter(object): 107 | def __init__(self, std=0.01, clip=0.05): 108 | self.std, self.clip = std, clip 109 | 110 | def __call__(self, points): 111 | jittered_data = ( 112 | points.new(points.size(0), 3) 113 | .normal_(mean=0.0, std=self.std) 114 | .clamp_(-self.clip, self.clip) 115 | ) 116 | points[:, 0:3] += jittered_data 117 | return points 118 | 119 | 120 | class PointcloudTranslate(object): 121 | def __init__(self, translate_range=0.1): 122 | self.translate_range = translate_range 123 | 124 | def __call__(self, points): 125 | translation = np.random.uniform(-self.translate_range, self.translate_range) 126 | points[:, 0:3] += translation 127 | return points 128 | 129 | 130 | class PointcloudToTensor(object): 131 | def __call__(self, points): 132 | return torch.from_numpy(points).float() 133 | 134 | 135 | class PointcloudRandomInputDropout(object): 136 | def __init__(self, max_dropout_ratio=0.875): 137 | assert max_dropout_ratio >= 0 and max_dropout_ratio < 1 138 | self.max_dropout_ratio = max_dropout_ratio 139 | 140 | def __call__(self, points): 141 | pc = points.numpy() 142 | 143 | dropout_ratio = np.random.random() * self.max_dropout_ratio # 0~0.875 144 | drop_idx = np.where(np.random.random((pc.shape[0])) <= dropout_ratio)[0] 145 | if len(drop_idx) > 0: 146 | pc[drop_idx] = pc[0] # set to the first point 147 | 148 | return torch.from_numpy(pc).float() 149 | 150 | 151 | class OnUnitCube(object): 152 | def __init__(self): 153 | pass 154 | 155 | def method1(self, tensor): 156 | m = tensor.mean(dim=0, keepdim=True) # [N, D] -> [1, D] 157 | v = tensor - m 158 | s = torch.max(v.abs()) 159 | v = v / s * 0.5 160 | return v 161 | 162 | def method2(self, tensor): 163 | c = torch.max(tensor, dim=0)[0] - torch.min(tensor, dim=0)[0] # [N, D] -> [D] 164 | s = torch.max(c) # -> scalar 165 | v = tensor / s 166 | return v - v.mean(dim=0, keepdim=True) 167 | 168 | def __call__(self, tensor): 169 | # return self.method1(tensor) 170 | return self.method2(tensor) 171 | -------------------------------------------------------------------------------- /classification/grouping/tf_grouping_g.cu: -------------------------------------------------------------------------------- 1 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 2 | // output: idx (b,m,nsample), pts_cnt (b,m) 3 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx, int *pts_cnt) { 4 | int batch_index = blockIdx.x; 5 | xyz1 += n*3*batch_index; 6 | xyz2 += m*3*batch_index; 7 | idx += m*nsample*batch_index; 8 | pts_cnt += m*batch_index; // counting how many unique points selected in local region 9 | 10 | int index = threadIdx.x; 11 | int stride = blockDim.x; 12 | 13 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx,pts_cnt); 127 | //cudaDeviceSynchronize(); 128 | } 129 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 130 | selection_sort_gpu<<>>(b,n,m,k,dist,outi,out); 131 | //cudaDeviceSynchronize(); 132 | } 133 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){ 134 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 135 | //cudaDeviceSynchronize(); 136 | } 137 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){ 138 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 139 | //group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 140 | //cudaDeviceSynchronize(); 141 | } 142 | -------------------------------------------------------------------------------- /reconstruction/external/grouping/tf_grouping_g.cu: -------------------------------------------------------------------------------- 1 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 2 | // output: idx (b,m,nsample), pts_cnt (b,m) 3 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx, int *pts_cnt) { 4 | int batch_index = blockIdx.x; 5 | xyz1 += n*3*batch_index; 6 | xyz2 += m*3*batch_index; 7 | idx += m*nsample*batch_index; 8 | pts_cnt += m*batch_index; // counting how many unique points selected in local region 9 | 10 | int index = threadIdx.x; 11 | int stride = blockDim.x; 12 | 13 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx,pts_cnt); 127 | //cudaDeviceSynchronize(); 128 | } 129 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 130 | selection_sort_gpu<<>>(b,n,m,k,dist,outi,out); 131 | //cudaDeviceSynchronize(); 132 | } 133 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){ 134 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 135 | //cudaDeviceSynchronize(); 136 | } 137 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){ 138 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 139 | //group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 140 | //cudaDeviceSynchronize(); 141 | } 142 | -------------------------------------------------------------------------------- /reconstruction/README.md: -------------------------------------------------------------------------------- 1 | ## Installation 2 | 3 | The code has been tested with Python 3.6.9, TensorFlow 1.13.2, TFLearn 0.3.2, CUDA 10.0 and cuDNN 7.6.2 on Ubuntu 16.04. 4 | 5 | Install TensorFlow. You can also use a TensorFlow Docker image. A Docker image that meets the TensorFlow, CUDA and cuDNN version that we used is **tensorflow/tensorflow:1.13.2-gpu-py3**. 6 | 7 | Install TFLearn. 8 | 9 | Install matplotlib (for visualization of point clouds). 10 | 11 | In order to download the dataset for training and evaluation, wget package is required. To install wget: 12 | ```bash 13 | sudo apt-get update 14 | sudo apt-get install wget 15 | ``` 16 | 17 | ### Compilation of TensorFlow ops 18 | 19 | Compile TensorFlow ops: nearest neighbor grouping and farthest point sampling, implemented by [Qi et al.](https://github.com/charlesq34/pointnet2); structural losses, implemented by [Fan et al.](https://github.com/fanhqme/PointSetGeneration) The ops are located under `reconstruction/external` at `grouping`, `sampling`, and `structural losses` folders, respectively. If needed, use a text editor and modify the corresponding `sh` file of each op to point to your `nvcc` path. Then, use: 20 | ```bash 21 | cd reconstruction/ 22 | sh compile_ops.sh 23 | ``` 24 | 25 | An `o` and `so` files should be created in the corresponding folder of each op. 26 | 27 | ## Usage 28 | For a quick start please use: 29 | ```bash 30 | sh runner_samplenet.sh 31 | ``` 32 | or: 33 | ```bash 34 | sh runner_samplenet_progressive.sh 35 | ``` 36 | 37 | These scripts train and evaluate an Autoencoder model with complete point clouds, use it to train a sampler (SampleNet or SampleNetProgressive), and then evaluate sampler by running its sampled points through the Autoencoder model. In the following sections, we explain how to run each part of this pipeline separately. 38 | 39 | ### Data Set 40 | 41 | Point clouds of ShapeNetCore models in `ply` files (provided by Achlioptas et al.) will be automatically downloaded (1.4GB) on the first training of an Autoencoder model. Each point cloud contains 2048 points, uniformly sampled from a shape surface. The data will be downloaded to the folder `reconstruction/data/shape_net_core_uniform_samples_2048`. 42 | 43 | Alternatively, you can download the data before training by using: 44 | ```bash 45 | sh download_data.sh 46 | ``` 47 | 48 | ### Autoencoder 49 | 50 | To train an Autoencoder model, use: 51 | ```bash 52 | python autoencoder/train_ae.py --train_folder log/autoencoder 53 | ``` 54 | 55 | To evaluate the Autoencoder model, use: 56 | ```bash 57 | python autoencoder/evaluate_ae.py --train_folder log/autoencoder 58 | ``` 59 | 60 | This evaluation script saves the reconstructed point clouds from complete input point clouds of the test set, and the reconstruction error per point cloud (Chamfer distance between the input and reconstruction). The results are saved to the `train_folder`. 61 | 62 | To evaluate reconstruction with FPS sampled points (with sample size 64 in this example), use: 63 | ```bash 64 | python autoencoder/evaluate_ae.py --train_folder log/autoencoder --use_fps 1 --n_sample_points 64 65 | ``` 66 | 67 | This evaluation script saves the sampled point clouds, sample indices and reconstructed point clouds of the test set, and the reconstruction error per point cloud (Chamfer distance between the input and reconstruction). It also computes the normalized reconstruction error, as explained in the paper. The results are saved to the `train_folder`. 68 | 69 | ### SampleNet 70 | To train SampleNet (with sample size 64 in this example), using an existing Autoencoder model as the task network (provided in `ae_folder` argument), use: 71 | ```bash 72 | python sampler/train_samplenet.py --ae_folder log/autoencoder --n_sample_points 64 --train_folder log/SampleNet64 73 | ``` 74 | 75 | To evaluate reconstruction with SampleNet's sampled points (with sample size 64 in this example), use: 76 | ```bash 77 | python sampler/evaluate_samplenet.py --train_folder log/SampleNet64 78 | ``` 79 | 80 | This script operates similarly to the evaluation script for the Autoencoder with FPS sampled points. 81 | 82 | ### SampleNetProgressive 83 | To train SampleNetProgressive, using an existing Autoencoder model as the task network (provided in `ae_folder` argument), use: 84 | ```bash 85 | python sampler/train_samplenet_progressive.py --ae_folder log/autoencoder --n_sample_points 64 --train_folder log/SampleNetProgressive 86 | ``` 87 | 88 | To evaluate reconstruction with SampleNetProgressive's sampled points (with sample size 64 in this example), use: 89 | ```bash 90 | python sampler/evaluate_samplenet_progressive.py --n_sample_points 64 --train_folder log/SampleNetProgressive 91 | ``` 92 | 93 | This script operates similarly to the evaluation script for the Autoencoder with FPS sampled points. 94 | 95 | ### Visualization 96 | You can visualized point clouds (input, reconstructed, or sampled) by adding the flag `--visualize_results` to the evaluation script of the Autoencoder, SampleNet or SampleNetProgressive. 97 | 98 | ## Acknowledgment 99 | Our code builds upon the code provided by Achlioptas et al. and Dovrat et al. We thank the authors for sharing their code. 100 | -------------------------------------------------------------------------------- /registration/src/chamfer_distance/chamfer_distance.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | __global__ 7 | void ChamferDistanceKernel( 8 | int b, 9 | int n, 10 | const float* xyz, 11 | int m, 12 | const float* xyz2, 13 | float* result, 14 | int* result_i) 15 | { 16 | const int batch=512; 17 | __shared__ float buf[batch*3]; 18 | for (int i=blockIdx.x;ibest){ 130 | result[(i*n+j)]=best; 131 | result_i[(i*n+j)]=best_i; 132 | } 133 | } 134 | __syncthreads(); 135 | } 136 | } 137 | } 138 | 139 | void ChamferDistanceKernelLauncher( 140 | const int b, const int n, 141 | const float* xyz, 142 | const int m, 143 | const float* xyz2, 144 | float* result, 145 | int* result_i, 146 | float* result2, 147 | int* result2_i) 148 | { 149 | ChamferDistanceKernel<<>>(b, n, xyz, m, xyz2, result, result_i); 150 | ChamferDistanceKernel<<>>(b, m, xyz2, n, xyz, result2, result2_i); 151 | 152 | cudaError_t err = cudaGetLastError(); 153 | if (err != cudaSuccess) 154 | printf("error in chamfer distance updateOutput: %s\n", cudaGetErrorString(err)); 155 | } 156 | 157 | 158 | __global__ 159 | void ChamferDistanceGradKernel( 160 | int b, int n, 161 | const float* xyz1, 162 | int m, 163 | const float* xyz2, 164 | const float* grad_dist1, 165 | const int* idx1, 166 | float* grad_xyz1, 167 | float* grad_xyz2) 168 | { 169 | for (int i = blockIdx.x; i>>(b, n, xyz1, m, xyz2, grad_dist1, idx1, grad_xyz1, grad_xyz2); 204 | ChamferDistanceGradKernel<<>>(b, m, xyz2, n, xyz1, grad_dist2, idx2, grad_xyz2, grad_xyz1); 205 | 206 | cudaError_t err = cudaGetLastError(); 207 | if (err != cudaSuccess) 208 | printf("error in chamfer distance get grad: %s\n", cudaGetErrorString(err)); 209 | } 210 | -------------------------------------------------------------------------------- /classification/models/samplenet_model.py: -------------------------------------------------------------------------------- 1 | from builtins import range 2 | import tensorflow as tf 3 | import numpy as np 4 | import math 5 | import sys 6 | import os 7 | 8 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(os.path.join(BASE_DIR, "../utils")) 11 | import tf_util 12 | from structural_losses.tf_nndistance import nn_distance 13 | from structural_losses.tf_approxmatch import approx_match 14 | 15 | 16 | def placeholder_inputs(batch_size, num_point): 17 | pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 18 | labels_pl = tf.placeholder(tf.int32, shape=(batch_size)) 19 | return pointclouds_pl, labels_pl 20 | 21 | 22 | def get_model( 23 | point_cloud, is_training, num_output_points, bottleneck_size, bn_decay=None 24 | ): 25 | """ Classification PointNet, input is BxNx3, output Bx40 """ 26 | batch_size = point_cloud.get_shape()[0].value 27 | num_point = point_cloud.get_shape()[1].value 28 | input_image = tf.expand_dims(point_cloud, -1) 29 | 30 | # Point functions (MLP implemented as conv2d) 31 | net = tf_util.conv2d( 32 | input_image, 33 | 64, 34 | [1, 3], 35 | padding="VALID", 36 | stride=[1, 1], 37 | bn=True, 38 | is_training=is_training, 39 | scope="conv1", 40 | bn_decay=bn_decay, 41 | ) 42 | net = tf_util.conv2d( 43 | net, 44 | 64, 45 | [1, 1], 46 | padding="VALID", 47 | stride=[1, 1], 48 | bn=True, 49 | is_training=is_training, 50 | scope="conv2", 51 | bn_decay=bn_decay, 52 | ) 53 | net = tf_util.conv2d( 54 | net, 55 | 64, 56 | [1, 1], 57 | padding="VALID", 58 | stride=[1, 1], 59 | bn=True, 60 | is_training=is_training, 61 | scope="conv3", 62 | bn_decay=bn_decay, 63 | ) 64 | net = tf_util.conv2d( 65 | net, 66 | 128, 67 | [1, 1], 68 | padding="VALID", 69 | stride=[1, 1], 70 | bn=True, 71 | is_training=is_training, 72 | scope="conv4", 73 | bn_decay=bn_decay, 74 | ) 75 | net = tf_util.conv2d( 76 | net, 77 | bottleneck_size, 78 | [1, 1], 79 | padding="VALID", 80 | stride=[1, 1], 81 | bn=True, 82 | is_training=is_training, 83 | scope="conv5", 84 | bn_decay=bn_decay, 85 | ) 86 | 87 | net = tf_util.max_pool2d(net, [num_point, 1], padding="VALID", scope="maxpool") 88 | 89 | net = tf.reshape(net, [batch_size, -1]) 90 | 91 | net = tf_util.fully_connected( 92 | net, 256, bn=True, is_training=is_training, scope="fc11b", bn_decay=bn_decay 93 | ) 94 | net = tf_util.fully_connected( 95 | net, 256, bn=True, is_training=is_training, scope="fc12b", bn_decay=bn_decay 96 | ) 97 | net = tf_util.fully_connected( 98 | net, 256, bn=True, is_training=is_training, scope="fc13b", bn_decay=bn_decay 99 | ) 100 | net = tf_util.fully_connected( 101 | net, 102 | 3 * num_output_points, 103 | bn=True, 104 | is_training=is_training, 105 | scope="fc14b", 106 | bn_decay=bn_decay, 107 | activation_fn=None, 108 | ) 109 | 110 | out_point_cloud = tf.reshape(net, [batch_size, -1, 3]) 111 | 112 | return out_point_cloud 113 | 114 | 115 | def calc_distances(p0, points): 116 | return ((p0 - points) ** 2).sum(axis=1) 117 | 118 | 119 | def fps_from_given_pc(pts, k, given_pc): 120 | farthest_pts = np.zeros((k, 3)) 121 | t = np.size(given_pc) // 3 122 | farthest_pts[0:t] = given_pc 123 | 124 | distances = calc_distances(farthest_pts[0], pts) 125 | for i in range(1, t): 126 | distances = np.minimum(distances, calc_distances(farthest_pts[i], pts)) 127 | 128 | for i in range(t, k): 129 | farthest_pts[i] = pts[np.argmax(distances)] 130 | distances = np.minimum(distances, calc_distances(farthest_pts[i], pts)) 131 | return farthest_pts 132 | 133 | 134 | def unique(arr): 135 | _, idx = np.unique(arr, return_index=True) 136 | return arr[np.sort(idx)] 137 | 138 | 139 | def nn_matching(full_pc, idx, k, complete_fps=True): 140 | batch_size = np.size(full_pc, 0) 141 | out_pc = np.zeros((full_pc.shape[0], k, 3)) 142 | for ii in range(0, batch_size): 143 | best_idx = idx[ii] 144 | if complete_fps: 145 | best_idx = unique(best_idx) 146 | out_pc[ii] = fps_from_given_pc(full_pc[ii], k, full_pc[ii][best_idx]) 147 | else: 148 | out_pc[ii] = full_pc[ii][best_idx] 149 | return out_pc[:, 0:k, :] 150 | 151 | 152 | def emd_matching(full_pc, gen_pc, sess): 153 | batch_size = np.size(full_pc, 0) 154 | k = np.size(gen_pc, 1) 155 | out_pc = np.zeros_like(gen_pc) 156 | 157 | match_mat_tensor = approx_match( 158 | tf.convert_to_tensor(full_pc), tf.convert_to_tensor(gen_pc) 159 | ) 160 | pc1_match_idx_tensor = tf.cast(tf.argmax(match_mat_tensor, axis=2), dtype=tf.int32) 161 | 162 | pc1_match_idx = pc1_match_idx_tensor.eval(session=sess) 163 | 164 | for ii in range(0, batch_size): 165 | best_idx = unique(pc1_match_idx[ii]) 166 | out_pc[ii] = fps_from_given_pc(full_pc[ii], k, full_pc[ii][best_idx]) 167 | 168 | return out_pc 169 | 170 | 171 | def get_nn_indices(ref_pc, samp_pc): 172 | _, idx, _, _ = nn_distance(samp_pc, ref_pc) 173 | return idx 174 | 175 | 176 | def get_simplification_loss(ref_pc, samp_pc, pc_size, gamma=1, delta=0): 177 | cost_p1_p2, _, cost_p2_p1, _ = nn_distance(samp_pc, ref_pc) 178 | max_cost = tf.reduce_max(cost_p1_p2, axis=1) 179 | max_cost = tf.reduce_mean(max_cost) 180 | cost_p1_p2 = tf.reduce_mean(cost_p1_p2) 181 | cost_p2_p1 = tf.reduce_mean(cost_p2_p1) 182 | loss = cost_p1_p2 + max_cost + (gamma + delta * pc_size) * cost_p2_p1 183 | 184 | tf.summary.scalar("cost_p1_p2", cost_p1_p2) 185 | tf.summary.scalar("cost_p2_p1", cost_p2_p1) 186 | tf.summary.scalar("max_cost", max_cost) 187 | 188 | return loss 189 | -------------------------------------------------------------------------------- /reconstruction/sampler/train_samplenet.py: -------------------------------------------------------------------------------- 1 | """ 2 | Created on September 24th, 2019 3 | 4 | @author: itailang 5 | """ 6 | from __future__ import print_function 7 | 8 | # import system modules 9 | from builtins import str 10 | from builtins import range 11 | import os.path as osp 12 | import sys 13 | import argparse 14 | 15 | # add paths 16 | parent_dir = osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))) 17 | if parent_dir not in sys.path: 18 | sys.path.append(parent_dir) 19 | 20 | # import modules 21 | from reconstruction.src.samplers import sampler_with_convs_and_symmetry_and_fc 22 | from reconstruction.src.autoencoder import Configuration as Conf 23 | from reconstruction.src.samplenet_pointnet_ae import SampleNetPointNetAE 24 | 25 | from reconstruction.src.in_out import ( 26 | snc_category_to_synth_id, 27 | create_dir, 28 | load_and_split_all_point_clouds_under_folder, 29 | ) 30 | 31 | from reconstruction.src.tf_utils import reset_tf_graph 32 | 33 | # command line arguments 34 | # fmt: off 35 | parser = argparse.ArgumentParser() 36 | parser.add_argument('--learning_rate', type=float, default=0.0005, help='Learning rate [default: 0.0005]') 37 | parser.add_argument('--training_epochs', type=int, default=400, help='Number of training epochs [default: 400]') 38 | parser.add_argument('--restore_ae', type=bool, default=True, help='Restore a trained autoencoder model [default: True]') 39 | parser.add_argument('--ae_restore_epoch', type=int, default=500, help='Epoch to restore a trained autoencoder model [default: 500]') 40 | parser.add_argument('--fixed_ae', type=bool, default=True, help='Fixed autoencoder model [default: True]') 41 | parser.add_argument('--ae_folder', type=str, default='log/autoencoder', help='Folder for loading a trained autoencoder model [default: log/autoencoder]') 42 | parser.add_argument('--object_class', type=str, default='multi', help='Single class name (for example: chair) or multi [default: multi]') 43 | 44 | # sampler arguments 45 | parser.add_argument('--n_sample_points', type=int, default=64, help='Number of sample points [default: 64]') 46 | parser.add_argument('--alpha', type=float, default=0.01, help='Sampling regularization loss weight [default: 0.01]') 47 | parser.add_argument('--train_folder', type=str, default='log/SampleNet64', help='Folder for saving data form the training [default: log/SampleNet64]') 48 | 49 | # projection arguments 50 | parser.add_argument('--projection_group_size', type=int, default=16, help='Neighborhood size in Soft Projection [default: 16]') 51 | parser.add_argument('--lmbda', type=float, default=0.0001, help='Temperature regularization loss weight [default: 0.0001]') 52 | flags = parser.parse_args() 53 | # fmt: on 54 | 55 | print(("Train flags:", flags)) 56 | 57 | # define basic parameters 58 | project_dir = osp.dirname(osp.dirname(osp.abspath(__file__))) 59 | top_in_dir = osp.join( 60 | project_dir, "data", "shape_net_core_uniform_samples_2048" 61 | ) # top-dir of where point-clouds are stored. 62 | top_out_dir = osp.join(project_dir) # use to save Neural-Net check-points etc. 63 | 64 | if flags.object_class == "multi": 65 | class_name = ["chair", "table", "car", "airplane"] 66 | else: 67 | class_name = [str(flags.object_class)] 68 | 69 | # load Point-Clouds 70 | syn_id = snc_category_to_synth_id()[class_name[0]] 71 | class_dir = osp.join(top_in_dir, syn_id) 72 | pc_data_train, pc_data_val, _ = load_and_split_all_point_clouds_under_folder( 73 | class_dir, n_threads=8, file_ending=".ply", verbose=True 74 | ) 75 | 76 | for i in range(1, len(class_name)): 77 | syn_id = snc_category_to_synth_id()[class_name[i]] 78 | class_dir = osp.join(top_in_dir, syn_id) 79 | ( 80 | pc_data_train_curr, 81 | pc_data_val_curr, 82 | _, 83 | ) = load_and_split_all_point_clouds_under_folder( 84 | class_dir, n_threads=8, file_ending=".ply", verbose=True 85 | ) 86 | pc_data_train.merge(pc_data_train_curr) 87 | pc_data_val.merge(pc_data_val_curr) 88 | 89 | if flags.object_class == "multi": 90 | pc_data_train.shuffle_data(seed=55) 91 | pc_data_val.shuffle_data(seed=55) 92 | 93 | ae_dir = osp.join(top_out_dir, flags.ae_folder) 94 | 95 | # load autoencoder configuration 96 | conf = Conf.load(osp.join(ae_dir, "configuration")) 97 | 98 | # update autoencoder configuration 99 | conf.ae_dir = ae_dir 100 | conf.ae_name = "autoencoder" 101 | conf.restore_ae = flags.restore_ae 102 | conf.ae_restore_epoch = flags.ae_restore_epoch 103 | conf.fixed_ae = flags.fixed_ae 104 | if conf.fixed_ae: 105 | conf.encoder_args[ 106 | "b_norm_decay" 107 | ] = 1.0 # for avoiding the update of batch normalization moving_mean and moving_variance parameters 108 | conf.decoder_args[ 109 | "b_norm_decay" 110 | ] = 1.0 # for avoiding the update of batch normalization moving_mean and moving_variance parameters 111 | conf.decoder_args[ 112 | "b_norm_decay_finish" 113 | ] = 1.0 # for avoiding the update of batch normalization moving_mean and moving_variance parameters 114 | 115 | conf.use_batch_size_for_place_holder = True 116 | 117 | # sampler configuration 118 | conf.experiment_name = "sampler" 119 | conf.n_samp = [flags.n_sample_points, 3] # Dimensionality of sampled points. 120 | conf.sampler = sampler_with_convs_and_symmetry_and_fc 121 | conf.alpha = flags.alpha 122 | conf.learning_rate = flags.learning_rate 123 | conf.training_epochs = flags.training_epochs 124 | 125 | # projector configuration 126 | conf.lmbda = flags.lmbda 127 | conf.projection_group_size = flags.projection_group_size 128 | conf.hard_projection = 0 129 | 130 | train_dir = create_dir(osp.join(top_out_dir, flags.train_folder)) 131 | conf.train_dir = train_dir 132 | 133 | conf.save(osp.join(train_dir, "configuration")) 134 | 135 | # build Sampler and AE Model 136 | reset_tf_graph() 137 | ae = SampleNetPointNetAE(conf.experiment_name, conf) 138 | 139 | # train the sampler (save output to train_stats.txt) 140 | buf_size = 1 # Make 'training_stats' file to flush each output line regarding training. 141 | fout = open(osp.join(conf.train_dir, "train_stats.txt"), "a", buf_size) 142 | fout.write("Train flags: %s\n" % flags) 143 | train_stats = ae.train(pc_data_train, conf, log_file=fout, held_out_data=pc_data_val) 144 | fout.close() 145 | -------------------------------------------------------------------------------- /classification/structural_losses/tf_approxmatch.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | from builtins import range 4 | import tensorflow as tf 5 | from tensorflow.python.framework import ops 6 | import os.path as osp 7 | 8 | base_dir = osp.dirname(osp.abspath(__file__)) 9 | 10 | approxmatch_module = tf.load_op_library(osp.join(base_dir, "tf_approxmatch_so.so")) 11 | 12 | 13 | def approx_match(xyz1, xyz2): 14 | """ 15 | input: 16 | xyz1 : batch_size * #dataset_points * 3 17 | xyz2 : batch_size * #query_points * 3 18 | returns: 19 | match : batch_size * #query_points * #dataset_points 20 | """ 21 | return approxmatch_module.approx_match(xyz1, xyz2) 22 | 23 | 24 | ops.NoGradient("ApproxMatch") 25 | # @tf.RegisterShape('ApproxMatch') 26 | @ops.RegisterShape("ApproxMatch") 27 | def _approx_match_shape(op): 28 | shape1 = op.inputs[0].get_shape().with_rank(3) 29 | shape2 = op.inputs[1].get_shape().with_rank(3) 30 | return [tf.TensorShape([shape1.dims[0], shape2.dims[1], shape1.dims[1]])] 31 | 32 | 33 | def match_cost(xyz1, xyz2, match): 34 | """ 35 | input: 36 | xyz1 : batch_size * #dataset_points * 3 37 | xyz2 : batch_size * #query_points * 3 38 | match : batch_size * #query_points * #dataset_points 39 | returns: 40 | cost : batch_size 41 | """ 42 | return approxmatch_module.match_cost(xyz1, xyz2, match) 43 | 44 | 45 | # @tf.RegisterShape('MatchCost') 46 | @ops.RegisterShape("MatchCost") 47 | def _match_cost_shape(op): 48 | shape1 = op.inputs[0].get_shape().with_rank(3) 49 | shape2 = op.inputs[1].get_shape().with_rank(3) 50 | shape3 = op.inputs[2].get_shape().with_rank(3) 51 | return [tf.TensorShape([shape1.dims[0]])] 52 | 53 | 54 | @tf.RegisterGradient("MatchCost") 55 | def _match_cost_grad(op, grad_cost): 56 | xyz1 = op.inputs[0] 57 | xyz2 = op.inputs[1] 58 | match = op.inputs[2] 59 | grad_1, grad_2 = approxmatch_module.match_cost_grad(xyz1, xyz2, match) 60 | return [ 61 | grad_1 * tf.expand_dims(tf.expand_dims(grad_cost, 1), 2), 62 | grad_2 * tf.expand_dims(tf.expand_dims(grad_cost, 1), 2), 63 | None, 64 | ] 65 | 66 | 67 | if __name__ == "__main__": 68 | alpha = 0.5 69 | beta = 2.0 70 | import bestmatch 71 | import numpy as np 72 | import math 73 | import random 74 | import cv2 75 | 76 | from . import tf_nndistance 77 | 78 | npoint = 100 79 | 80 | with tf.device("/gpu:2"): 81 | pt_in = tf.placeholder(tf.float32, shape=(1, npoint * 4, 3)) 82 | mypoints = tf.Variable(np.random.randn(1, npoint, 3).astype("float32")) 83 | match = approx_match(pt_in, mypoints) 84 | loss = tf.reduce_sum(match_cost(pt_in, mypoints, match)) 85 | # match=approx_match(mypoints,pt_in) 86 | # loss=tf.reduce_sum(match_cost(mypoints,pt_in,match)) 87 | # distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints) 88 | # loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5 89 | # loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint 90 | 91 | optimizer = tf.train.GradientDescentOptimizer(1e-4).minimize(loss) 92 | with tf.Session("") as sess: 93 | sess.run(tf.initialize_all_variables()) 94 | while True: 95 | meanloss = 0 96 | meantrueloss = 0 97 | for i in range(1001): 98 | # phi=np.random.rand(4*npoint)*math.pi*2 99 | # tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:] 100 | # tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:] 101 | tpoints = np.hstack( 102 | [ 103 | np.linspace(-1, 1, 400)[:, None], 104 | (random.random() * 2 * np.linspace(1, 0, 400) ** 2)[:, None], 105 | np.zeros((400, 1)), 106 | ] 107 | )[None, :, :] 108 | trainloss, _ = sess.run( 109 | [loss, optimizer], feed_dict={pt_in: tpoints.astype("float32")} 110 | ) 111 | trainloss, trainmatch = sess.run( 112 | [loss, match], feed_dict={pt_in: tpoints.astype("float32")} 113 | ) 114 | # trainmatch=trainmatch.transpose((0,2,1)) 115 | show = np.zeros((400, 400, 3), dtype="uint8") ^ 255 116 | trainmypoints = sess.run(mypoints) 117 | for i in range(len(tpoints[0])): 118 | u = np.random.choice( 119 | list(range(len(trainmypoints[0]))), p=trainmatch[0].T[i] 120 | ) 121 | cv2.line( 122 | show, 123 | ( 124 | int(tpoints[0][i, 1] * 100 + 200), 125 | int(tpoints[0][i, 0] * 100 + 200), 126 | ), 127 | ( 128 | int(trainmypoints[0][u, 1] * 100 + 200), 129 | int(trainmypoints[0][u, 0] * 100 + 200), 130 | ), 131 | cv2.cv.CV_RGB(0, 255, 0), 132 | ) 133 | for x, y, z in tpoints[0]: 134 | cv2.circle( 135 | show, 136 | (int(y * 100 + 200), int(x * 100 + 200)), 137 | 2, 138 | cv2.cv.CV_RGB(255, 0, 0), 139 | ) 140 | for x, y, z in trainmypoints[0]: 141 | cv2.circle( 142 | show, 143 | (int(y * 100 + 200), int(x * 100 + 200)), 144 | 3, 145 | cv2.cv.CV_RGB(0, 0, 255), 146 | ) 147 | cost = ( 148 | ( 149 | tpoints[0][:, None, :] 150 | - np.repeat(trainmypoints[0][None, :, :], 4, axis=1) 151 | ) 152 | ** 2 153 | ).sum(axis=2) ** 0.5 154 | # trueloss=bestmatch.bestmatch(cost)[0] 155 | print(trainloss) # ,trueloss) 156 | cv2.imshow("show", show) 157 | cmd = cv2.waitKey(10) % 256 158 | if cmd == ord("q"): 159 | break 160 | -------------------------------------------------------------------------------- /reconstruction/external/structural_losses/tf_approxmatch.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from __future__ import absolute_import 3 | from builtins import range 4 | import tensorflow as tf 5 | from tensorflow.python.framework import ops 6 | import os.path as osp 7 | 8 | base_dir = osp.dirname(osp.abspath(__file__)) 9 | 10 | approxmatch_module = tf.load_op_library(osp.join(base_dir, "tf_approxmatch_so.so")) 11 | 12 | 13 | def approx_match(xyz1, xyz2): 14 | """ 15 | input: 16 | xyz1 : batch_size * #dataset_points * 3 17 | xyz2 : batch_size * #query_points * 3 18 | returns: 19 | match : batch_size * #query_points * #dataset_points 20 | """ 21 | return approxmatch_module.approx_match(xyz1, xyz2) 22 | 23 | 24 | ops.NoGradient("ApproxMatch") 25 | # @tf.RegisterShape('ApproxMatch') 26 | @ops.RegisterShape("ApproxMatch") 27 | def _approx_match_shape(op): 28 | shape1 = op.inputs[0].get_shape().with_rank(3) 29 | shape2 = op.inputs[1].get_shape().with_rank(3) 30 | return [tf.TensorShape([shape1.dims[0], shape2.dims[1], shape1.dims[1]])] 31 | 32 | 33 | def match_cost(xyz1, xyz2, match): 34 | """ 35 | input: 36 | xyz1 : batch_size * #dataset_points * 3 37 | xyz2 : batch_size * #query_points * 3 38 | match : batch_size * #query_points * #dataset_points 39 | returns: 40 | cost : batch_size 41 | """ 42 | return approxmatch_module.match_cost(xyz1, xyz2, match) 43 | 44 | 45 | # @tf.RegisterShape('MatchCost') 46 | @ops.RegisterShape("MatchCost") 47 | def _match_cost_shape(op): 48 | shape1 = op.inputs[0].get_shape().with_rank(3) 49 | shape2 = op.inputs[1].get_shape().with_rank(3) 50 | shape3 = op.inputs[2].get_shape().with_rank(3) 51 | return [tf.TensorShape([shape1.dims[0]])] 52 | 53 | 54 | @tf.RegisterGradient("MatchCost") 55 | def _match_cost_grad(op, grad_cost): 56 | xyz1 = op.inputs[0] 57 | xyz2 = op.inputs[1] 58 | match = op.inputs[2] 59 | grad_1, grad_2 = approxmatch_module.match_cost_grad(xyz1, xyz2, match) 60 | return [ 61 | grad_1 * tf.expand_dims(tf.expand_dims(grad_cost, 1), 2), 62 | grad_2 * tf.expand_dims(tf.expand_dims(grad_cost, 1), 2), 63 | None, 64 | ] 65 | 66 | 67 | if __name__ == "__main__": 68 | alpha = 0.5 69 | beta = 2.0 70 | import bestmatch 71 | import numpy as np 72 | import math 73 | import random 74 | import cv2 75 | 76 | from . import tf_nndistance 77 | 78 | npoint = 100 79 | 80 | with tf.device("/gpu:2"): 81 | pt_in = tf.placeholder(tf.float32, shape=(1, npoint * 4, 3)) 82 | mypoints = tf.Variable(np.random.randn(1, npoint, 3).astype("float32")) 83 | match = approx_match(pt_in, mypoints) 84 | loss = tf.reduce_sum(match_cost(pt_in, mypoints, match)) 85 | # match=approx_match(mypoints,pt_in) 86 | # loss=tf.reduce_sum(match_cost(mypoints,pt_in,match)) 87 | # distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints) 88 | # loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5 89 | # loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint 90 | 91 | optimizer = tf.train.GradientDescentOptimizer(1e-4).minimize(loss) 92 | with tf.Session("") as sess: 93 | sess.run(tf.initialize_all_variables()) 94 | while True: 95 | meanloss = 0 96 | meantrueloss = 0 97 | for i in range(1001): 98 | # phi=np.random.rand(4*npoint)*math.pi*2 99 | # tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:] 100 | # tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:] 101 | tpoints = np.hstack( 102 | [ 103 | np.linspace(-1, 1, 400)[:, None], 104 | (random.random() * 2 * np.linspace(1, 0, 400) ** 2)[:, None], 105 | np.zeros((400, 1)), 106 | ] 107 | )[None, :, :] 108 | trainloss, _ = sess.run( 109 | [loss, optimizer], feed_dict={pt_in: tpoints.astype("float32")} 110 | ) 111 | trainloss, trainmatch = sess.run( 112 | [loss, match], feed_dict={pt_in: tpoints.astype("float32")} 113 | ) 114 | # trainmatch=trainmatch.transpose((0,2,1)) 115 | show = np.zeros((400, 400, 3), dtype="uint8") ^ 255 116 | trainmypoints = sess.run(mypoints) 117 | for i in range(len(tpoints[0])): 118 | u = np.random.choice( 119 | list(range(len(trainmypoints[0]))), p=trainmatch[0].T[i] 120 | ) 121 | cv2.line( 122 | show, 123 | ( 124 | int(tpoints[0][i, 1] * 100 + 200), 125 | int(tpoints[0][i, 0] * 100 + 200), 126 | ), 127 | ( 128 | int(trainmypoints[0][u, 1] * 100 + 200), 129 | int(trainmypoints[0][u, 0] * 100 + 200), 130 | ), 131 | cv2.cv.CV_RGB(0, 255, 0), 132 | ) 133 | for x, y, z in tpoints[0]: 134 | cv2.circle( 135 | show, 136 | (int(y * 100 + 200), int(x * 100 + 200)), 137 | 2, 138 | cv2.cv.CV_RGB(255, 0, 0), 139 | ) 140 | for x, y, z in trainmypoints[0]: 141 | cv2.circle( 142 | show, 143 | (int(y * 100 + 200), int(x * 100 + 200)), 144 | 3, 145 | cv2.cv.CV_RGB(0, 0, 255), 146 | ) 147 | cost = ( 148 | ( 149 | tpoints[0][:, None, :] 150 | - np.repeat(trainmypoints[0][None, :, :], 4, axis=1) 151 | ) 152 | ** 2 153 | ).sum(axis=2) ** 0.5 154 | # trueloss=bestmatch.bestmatch(cost)[0] 155 | print(trainloss) # ,trueloss) 156 | cv2.imshow("show", show) 157 | cmd = cv2.waitKey(10) % 256 158 | if cmd == ord("q"): 159 | break 160 | --------------------------------------------------------------------------------