├── LICENSE ├── README.md ├── experiments ├── .DS_Store ├── ex1_genpert.sh ├── ex1_genrot.sh ├── ex1_result_stat.sh ├── ex1_test_icp.sh ├── ex1_test_pointlk.sh ├── ex1_train.sh ├── generate_perturbations.py ├── generate_rotations.py ├── icp.py ├── result_stat.py ├── sampledata │ ├── modelnet40.txt │ ├── modelnet40_half1.txt │ ├── modelnet40_half2.txt │ └── shapenet2_20.txt ├── test_icp.py ├── test_icp_vs_pointlk.py ├── test_pointlk.py ├── test_utils.py ├── train_classifier.py └── train_pointlk.py └── ptlk ├── __init__.py ├── __pycache__ ├── __init__.cpython-36.pyc ├── invmat.cpython-36.pyc ├── pointlk.cpython-36.pyc ├── pointnet.cpython-36.pyc ├── se3.cpython-36.pyc ├── sinc.cpython-36.pyc └── so3.cpython-36.pyc ├── data ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── datasets.cpython-36.pyc │ ├── globset.cpython-36.pyc │ ├── mesh.cpython-36.pyc │ └── transforms.cpython-36.pyc ├── datasets.py ├── globset.py ├── mesh.py └── transforms.py ├── invmat.py ├── pointlk.py ├── pointnet.py ├── se3.py ├── sinc.py └── so3.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Hunter Goforth 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PointNetLK: Point Cloud Registration using PointNet 2 | 3 | ### [Video](https://youtu.be/J2ClR5OZuLc) 4 | 5 | Source Code Author: 6 | Yasuhiro Aoki 7 | 8 | ### Requires: 9 | * PyTorch 0.4.0 (perhaps, 0.4.1 (the latest) will be OK.) and torchvision 10 | * NumPy 11 | * SciPy 12 | * MatPlotLib 13 | * ModelNet40 14 | 15 | ### Main files for experiments: 16 | * train_classifier.py: train PointNet classifier (used for transfer learning) 17 | * train_pointlk.py: train PointNet-LK 18 | * generate_rotation.py: generate 6-dim perturbations (rotation and translation) (for testing) 19 | * test_pointlk.py: test PointNet-LK 20 | * test_icp.py: test ICP 21 | * result_stat.py: compute mean errors of above tests 22 | 23 | ### Examples (Bash shell scripts): 24 | * ex1_train.sh: train PointNet classifier and transfer to PointNet-LK. 25 | * ex1_genrot.sh: generate perturbations for testing 26 | * ex1_test_pointlk.sh: test PointNet-LK 27 | * ex1_test_icp.sh: test ICP 28 | * ex1_result_stat.sh: compute mean errors of above tests 29 | 30 | ### Citation 31 | 32 | ``` 33 | @InProceedings{yaoki2019pointnetlk, 34 | author = {Aoki, Yasuhiro and Goforth, Hunter and Arun Srivatsan, Rangaprasad and Lucey, Simon}, 35 | title = {PointNetLK: Robust & Efficient Point Cloud Registration Using PointNet}, 36 | booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 37 | month = {June}, 38 | year = {2019} 39 | } 40 | ``` 41 | -------------------------------------------------------------------------------- /experiments/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/experiments/.DS_Store -------------------------------------------------------------------------------- /experiments/ex1_genpert.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # generate perturbations for each object (for each 'ModelNet40/[category]/test/*') 4 | 5 | # for output 6 | OUTDIR=${HOME}/results/gt 7 | mkdir -p ${OUTDIR} 8 | 9 | # Python3 command 10 | PY3="nice -n 10 python" 11 | 12 | # categories for testing 13 | CMN="-i /home/yasuhiro/work/pointnet/ModelNet40 -c ./sampledata/modelnet40_half1.txt --no-translation" 14 | 15 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_00.csv --mag 0.0 16 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_01.csv --mag 0.1 17 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_02.csv --mag 0.2 18 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_03.csv --mag 0.3 19 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_04.csv --mag 0.4 20 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_05.csv --mag 0.5 21 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_06.csv --mag 0.6 22 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_07.csv --mag 0.7 23 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_08.csv --mag 0.8 24 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_09.csv --mag 0.9 25 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_10.csv --mag 1.0 26 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_11.csv --mag 1.1 27 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_12.csv --mag 1.2 28 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_13.csv --mag 1.3 29 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_14.csv --mag 1.4 30 | ${PY3} generate_perturbations.py ${CMN} -o ${OUTDIR}/pert_15.csv --mag 1.5 31 | 32 | 33 | #EOF 34 | -------------------------------------------------------------------------------- /experiments/ex1_genrot.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # generate perturbations for each object (for each 'ModelNet40/[category]/test/*') 4 | 5 | # for output 6 | OUTDIR=${HOME}/results/ex1/gt 7 | mkdir -p ${OUTDIR} 8 | 9 | # Python3 command 10 | PY3="nice -n 10 python" 11 | 12 | # categories for testing 13 | CMN="-i /home/yasuhiro/work/pointnet/ModelNet40 -c ./sampledata/modelnet40_half1.txt --format wt" 14 | 15 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_000.csv --deg 0.0 16 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_010.csv --deg 10.0 17 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_020.csv --deg 20.0 18 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_030.csv --deg 30.0 19 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_040.csv --deg 40.0 20 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_050.csv --deg 50.0 21 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_060.csv --deg 60.0 22 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_070.csv --deg 70.0 23 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_080.csv --deg 80.0 24 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_090.csv --deg 90.0 25 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_100.csv --deg 100.0 26 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_110.csv --deg 110.0 27 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_120.csv --deg 120.0 28 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_130.csv --deg 130.0 29 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_140.csv --deg 140.0 30 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_150.csv --deg 150.0 31 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_160.csv --deg 160.0 32 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_170.csv --deg 170.0 33 | ${PY3} generate_rotations.py ${CMN} -o ${OUTDIR}/pert_180.csv --deg 180.0 34 | 35 | 36 | #EOF 37 | -------------------------------------------------------------------------------- /experiments/ex1_result_stat.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # Python3 command 4 | PY3="nice -n 10 python" 5 | 6 | # for output 7 | #RES="${HOME}/results/ex1/icp" 8 | RES="${HOME}/results/ex1/plk" 9 | 10 | # gather 'result_*.csv' to 'result.csv' 11 | ${PY3} result_stat.py --hdr > ${RES}/result.csv 12 | ${PY3} result_stat.py -i ${RES}/result_000.csv --val 0 >> ${RES}/result.csv 13 | ${PY3} result_stat.py -i ${RES}/result_010.csv --val 10 >> ${RES}/result.csv 14 | ${PY3} result_stat.py -i ${RES}/result_020.csv --val 20 >> ${RES}/result.csv 15 | ${PY3} result_stat.py -i ${RES}/result_030.csv --val 30 >> ${RES}/result.csv 16 | ${PY3} result_stat.py -i ${RES}/result_040.csv --val 40 >> ${RES}/result.csv 17 | ${PY3} result_stat.py -i ${RES}/result_050.csv --val 50 >> ${RES}/result.csv 18 | ${PY3} result_stat.py -i ${RES}/result_060.csv --val 60 >> ${RES}/result.csv 19 | ${PY3} result_stat.py -i ${RES}/result_070.csv --val 70 >> ${RES}/result.csv 20 | ${PY3} result_stat.py -i ${RES}/result_080.csv --val 80 >> ${RES}/result.csv 21 | ${PY3} result_stat.py -i ${RES}/result_090.csv --val 90 >> ${RES}/result.csv 22 | ${PY3} result_stat.py -i ${RES}/result_100.csv --val 100 >> ${RES}/result.csv 23 | ${PY3} result_stat.py -i ${RES}/result_110.csv --val 110 >> ${RES}/result.csv 24 | ${PY3} result_stat.py -i ${RES}/result_120.csv --val 120 >> ${RES}/result.csv 25 | ${PY3} result_stat.py -i ${RES}/result_130.csv --val 130 >> ${RES}/result.csv 26 | ${PY3} result_stat.py -i ${RES}/result_140.csv --val 140 >> ${RES}/result.csv 27 | ${PY3} result_stat.py -i ${RES}/result_150.csv --val 150 >> ${RES}/result.csv 28 | ${PY3} result_stat.py -i ${RES}/result_160.csv --val 160 >> ${RES}/result.csv 29 | ${PY3} result_stat.py -i ${RES}/result_170.csv --val 170 >> ${RES}/result.csv 30 | ${PY3} result_stat.py -i ${RES}/result_180.csv --val 180 >> ${RES}/result.csv 31 | 32 | #EOF -------------------------------------------------------------------------------- /experiments/ex1_test_icp.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # for output 4 | OUTDIR=${HOME}/results/ex1/icp 5 | mkdir -p ${OUTDIR} 6 | 7 | # Python3 command 8 | PY3="nice -n 10 python" 9 | 10 | # categories for testing 11 | CMN="-i /home/yasuhiro/work/pointnet/ModelNet40 -c ./sampledata/modelnet40_half1.txt --format wt" 12 | 13 | # NOTICE! before testing, generate perturbations 14 | # see) 'ex1_genrot.sh' 15 | 16 | # perturbations 17 | PERDIR=${HOME}/results/ex1/gt 18 | 19 | # test ICP with given perturbations. 20 | # probably, this will take several days... 21 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_000.csv -p ${PERDIR}/pert_000.csv -l ${OUTDIR}/log_000.log & 22 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_010.csv -p ${PERDIR}/pert_010.csv -l ${OUTDIR}/log_010.log 23 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_020.csv -p ${PERDIR}/pert_020.csv -l ${OUTDIR}/log_020.log & 24 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_030.csv -p ${PERDIR}/pert_030.csv -l ${OUTDIR}/log_030.log 25 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_040.csv -p ${PERDIR}/pert_040.csv -l ${OUTDIR}/log_040.log & 26 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_050.csv -p ${PERDIR}/pert_050.csv -l ${OUTDIR}/log_050.log 27 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_060.csv -p ${PERDIR}/pert_060.csv -l ${OUTDIR}/log_060.log & 28 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_070.csv -p ${PERDIR}/pert_070.csv -l ${OUTDIR}/log_070.log 29 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_080.csv -p ${PERDIR}/pert_080.csv -l ${OUTDIR}/log_080.log & 30 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_090.csv -p ${PERDIR}/pert_090.csv -l ${OUTDIR}/log_090.log 31 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_100.csv -p ${PERDIR}/pert_100.csv -l ${OUTDIR}/log_100.log & 32 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_110.csv -p ${PERDIR}/pert_110.csv -l ${OUTDIR}/log_110.log 33 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_120.csv -p ${PERDIR}/pert_120.csv -l ${OUTDIR}/log_120.log & 34 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_130.csv -p ${PERDIR}/pert_130.csv -l ${OUTDIR}/log_130.log 35 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_140.csv -p ${PERDIR}/pert_140.csv -l ${OUTDIR}/log_140.log & 36 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_150.csv -p ${PERDIR}/pert_150.csv -l ${OUTDIR}/log_150.log 37 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_160.csv -p ${PERDIR}/pert_160.csv -l ${OUTDIR}/log_160.log & 38 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_170.csv -p ${PERDIR}/pert_170.csv -l ${OUTDIR}/log_170.log 39 | ${PY3} test_icp.py ${CMN} -o ${OUTDIR}/result_180.csv -p ${PERDIR}/pert_180.csv -l ${OUTDIR}/log_180.log 40 | 41 | #EOF 42 | -------------------------------------------------------------------------------- /experiments/ex1_test_pointlk.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # for output 4 | OUTDIR=${HOME}/results/ex1/plk 5 | mkdir -p ${OUTDIR} 6 | 7 | # Python3 command 8 | PY3="nice -n 10 python" 9 | 10 | # categories for testing and the trained model 11 | MODEL=${HOME}/results/ex1_pointlk_0915_model_best.pth 12 | CMN="-i /home/yasuhiro/work/pointnet/ModelNet40 -c ./sampledata/modelnet40_half1.txt --format wt --pretrained ${MODEL}" 13 | 14 | # perturbations 15 | PERDIR=${HOME}/results/ex1/gt 16 | 17 | # test PointNet-LK with given perturbations (see. 'ex1_genrot.sh' for perturbations) 18 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_010.csv -p ${PERDIR}/pert_010.csv -l ${OUTDIR}/log_010.log 19 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_000.csv -p ${PERDIR}/pert_000.csv -l ${OUTDIR}/log_000.log 20 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_090.csv -p ${PERDIR}/pert_090.csv -l ${OUTDIR}/log_090.log 21 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_060.csv -p ${PERDIR}/pert_060.csv -l ${OUTDIR}/log_060.log 22 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_120.csv -p ${PERDIR}/pert_120.csv -l ${OUTDIR}/log_120.log 23 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_030.csv -p ${PERDIR}/pert_030.csv -l ${OUTDIR}/log_030.log 24 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_180.csv -p ${PERDIR}/pert_180.csv -l ${OUTDIR}/log_180.log 25 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_020.csv -p ${PERDIR}/pert_020.csv -l ${OUTDIR}/log_020.log 26 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_040.csv -p ${PERDIR}/pert_040.csv -l ${OUTDIR}/log_040.log 27 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_050.csv -p ${PERDIR}/pert_050.csv -l ${OUTDIR}/log_050.log 28 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_070.csv -p ${PERDIR}/pert_070.csv -l ${OUTDIR}/log_070.log 29 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_080.csv -p ${PERDIR}/pert_080.csv -l ${OUTDIR}/log_080.log 30 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_100.csv -p ${PERDIR}/pert_100.csv -l ${OUTDIR}/log_100.log 31 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_110.csv -p ${PERDIR}/pert_110.csv -l ${OUTDIR}/log_110.log 32 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_130.csv -p ${PERDIR}/pert_130.csv -l ${OUTDIR}/log_130.log 33 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_140.csv -p ${PERDIR}/pert_140.csv -l ${OUTDIR}/log_140.log 34 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_150.csv -p ${PERDIR}/pert_150.csv -l ${OUTDIR}/log_150.log 35 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_160.csv -p ${PERDIR}/pert_160.csv -l ${OUTDIR}/log_160.log 36 | ${PY3} test_pointlk.py ${CMN} -o ${OUTDIR}/result_170.csv -p ${PERDIR}/pert_170.csv -l ${OUTDIR}/log_170.log 37 | 38 | 39 | #EOF 40 | -------------------------------------------------------------------------------- /experiments/ex1_train.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | # for output 4 | mkdir -p ${HOME}/results 5 | 6 | # Python3 command 7 | PY3="nice -n 10 python" 8 | 9 | # first, traing a classifier. 10 | # ModelNet categories are given in './sampledata/modelnet40_half1.txt' (as examaple) 11 | ${PY3} train_classifier.py \ 12 | -o ${HOME}/results/ex1_classifier_0915 \ 13 | -i /home/yasuhiro/work/pointnet/ModelNet40 \ 14 | -c ./sampledata/modelnet40_half1.txt \ 15 | -l ${HOME}/results/ex1_classifier_0915.log \ 16 | --device cuda:2 17 | 18 | # the one of the results is '${HOME}/results/ex1_classifier_0915_feat_best.pth' 19 | # this file is the model that computes PointNet feature. 20 | 21 | # train PointNet-LK. fine-tune the PointNet feature for classification (the above file). 22 | ${PY3} train_pointlk.py \ 23 | -o ${HOME}/results/ex1_pointlk_0915 \ 24 | -i /home/yasuhiro/work/pointnet/ModelNet40 \ 25 | -c ./sampledata/modelnet40_half1.txt \ 26 | -l ${HOME}/results/ex1_pointlk_0915.log \ 27 | --transfer-from ${HOME}/results/ex1_classifier_0915_feat_best.pth \ 28 | --epochs 400 \ 29 | --device cuda:2 30 | 31 | # the trained model: 32 | # ${HOME}/results/ex1_pointlk_0915_model_best.pth 33 | 34 | #EOF 35 | -------------------------------------------------------------------------------- /experiments/generate_perturbations.py: -------------------------------------------------------------------------------- 1 | """ 2 | generate perturbations for testing PointNet-LK and/or ICP 3 | """ 4 | 5 | import argparse 6 | import os 7 | import sys 8 | import logging 9 | import numpy 10 | import torch 11 | import torch.utils.data 12 | import torchvision 13 | 14 | # addpath('../') 15 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 16 | import ptlk 17 | 18 | def options(argv=None): 19 | parser = argparse.ArgumentParser(description='PointNet-LK') 20 | 21 | # required. 22 | parser.add_argument('-o', '--outfile', required=True, type=str, 23 | metavar='FILENAME', help='output filename (.csv)') # the perturbation file for 'test_pointlk.py' 24 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 25 | metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40' 26 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 27 | metavar='PATH', help='path to the categories to be tested') # eg. './sampledata/modelnet40_half1.txt' 28 | 29 | # settings for input data 30 | parser.add_argument('--mag', default=1.0, type=float, 31 | metavar='T', help='mag. of twist-vectors (perturbations) (default: 1.0)') 32 | parser.add_argument('--mag-randomry', dest='mag_randomly', action='store_true', 33 | help='flag to no-fix mag. of twist-vectors') 34 | parser.add_argument('--no-translation', dest='no_translation', action='store_true', 35 | help='generate twist-vectors as (w, 0) (rotation-only)') 36 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet'], 37 | metavar='DATASET', help='dataset type (default: modelnet)') 38 | 39 | args = parser.parse_args(argv) 40 | return args 41 | 42 | def main(args): 43 | # dataset 44 | testset = get_datasets(args) 45 | batch_size = len(testset) 46 | 47 | if not args.no_translation: 48 | x = ptlk.data.datasets.CADset4tracking_fixed_perturbation.generate_perturbations(\ 49 | batch_size, args.mag, \ 50 | randomly=args.mag_randomly) 51 | else: 52 | x = ptlk.data.datasets.CADset4tracking_fixed_perturbation.generate_rotations(\ 53 | batch_size, args.mag, \ 54 | randomly=args.mag_randomly) 55 | 56 | numpy.savetxt(args.outfile, x, delimiter=',') 57 | 58 | def get_datasets(args): 59 | 60 | cinfo = None 61 | if args.categoryfile: 62 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 63 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 64 | categories.sort() 65 | c_to_idx = {categories[i]: i for i in range(len(categories))} 66 | cinfo = (categories, c_to_idx) 67 | 68 | if args.dataset_type == 'modelnet': 69 | transform = torchvision.transforms.Compose([\ 70 | ptlk.data.transforms.Mesh2Points(),\ 71 | ptlk.data.transforms.OnUnitCube(),\ 72 | ]) 73 | 74 | testdata = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 75 | 76 | return testdata 77 | 78 | 79 | if __name__ == '__main__': 80 | ARGS = options() 81 | main(ARGS) 82 | 83 | #EOF 84 | -------------------------------------------------------------------------------- /experiments/generate_rotations.py: -------------------------------------------------------------------------------- 1 | """ 2 | generate perturbations for testing PointNet-LK and/or ICP 3 | for ex.1 4 | """ 5 | 6 | import argparse 7 | import os 8 | import sys 9 | import logging 10 | import math 11 | import numpy 12 | import torch 13 | import torch.utils.data 14 | import torchvision 15 | 16 | # addpath('../') 17 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 18 | import ptlk 19 | 20 | def options(argv=None): 21 | parser = argparse.ArgumentParser(description='PointNet-LK') 22 | 23 | # required. 24 | parser.add_argument('-o', '--outfile', required=True, type=str, 25 | metavar='FILENAME', help='output filename (.csv)') # the perturbation file for 'test_pointlk.py' 26 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 27 | metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40' 28 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 29 | metavar='PATH', help='path to the categories to be tested') # eg. './sampledata/modelnet40_half1.txt' 30 | 31 | # settings for input data 32 | parser.add_argument('--deg', default=60, type=float, 33 | metavar='T', help='fixed degree of rotation (perturbations) (default: 60)') 34 | parser.add_argument('--max-trans', default=0.3, type=float, 35 | help='max translation in each axis (default: 0.3)') 36 | parser.add_argument('--format', default='wv', choices=['wv', 'wt'], 37 | help='output format (default: wv (twist-vector), wt means rotation- and translation-vector)') 38 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet'], 39 | metavar='DATASET', help='dataset type (default: modelnet)') 40 | 41 | args = parser.parse_args(argv) 42 | return args 43 | 44 | def main(args): 45 | # dataset 46 | testset = get_datasets(args) 47 | batch_size = len(testset) 48 | 49 | amp = args.deg * math.pi / 180.0 50 | w = torch.randn(batch_size, 3) 51 | w = w / w.norm(p=2, dim=1, keepdim=True) * amp 52 | t = torch.rand(batch_size, 3) * args.max_trans 53 | 54 | if args.format == 'wv': 55 | # the output: twist vectors. 56 | R = ptlk.so3.exp(w) # (N, 3) --> (N, 3, 3) 57 | G = torch.zeros(batch_size, 4, 4) 58 | G[:, 3, 3] = 1 59 | G[:, 0:3, 0:3] = R 60 | G[:, 0:3, 3] = t 61 | 62 | x = ptlk.se3.log(G) # --> (N, 6) 63 | else: 64 | # rotation-vector and translation-vector 65 | x = torch.cat((w, t), dim=1) 66 | 67 | numpy.savetxt(args.outfile, x, delimiter=',') 68 | 69 | def get_datasets(args): 70 | 71 | cinfo = None 72 | if args.categoryfile: 73 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 74 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 75 | categories.sort() 76 | c_to_idx = {categories[i]: i for i in range(len(categories))} 77 | cinfo = (categories, c_to_idx) 78 | 79 | if args.dataset_type == 'modelnet': 80 | transform = torchvision.transforms.Compose([\ 81 | ptlk.data.transforms.Mesh2Points(),\ 82 | ptlk.data.transforms.OnUnitCube(),\ 83 | ]) 84 | 85 | testdata = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 86 | 87 | return testdata 88 | 89 | 90 | if __name__ == '__main__': 91 | ARGS = options() 92 | main(ARGS) 93 | 94 | #EOF 95 | -------------------------------------------------------------------------------- /experiments/icp.py: -------------------------------------------------------------------------------- 1 | """ ICP algorithm 2 | 3 | References: 4 | (ICP) 5 | [1] Paul J. Besl and Neil D. McKay, 6 | "A method for registration of 3-D shapes", 7 | PAMI Vol. 14, Issue 2, pp. 239-256, 1992. 8 | (SVD) 9 | [2] K. S. Arun, T. S. Huang and S. D. Blostein, 10 | "Least-Squares Fitting of Two 3-D Point Sets", 11 | PAMI Vol. 9, Issue 5, pp.698--700, 1987 12 | """ 13 | import numpy as np 14 | from scipy.spatial import KDTree 15 | 16 | def _icp_find_rigid_transform(p_from, p_target): 17 | A, B = np.copy(p_from), np.copy(p_target) 18 | 19 | centroid_A = np.mean(A, axis=0) 20 | centroid_B = np.mean(B, axis=0) 21 | 22 | A -= centroid_A 23 | B -= centroid_B 24 | 25 | H = np.dot(A.T, B) 26 | U, S, Vt = np.linalg.svd(H) 27 | R = np.dot(Vt.T, U.T) 28 | 29 | # special reflection case 30 | if np.linalg.det(R) < 0: 31 | Vt[2,:] *= -1 32 | R = np.dot(Vt.T, U.T) 33 | 34 | t = np.dot(-R, centroid_A) + centroid_B 35 | 36 | return R, t 37 | 38 | def _icp_Rt_to_matrix(R, t): 39 | # matrix M = [R, t; 0, 1] 40 | Rt = np.concatenate((R, np.expand_dims(t.T, axis=-1)), axis=1) 41 | a = np.concatenate((np.zeros_like(t), np.ones(1))) 42 | M = np.concatenate((Rt, np.expand_dims(a, axis=0)), axis=0) 43 | return M 44 | 45 | class ICP: 46 | """ Estimate a rigid-body transform g such that: 47 | p0 = g.p1 48 | """ 49 | def __init__(self, p0, p1): 50 | """ p0.shape == (N, 3) 51 | p1.shape == (N, 3) 52 | """ 53 | self.p0 = p0 54 | self.p1 = p1 55 | leafsize = 1000 56 | self.nearest = KDTree(self.p0, leafsize=leafsize) 57 | self.g_series = None 58 | 59 | def compute(self, max_iter): 60 | ftol = 1.0e-7 61 | dim_k = self.p0.shape[1] 62 | g = np.eye(dim_k + 1, dtype=self.p0.dtype) 63 | p = np.copy(self.p1) 64 | 65 | self.g_series = np.zeros((max_iter + 1, dim_k + 1, dim_k + 1), dtype=g.dtype) 66 | self.g_series[0, :, :] = g 67 | 68 | itr = -1 69 | for itr in range(max_iter): 70 | neighbor_idx = self.nearest.query(p)[1] 71 | targets = self.p0[neighbor_idx] 72 | R, t = _icp_find_rigid_transform(p, targets) 73 | 74 | new_p = np.dot(R, p.T).T + t 75 | if np.sum(np.abs(p - new_p)) < ftol: 76 | break 77 | 78 | p = np.copy(new_p) 79 | dg = _icp_Rt_to_matrix(R, t) 80 | new_g = np.dot(dg, g) 81 | g = np.copy(new_g) 82 | self.g_series[itr + 1, :, :] = g 83 | 84 | self.g_series[(itr+1):, :, :] = g 85 | 86 | return g, p, (itr + 1) 87 | 88 | 89 | 90 | def icp_test(): 91 | from math import sin, cos 92 | import matplotlib.pyplot as plt 93 | from mpl_toolkits.mplot3d import Axes3D 94 | 95 | Y, X = np.mgrid[0:100:5, 0:100:5] 96 | Z = Y ** 2 + X ** 2 97 | A = np.vstack([Y.reshape(-1), X.reshape(-1), Z.reshape(-1)]).T 98 | 99 | R = np.array([ 100 | [cos(-0.279), -sin(-0.279), 0], 101 | [sin(-0.279), cos(-0.279), 0], 102 | [0, 0, 1] 103 | ]) 104 | #R = np.eye(3) 105 | 106 | t = np.array([5.0, 20.0, 10.0]) 107 | #t = np.array([0.0, 0.0, 0.0]) 108 | 109 | B = np.dot(R, A.T).T + t 110 | A = A.astype(B.dtype) 111 | 112 | icp = ICP(A, B) 113 | matrix, points, itr = icp.compute(10) 114 | 115 | print(itr) 116 | print(icp.g_series) 117 | print(icp.g_series[itr]) 118 | print(matrix) 119 | print(R.T) 120 | print(np.dot(-R.T, t)) 121 | 122 | fig = plt.figure() 123 | #ax = Axes3D(fig) 124 | ax = fig.add_subplot(111, projection='3d') 125 | 126 | ax.set_label("x - axis") 127 | ax.set_label("y - axis") 128 | ax.set_label("z - axis") 129 | 130 | ax.plot(A[:,1], A[:,0], A[:,2], "o", color="#cccccc", ms=4, mew=0.5) 131 | ax.plot(points[:,1], points[:,0], points[:,2], "o", color="#00cccc", ms=4, mew=0.5) 132 | ax.plot(B[:,0], B[:,1], B[:,2], "o", color="#ff0000", ms=4, mew=0.5) 133 | 134 | plt.show() 135 | 136 | if __name__ == '__main__': 137 | icp_test() 138 | 139 | #EOF -------------------------------------------------------------------------------- /experiments/result_stat.py: -------------------------------------------------------------------------------- 1 | """ 2 | analyze results... 3 | """ 4 | 5 | import argparse 6 | import os 7 | import sys 8 | import numpy 9 | import torch 10 | import torch.utils.data 11 | import torchvision 12 | 13 | # addpath('../') 14 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 15 | import ptlk 16 | 17 | 18 | def options(argv=None): 19 | parser = argparse.ArgumentParser(description='PointNet-LK') 20 | 21 | parser.add_argument('-i', '--infile', default='', type=str, 22 | metavar='PATH', help='path to the result of the test (.csv)') # see. 'test_pointlk.py' 23 | parser.add_argument('--val', default='#N/A', type=str, 24 | metavar='VAL', help='info. for the x-axis') 25 | parser.add_argument('--hdr', dest='hdr', action='store_true', 26 | help='outputs header line') 27 | 28 | args = parser.parse_args(argv) 29 | return args 30 | 31 | def main(args): 32 | if args.hdr: 33 | # val: given value (for x-axis) 34 | # me_pos: mean error of estimated position (distance) 35 | # me_rad: mean error of estimated rotation (angle in radian) 36 | # me_twist: mean error represented as norm of twist vector 37 | # me_vel: translation part of the twist. (rotation part is me_rad) 38 | print('val, me_pos, me_rad, me_twist, me_vel') 39 | 40 | if args.infile: 41 | npdata = numpy.loadtxt(args.infile, delimiter=',', skiprows=1) # --> (N, 12) 42 | res = torch.from_numpy(npdata).view(-1, 12) 43 | x_hat = res[:, 0:6] # estimated twist vector 44 | x_mgt = -res[:, 6:12] # (minus) ground-truth 45 | 46 | g_hat = ptlk.se3.exp(x_hat) # [N, 4, 4], estimated matrices 47 | g_igt = ptlk.se3.exp(x_mgt) # [N, 4, 4], inverse of ground-truth 48 | dg = g_hat.bmm(g_igt) # [N, 4, 4]. if correct, dg == identity matrices. 49 | 50 | dp = dg[:, 0:3, 3] # [N, 3], position error 51 | dx = ptlk.se3.log(dg) # [N, 6], twist error 52 | dw = dx[:, 0:3] # [N, 3], rotation part of the twist error 53 | dv = dx[:, 3:6] # [N, 3], translation part 54 | 55 | ep = dp.norm(p=2, dim=1) # [N] 56 | ex = dx.norm(p=2, dim=1) # [N] 57 | ew = dw.norm(p=2, dim=1) # [N] 58 | ev = dv.norm(p=2, dim=1) # [N] 59 | 60 | e = torch.stack((ep, ew, ex, ev)) # [4, N] 61 | me = torch.mean(e, dim=1) # [4] 62 | 63 | line = ','.join(map(str, me.numpy().tolist())) 64 | print('{},{}'.format(args.val, line)) 65 | 66 | if __name__ == '__main__': 67 | ARGS = options() 68 | main(ARGS) 69 | 70 | 71 | #EOF -------------------------------------------------------------------------------- /experiments/sampledata/modelnet40.txt: -------------------------------------------------------------------------------- 1 | airplane 2 | bathtub 3 | bed 4 | bench 5 | bookshelf 6 | bottle 7 | bowl 8 | car 9 | chair 10 | cone 11 | cup 12 | curtain 13 | desk 14 | door 15 | dresser 16 | flower_pot 17 | glass_box 18 | guitar 19 | keyboard 20 | lamp 21 | laptop 22 | mantel 23 | monitor 24 | night_stand 25 | person 26 | piano 27 | plant 28 | radio 29 | range_hood 30 | sink 31 | sofa 32 | stairs 33 | stool 34 | table 35 | tent 36 | toilet 37 | tv_stand 38 | vase 39 | wardrobe 40 | xbox 41 | -------------------------------------------------------------------------------- /experiments/sampledata/modelnet40_half1.txt: -------------------------------------------------------------------------------- 1 | airplane 2 | bathtub 3 | bed 4 | bench 5 | bookshelf 6 | bottle 7 | bowl 8 | car 9 | chair 10 | cone 11 | cup 12 | curtain 13 | desk 14 | door 15 | dresser 16 | flower_pot 17 | glass_box 18 | guitar 19 | keyboard 20 | lamp 21 | -------------------------------------------------------------------------------- /experiments/sampledata/modelnet40_half2.txt: -------------------------------------------------------------------------------- 1 | laptop 2 | mantel 3 | monitor 4 | night_stand 5 | person 6 | piano 7 | plant 8 | radio 9 | range_hood 10 | sink 11 | sofa 12 | stairs 13 | stool 14 | table 15 | tent 16 | toilet 17 | tv_stand 18 | vase 19 | wardrobe 20 | xbox 21 | -------------------------------------------------------------------------------- /experiments/sampledata/shapenet2_20.txt: -------------------------------------------------------------------------------- 1 | 02691156 2 | 02747177 3 | 02773838 4 | 02801938 5 | 02808440 6 | 02818832 7 | 02828884 8 | 02843684 9 | 02871439 10 | 02876657 11 | 02880940 12 | 02924116 13 | 02933112 14 | 02942699 15 | 02946921 16 | 02954340 17 | 02958343 18 | 02992529 19 | 03001627 20 | 03046257 21 | -------------------------------------------------------------------------------- /experiments/test_icp.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example for testing ICP. 3 | 4 | No-noise version. 5 | 6 | This ICP-test is very slow. use faster ones like Matlab or C++... 7 | """ 8 | 9 | import argparse 10 | import os 11 | import sys 12 | import logging 13 | import numpy 14 | import torch 15 | import torch.utils.data 16 | import torchvision 17 | 18 | # addpath('../') 19 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 20 | import ptlk 21 | 22 | import icp 23 | 24 | LOGGER = logging.getLogger(__name__) 25 | LOGGER.addHandler(logging.NullHandler()) 26 | 27 | 28 | def options(argv=None): 29 | parser = argparse.ArgumentParser(description='ICP') 30 | 31 | # required. 32 | parser.add_argument('-o', '--outfile', required=True, type=str, 33 | metavar='FILENAME', help='output filename (.csv)') 34 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 35 | metavar='PATH', help='path to the input dataset') 36 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 37 | metavar='PATH', help='path to the categories to be tested') 38 | parser.add_argument('-p', '--perturbations', required=True, type=str, 39 | metavar='PATH', help='path to the perturbations') 40 | 41 | # settings for input data 42 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet'], 43 | metavar='DATASET', help='dataset type (default: modelnet)') 44 | parser.add_argument('--format', default='wv', choices=['wv', 'wt'], 45 | help='perturbation format (default: wv (twist)) (wt: rotation and translation)') # the output is always in twist format 46 | 47 | # settings for ICP 48 | parser.add_argument('--max-iter', default=20, type=int, 49 | metavar='N', help='max-iter on ICP. (default: 20)') 50 | 51 | # settings for on testing 52 | parser.add_argument('-l', '--logfile', default='', type=str, 53 | metavar='LOGNAME', help='path to logfile (default: null (no logging))') 54 | parser.add_argument('-j', '--workers', default=4, type=int, 55 | metavar='N', help='number of data loading workers (default: 4)') 56 | 57 | args = parser.parse_args(argv) 58 | return args 59 | 60 | def main(args): 61 | # dataset 62 | testset = get_datasets(args) 63 | 64 | # testing 65 | act = Action(args) 66 | run(args, testset, act) 67 | 68 | 69 | def run(args, testset, action): 70 | LOGGER.debug('Testing (PID=%d), %s', os.getpid(), args) 71 | 72 | sys.setrecursionlimit(20000) 73 | 74 | # dataloader 75 | testloader = torch.utils.data.DataLoader( 76 | testset, 77 | batch_size=1, shuffle=False, num_workers=args.workers) 78 | 79 | # testing 80 | LOGGER.debug('tests, begin') 81 | action.eval_1(testloader) 82 | LOGGER.debug('tests, end') 83 | 84 | 85 | class Action: 86 | def __init__(self, args): 87 | self.filename = args.outfile 88 | # ICP 89 | self.max_iter = args.max_iter 90 | 91 | def eval_1__header(self, fout): 92 | cols = ['h_w1', 'h_w2', 'h_w3', 'h_v1', 'h_v2', 'h_v3', \ 93 | 'g_w1', 'g_w2', 'g_w3', 'g_v1', 'g_v2', 'g_v3'] # h: estimated, g: ground-truth twist vectors 94 | print(','.join(map(str, cols)), file=fout) 95 | fout.flush() 96 | 97 | def eval_1__write(self, fout, ig_gt, g_hat): 98 | x_hat = ptlk.se3.log(g_hat) # --> [-1, 6] 99 | mx_gt = ptlk.se3.log(ig_gt) # --> [-1, 6] 100 | for i in range(x_hat.size(0)): 101 | x_hat1 = x_hat[i] # [6] 102 | mx_gt1 = mx_gt[i] # [6] 103 | vals = torch.cat((x_hat1, -mx_gt1)) # [12] 104 | valn = vals.cpu().numpy().tolist() 105 | print(','.join(map(str, valn)), file=fout) 106 | fout.flush() 107 | 108 | def eval_1(self, testloader): 109 | with open(self.filename, 'w') as fout: 110 | self.eval_1__header(fout) 111 | with torch.no_grad(): 112 | for i, data in enumerate(testloader): 113 | p0, p1, igt = data 114 | res = self.do_estimate(p0, p1) # [1, N, 3] x [1, M, 3] --> [1, 4, 4] 115 | ig_gt = igt.cpu().contiguous().view(-1, 4, 4) # --> [1, 4, 4] 116 | g_hat = res.cpu().contiguous().view(-1, 4, 4) # --> [1, 4, 4] 117 | 118 | dg = g_hat.bmm(ig_gt) # if correct, dg == identity matrix. 119 | dx = ptlk.se3.log(dg) # --> [1, 6] (if corerct, dx == zero vector) 120 | dn = dx.norm(p=2, dim=1) # --> [1] 121 | dm = dn.mean() 122 | 123 | self.eval_1__write(fout, ig_gt, g_hat) 124 | LOGGER.info('test, %d/%d, %f', i, len(testloader), dm) 125 | 126 | 127 | def do_estimate(self, p0, p1): 128 | np_p0 = p0.cpu().contiguous().squeeze(0).numpy() # --> (N, 3) 129 | np_p1 = p1.cpu().contiguous().squeeze(0).numpy() # --> (M, 3) 130 | 131 | mod = icp.ICP(np_p0, np_p1) 132 | g, p, itr = mod.compute(self.max_iter) 133 | 134 | est_g = torch.from_numpy(g).view(-1, 4, 4).to(p0) # (1, 4, 4) 135 | 136 | return est_g 137 | 138 | 139 | def get_datasets(args): 140 | 141 | cinfo = None 142 | if args.categoryfile: 143 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 144 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 145 | categories.sort() 146 | c_to_idx = {categories[i]: i for i in range(len(categories))} 147 | cinfo = (categories, c_to_idx) 148 | 149 | perturbations = None 150 | fmt_trans = False 151 | if args.perturbations: 152 | perturbations = numpy.loadtxt(args.perturbations, delimiter=',') 153 | if args.format == 'wt': 154 | fmt_trans = True 155 | 156 | if args.dataset_type == 'modelnet': 157 | transform = torchvision.transforms.Compose([\ 158 | ptlk.data.transforms.Mesh2Points(),\ 159 | ptlk.data.transforms.OnUnitCube(),\ 160 | ]) 161 | 162 | testdata = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 163 | 164 | testset = ptlk.data.datasets.CADset4tracking_fixed_perturbation(testdata,\ 165 | perturbations, fmt_trans=fmt_trans) 166 | 167 | return testset 168 | 169 | 170 | if __name__ == '__main__': 171 | ARGS = options() 172 | 173 | logging.basicConfig( 174 | level=logging.DEBUG, 175 | format='%(levelname)s:%(name)s, %(asctime)s, %(message)s', 176 | filename=ARGS.logfile) 177 | LOGGER.debug('Testing (PID=%d), %s', os.getpid(), ARGS) 178 | 179 | main(ARGS) 180 | LOGGER.debug('done (PID=%d)', os.getpid()) 181 | 182 | #EOF -------------------------------------------------------------------------------- /experiments/test_icp_vs_pointlk.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/experiments/test_icp_vs_pointlk.py -------------------------------------------------------------------------------- /experiments/test_pointlk.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example for testing PointNet-LK. 3 | 4 | No-noise version. 5 | """ 6 | 7 | import argparse 8 | import os 9 | import sys 10 | import logging 11 | import numpy 12 | import torch 13 | import torch.utils.data 14 | import torchvision 15 | 16 | # addpath('../') 17 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 18 | import ptlk 19 | 20 | 21 | LOGGER = logging.getLogger(__name__) 22 | LOGGER.addHandler(logging.NullHandler()) 23 | 24 | 25 | def options(argv=None): 26 | parser = argparse.ArgumentParser(description='PointNet-LK') 27 | 28 | # required. 29 | parser.add_argument('-o', '--outfile', required=True, type=str, 30 | metavar='FILENAME', help='output filename (.csv)') 31 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 32 | metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40' 33 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 34 | metavar='PATH', help='path to the categories to be tested') # eg. './sampledata/modelnet40_half1.txt' 35 | parser.add_argument('-p', '--perturbations', required=True, type=str, 36 | metavar='PATH', help='path to the perturbation file') # see. generate_perturbations.py 37 | 38 | # settings for input data 39 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet'], 40 | metavar='DATASET', help='dataset type (default: modelnet)') 41 | parser.add_argument('--format', default='wv', choices=['wv', 'wt'], 42 | help='perturbation format (default: wv (twist)) (wt: rotation and translation)') # the output is always in twist format 43 | 44 | # settings for PointNet-LK 45 | parser.add_argument('--max-iter', default=20, type=int, 46 | metavar='N', help='max-iter on LK. (default: 20)') 47 | parser.add_argument('--pretrained', default='', type=str, 48 | metavar='PATH', help='path to trained model file (default: null (no-use))') 49 | parser.add_argument('--transfer-from', default='', type=str, 50 | metavar='PATH', help='path to classifier feature (default: null (no-use))') 51 | parser.add_argument('--dim-k', default=1024, type=int, 52 | metavar='K', help='dim. of the feature vector (default: 1024)') 53 | parser.add_argument('--symfn', default='max', choices=['max', 'avg'], 54 | help='symmetric function (default: max)') 55 | parser.add_argument('--delta', default=1.0e-2, type=float, 56 | metavar='D', help='step size for approx. Jacobian (default: 1.0e-2)') 57 | 58 | # settings for on testing 59 | parser.add_argument('-l', '--logfile', default='', type=str, 60 | metavar='LOGNAME', help='path to logfile (default: null (no logging))') 61 | parser.add_argument('-j', '--workers', default=4, type=int, 62 | metavar='N', help='number of data loading workers (default: 4)') 63 | parser.add_argument('--device', default='cpu', type=str, 64 | metavar='DEVICE', help='use CUDA if available (default: cpu)') 65 | 66 | args = parser.parse_args(argv) 67 | return args 68 | 69 | def main(args): 70 | # dataset 71 | testset = get_datasets(args) 72 | 73 | # testing 74 | act = Action(args) 75 | run(args, testset, act) 76 | 77 | 78 | def run(args, testset, action): 79 | if not torch.cuda.is_available(): 80 | args.device = 'cpu' 81 | args.device = torch.device(args.device) 82 | 83 | LOGGER.debug('Testing (PID=%d), %s', os.getpid(), args) 84 | 85 | model = action.create_model() 86 | if args.pretrained: 87 | assert os.path.isfile(args.pretrained) 88 | model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) 89 | model.to(args.device) 90 | 91 | # dataloader 92 | testloader = torch.utils.data.DataLoader( 93 | testset, 94 | batch_size=1, shuffle=False, num_workers=args.workers) 95 | 96 | # testing 97 | LOGGER.debug('tests, begin') 98 | action.eval_1(model, testloader, args.device) 99 | LOGGER.debug('tests, end') 100 | 101 | 102 | class Action: 103 | def __init__(self, args): 104 | self.filename = args.outfile 105 | # PointNet 106 | self.transfer_from = args.transfer_from 107 | self.dim_k = args.dim_k 108 | self.sym_fn = None 109 | if args.symfn == 'max': 110 | self.sym_fn = ptlk.pointnet.symfn_max 111 | elif args.symfn == 'avg': 112 | self.sym_fn = ptlk.pointnet.symfn_avg 113 | # LK 114 | self.delta = args.delta 115 | self.max_iter = args.max_iter 116 | self.xtol = 1.0e-7 117 | self.p0_zero_mean = True 118 | self.p1_zero_mean = True 119 | 120 | def create_model(self): 121 | ptnet = self.create_pointnet_features() 122 | return self.create_from_pointnet_features(ptnet) 123 | 124 | def create_pointnet_features(self): 125 | ptnet = ptlk.pointnet.PointNet_features(self.dim_k, use_tnet=False, sym_fn=self.sym_fn) 126 | if self.transfer_from and os.path.isfile(self.transfer_from): 127 | ptnet.load_state_dict(torch.load(self.transfer_from, map_location='cpu')) 128 | return ptnet 129 | 130 | def create_from_pointnet_features(self, ptnet): 131 | return ptlk.pointlk.PointLK(ptnet, self.delta) 132 | 133 | def eval_1__header(self, fout): 134 | cols = ['h_w1', 'h_w2', 'h_w3', 'h_v1', 'h_v2', 'h_v3', \ 135 | 'g_w1', 'g_w2', 'g_w3', 'g_v1', 'g_v2', 'g_v3'] # h: estimated, g: ground-truth twist vectors 136 | print(','.join(map(str, cols)), file=fout) 137 | fout.flush() 138 | 139 | def eval_1__write(self, fout, ig_gt, g_hat): 140 | x_hat = ptlk.se3.log(g_hat) # --> [-1, 6] 141 | mx_gt = ptlk.se3.log(ig_gt) # --> [-1, 6] 142 | for i in range(x_hat.size(0)): 143 | x_hat1 = x_hat[i] # [6] 144 | mx_gt1 = mx_gt[i] # [6] 145 | vals = torch.cat((x_hat1, -mx_gt1)) # [12] 146 | valn = vals.cpu().numpy().tolist() 147 | print(','.join(map(str, valn)), file=fout) 148 | fout.flush() 149 | 150 | def eval_1(self, model, testloader, device): 151 | model.eval() 152 | with open(self.filename, 'w') as fout: 153 | self.eval_1__header(fout) 154 | with torch.no_grad(): 155 | for i, data in enumerate(testloader): 156 | p0, p1, igt = data 157 | res = self.do_estimate(p0, p1, model, device) # --> [1, 4, 4] 158 | ig_gt = igt.cpu().contiguous().view(-1, 4, 4) # --> [1, 4, 4] 159 | g_hat = res.cpu().contiguous().view(-1, 4, 4) # --> [1, 4, 4] 160 | 161 | dg = g_hat.bmm(ig_gt) # if correct, dg == identity matrix. 162 | dx = ptlk.se3.log(dg) # --> [1, 6] (if corerct, dx == zero vector) 163 | dn = dx.norm(p=2, dim=1) # --> [1] 164 | dm = dn.mean() 165 | 166 | self.eval_1__write(fout, ig_gt, g_hat) 167 | LOGGER.info('test, %d/%d, %f', i, len(testloader), dm) 168 | 169 | 170 | def do_estimate(self, p0, p1, model, device): 171 | p0 = p0.to(device) # template (1, N, 3) 172 | p1 = p1.to(device) # source (1, M, 3) 173 | r = ptlk.pointlk.PointLK.do_forward(model, p0, p1, self.max_iter, self.xtol,\ 174 | self.p0_zero_mean, self.p1_zero_mean) 175 | #r = model(p0, p1, self.max_iter) 176 | est_g = model.g # (1, 4, 4) 177 | 178 | return est_g 179 | 180 | 181 | def get_datasets(args): 182 | 183 | cinfo = None 184 | if args.categoryfile: 185 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 186 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 187 | categories.sort() 188 | c_to_idx = {categories[i]: i for i in range(len(categories))} 189 | cinfo = (categories, c_to_idx) 190 | 191 | perturbations = None 192 | fmt_trans = False 193 | if args.perturbations: 194 | perturbations = numpy.loadtxt(args.perturbations, delimiter=',') 195 | if args.format == 'wt': 196 | fmt_trans = True 197 | 198 | if args.dataset_type == 'modelnet': 199 | transform = torchvision.transforms.Compose([\ 200 | ptlk.data.transforms.Mesh2Points(),\ 201 | ptlk.data.transforms.OnUnitCube(),\ 202 | ]) 203 | 204 | testdata = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 205 | 206 | testset = ptlk.data.datasets.CADset4tracking_fixed_perturbation(testdata,\ 207 | perturbations, fmt_trans=fmt_trans) 208 | 209 | return testset 210 | 211 | 212 | if __name__ == '__main__': 213 | ARGS = options() 214 | 215 | logging.basicConfig( 216 | level=logging.DEBUG, 217 | format='%(levelname)s:%(name)s, %(asctime)s, %(message)s', 218 | filename=ARGS.logfile) 219 | LOGGER.debug('Testing (PID=%d), %s', os.getpid(), ARGS) 220 | 221 | main(ARGS) 222 | LOGGER.debug('done (PID=%d)', os.getpid()) 223 | 224 | #EOF -------------------------------------------------------------------------------- /experiments/test_utils.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy 3 | import math 4 | import torch 5 | import torch.utils.data 6 | import torchvision 7 | 8 | # addpath('../') 9 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 10 | import ptlk 11 | 12 | 13 | # def rotation(): 14 | # c1 = sympy.Symbol("c1") 15 | # s1 = sympy.Symbol("s1") 16 | # c2 = sympy.Symbol("c2") 17 | # s2 = sympy.Symbol("s2") 18 | # c3 = sympy.Symbol("c3") 19 | # s3 = sympy.Symbol("s3") 20 | # 21 | # Rx = sympy.Matrix([[1, 0, 0], [0, c1, -s1], [0, s1, c1]]) 22 | # Ry = sympy.Matrix([[c2, 0, s2], [0, 1, 0], [-s2, 0, c2]]) 23 | # Rz = sympy.Matrix([[c3, -s3, 0], [s3, c3, 0], [0, 0, 1]]) 24 | # 25 | # print((Rz * Ry * Rx)) 26 | # # Matrix([[c2*c3, -c1*s3 + c3*s1*s2, c1*c3*s2 + s1*s3],\ 27 | # # [c2*s3, c1*c3 + s1*s2*s3, c1*s2*s3 - c3*s1],\ 28 | # # [ -s2, c2*s1, c1*c2]]) 29 | # 30 | # # singular case: c2 == 0, s2 == (-1, +1) 31 | # # --> Rz == identity. 32 | # # Matrix([[ c2, s1*s2, c1*s2], 33 | # # [ 0, c1, -s1], 34 | # # [-s2, 0, 0]]) 35 | 36 | 37 | def rotm2eul(m): 38 | """ m (3x3, rotation matrix) --> rotation m = Rz*Ry*Rx 39 | """ 40 | c = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0]) 41 | 42 | singular = c < 1e-6 43 | if not singular : 44 | x = math.atan2(R[2,1] , R[2,2]) 45 | y = math.atan2(-R[2,0], c) 46 | z = math.atan2(R[1,0], R[0,0]) 47 | else : 48 | x = math.atan2(-R[1,2], R[1,1]) 49 | y = math.atan2(-R[2,0], c) 50 | z = 0 51 | 52 | return numpy.array([x, y, z]) 53 | 54 | 55 | #EOF -------------------------------------------------------------------------------- /experiments/train_classifier.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example for training a classifier. 3 | 4 | It is better that before training PointNet-LK, 5 | train a classifier with same dataset 6 | so that we can use 'transfer-learning.' 7 | """ 8 | 9 | import argparse 10 | import os 11 | import sys 12 | import logging 13 | import numpy 14 | import torch 15 | import torch.utils.data 16 | import torchvision 17 | 18 | # addpath('../') 19 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 20 | import ptlk 21 | 22 | LOGGER = logging.getLogger(__name__) 23 | LOGGER.addHandler(logging.NullHandler()) 24 | 25 | 26 | def options(argv=None): 27 | parser = argparse.ArgumentParser(description='PointNet classifier') 28 | 29 | # required. 30 | parser.add_argument('-o', '--outfile', required=True, type=str, 31 | metavar='BASENAME', help='output filename (prefix)') # result: ${BASENAME}_feat_best.pth 32 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 33 | metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40' 34 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 35 | metavar='PATH', help='path to the categories to be trained') # eg. './sampledata/modelnet40_half1.txt' 36 | 37 | # settings for input data 38 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet', 'shapenet2'], 39 | metavar='DATASET', help='dataset type (default: modelnet)') 40 | parser.add_argument('--num-points', default=1024, type=int, 41 | metavar='N', help='points in point-cloud (default: 1024)') 42 | 43 | # settings for PointNet 44 | parser.add_argument('--use-tnet', dest='use_tnet', action='store_true', 45 | help='flag for setting up PointNet with TNet') 46 | parser.add_argument('--dim-k', default=1024, type=int, 47 | metavar='K', help='dim. of the feature vector (default: 1024)') 48 | parser.add_argument('--symfn', default='max', choices=['max', 'avg'], 49 | help='symmetric function (default: max)') 50 | 51 | # settings for on training 52 | parser.add_argument('-l', '--logfile', default='', type=str, 53 | metavar='LOGNAME', help='path to logfile (default: null (no logging))') 54 | parser.add_argument('-j', '--workers', default=4, type=int, 55 | metavar='N', help='number of data loading workers (default: 4)') 56 | parser.add_argument('-b', '--batch-size', default=32, type=int, 57 | metavar='N', help='mini-batch size (default: 32)') 58 | parser.add_argument('--epochs', default=200, type=int, 59 | metavar='N', help='number of total epochs to run') 60 | parser.add_argument('--start-epoch', default=0, type=int, 61 | metavar='N', help='manual epoch number (useful on restarts)') 62 | parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'], 63 | metavar='METHOD', help='name of an optimizer (default: Adam)') 64 | parser.add_argument('--resume', default='', type=str, 65 | metavar='PATH', help='path to latest checkpoint (default: null (no-use))') 66 | parser.add_argument('--pretrained', default='', type=str, 67 | metavar='PATH', help='path to pretrained model file (default: null (no-use))') 68 | parser.add_argument('--device', default='cuda:0', type=str, 69 | metavar='DEVICE', help='use CUDA if available') 70 | 71 | args = parser.parse_args(argv) 72 | return args 73 | 74 | def main(args): 75 | # dataset 76 | trainset, testset = get_datasets(args) 77 | num_classes = len(trainset.classes) 78 | 79 | # training 80 | act = Action(args, num_classes) 81 | run(args, trainset, testset, act) 82 | 83 | 84 | def run(args, trainset, testset, action): 85 | if not torch.cuda.is_available(): 86 | args.device = 'cpu' 87 | args.device = torch.device(args.device) 88 | 89 | LOGGER.debug('Trainer (PID=%d), %s', os.getpid(), args) 90 | 91 | model = action.create_model() 92 | if args.pretrained: 93 | assert os.path.isfile(args.pretrained) 94 | model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) 95 | model.to(args.device) 96 | 97 | checkpoint = None 98 | if args.resume: 99 | assert os.path.isfile(args.resume) 100 | checkpoint = torch.load(args.resume) 101 | args.start_epoch = checkpoint['epoch'] 102 | model.load_state_dict(checkpoint['model']) 103 | 104 | # dataloader 105 | testloader = torch.utils.data.DataLoader( 106 | testset, 107 | batch_size=args.batch_size, shuffle=False, num_workers=args.workers) 108 | trainloader = torch.utils.data.DataLoader( 109 | trainset, 110 | batch_size=args.batch_size, shuffle=True, num_workers=args.workers) 111 | 112 | # optimizer 113 | min_loss = float('inf') 114 | learnable_params = filter(lambda p: p.requires_grad, model.parameters()) 115 | if args.optimizer == 'Adam': 116 | optimizer = torch.optim.Adam(learnable_params) 117 | else: 118 | optimizer = torch.optim.SGD(learnable_params, lr=0.1) 119 | 120 | if checkpoint is not None: 121 | min_loss = checkpoint['min_loss'] 122 | optimizer.load_state_dict(checkpoint['optimizer']) 123 | 124 | # training 125 | LOGGER.debug('train, begin') 126 | for epoch in range(args.start_epoch, args.epochs): 127 | #scheduler.step() 128 | 129 | running_loss, running_info = action.train_1(model, trainloader, optimizer, args.device) 130 | val_loss, val_info = action.eval_1(model, testloader, args.device) 131 | 132 | is_best = val_loss < min_loss 133 | min_loss = min(val_loss, min_loss) 134 | 135 | LOGGER.info('epoch, %04d, %f, %f, %f, %f', epoch + 1, running_loss, val_loss, running_info, val_info) 136 | snap = {'epoch': epoch + 1, 137 | 'model': model.state_dict(), 138 | 'min_loss': min_loss, 139 | 'optimizer' : optimizer.state_dict(),} 140 | if is_best: 141 | save_checkpoint(snap, args.outfile, 'snap_best') 142 | save_checkpoint(model.state_dict(), args.outfile, 'model_best') 143 | save_checkpoint(model.features.state_dict(), args.outfile, 'feat_best') 144 | 145 | save_checkpoint(snap, args.outfile, 'snap_last') 146 | save_checkpoint(model.state_dict(), args.outfile, 'model_last') 147 | save_checkpoint(model.features.state_dict(), args.outfile, 'feat_last') 148 | 149 | LOGGER.debug('train, end') 150 | 151 | def save_checkpoint(state, filename, suffix): 152 | torch.save(state, '{}_{}.pth'.format(filename, suffix)) 153 | 154 | 155 | class Action: 156 | def __init__(self, args, num_classes): 157 | self.num_classes = num_classes 158 | self.dim_k = args.dim_k 159 | self.use_tnet = args.use_tnet 160 | self.sym_fn = None 161 | if args.symfn == 'max': 162 | self.sym_fn = ptlk.pointnet.symfn_max 163 | elif args.symfn == 'avg': 164 | self.sym_fn = ptlk.pointnet.symfn_avg 165 | 166 | def create_model(self): 167 | feat = ptlk.pointnet.PointNet_features(self.dim_k, self.use_tnet, self.sym_fn) 168 | return ptlk.pointnet.PointNet_classifier(self.num_classes, feat, self.dim_k) 169 | 170 | def train_1(self, model, trainloader, optimizer, device): 171 | model.train() 172 | vloss = 0.0 173 | pred = 0.0 174 | count = 0 175 | for i, data in enumerate(trainloader): 176 | target, output, loss = self.compute_loss(model, data, device) 177 | # forward + backward + optimize 178 | optimizer.zero_grad() 179 | loss.backward() 180 | optimizer.step() 181 | 182 | loss1 = loss.item() 183 | vloss += loss1 184 | count += output.size(0) 185 | 186 | _, pred1 = output.max(dim=1) 187 | ag = (pred1 == target) 188 | am = ag.sum() 189 | pred += am.item() 190 | 191 | running_loss = float(vloss)/count 192 | accuracy = float(pred)/count 193 | return running_loss, accuracy 194 | 195 | def eval_1(self, model, testloader, device): 196 | model.eval() 197 | vloss = 0.0 198 | pred = 0.0 199 | count = 0 200 | with torch.no_grad(): 201 | for i, data in enumerate(testloader): 202 | target, output, loss = self.compute_loss(model, data, device) 203 | 204 | loss1 = loss.item() 205 | vloss += loss1 206 | count += output.size(0) 207 | 208 | _, pred1 = output.max(dim=1) 209 | ag = (pred1 == target) 210 | am = ag.sum() 211 | pred += am.item() 212 | 213 | ave_loss = float(vloss)/count 214 | accuracy = float(pred)/count 215 | return ave_loss, accuracy 216 | 217 | def compute_loss(self, model, data, device): 218 | points, target = data 219 | 220 | points = points.to(device) 221 | target = target.to(device) 222 | 223 | output = model(points) 224 | loss = model.loss(output, target) 225 | 226 | return target, output, loss 227 | 228 | 229 | class ShapeNet2_transform_coordinate: 230 | def __init__(self): 231 | pass 232 | def __call__(self, mesh): 233 | return mesh.clone().rot_x() 234 | 235 | def get_datasets(args): 236 | 237 | cinfo = None 238 | if args.categoryfile: 239 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 240 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 241 | categories.sort() 242 | c_to_idx = {categories[i]: i for i in range(len(categories))} 243 | cinfo = (categories, c_to_idx) 244 | 245 | if args.dataset_type == 'modelnet': 246 | transform = torchvision.transforms.Compose([\ 247 | ptlk.data.transforms.Mesh2Points(),\ 248 | ptlk.data.transforms.OnUnitCube(),\ 249 | ptlk.data.transforms.Resampler(args.num_points),\ 250 | ptlk.data.transforms.RandomRotatorZ(),\ 251 | ptlk.data.transforms.RandomJitter()\ 252 | ]) 253 | 254 | trainset = ptlk.data.datasets.ModelNet(args.dataset_path, train=1, transform=transform, classinfo=cinfo) 255 | testset = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 256 | 257 | elif args.dataset_type == 'shapenet2': 258 | transform = torchvision.transforms.Compose([\ 259 | ShapeNet2_transform_coordinate(),\ 260 | ptlk.data.transforms.Mesh2Points(),\ 261 | ptlk.data.transforms.OnUnitCube(),\ 262 | ptlk.data.transforms.Resampler(args.num_points),\ 263 | ptlk.data.transforms.RandomRotatorZ(),\ 264 | ptlk.data.transforms.RandomJitter()\ 265 | ]) 266 | 267 | dataset = ptlk.data.datasets.ShapeNet2(args.dataset_path, transform=transform, classinfo=cinfo) 268 | trainset, testset = dataset.split(0.8) 269 | 270 | return trainset, testset 271 | 272 | 273 | if __name__ == '__main__': 274 | ARGS = options() 275 | 276 | logging.basicConfig( 277 | level=logging.DEBUG, 278 | format='%(levelname)s:%(name)s, %(asctime)s, %(message)s', 279 | filename=ARGS.logfile) 280 | LOGGER.debug('Training (PID=%d), %s', os.getpid(), ARGS) 281 | 282 | main(ARGS) 283 | LOGGER.debug('done (PID=%d)', os.getpid()) 284 | 285 | #EOF -------------------------------------------------------------------------------- /experiments/train_pointlk.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example for training a tracker (PointNet-LK). 3 | 4 | No-noise version. 5 | """ 6 | 7 | import argparse 8 | import os 9 | import sys 10 | import logging 11 | import numpy 12 | import torch 13 | import torch.utils.data 14 | import torchvision 15 | 16 | # addpath('../') 17 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) 18 | import ptlk 19 | 20 | LOGGER = logging.getLogger(__name__) 21 | LOGGER.addHandler(logging.NullHandler()) 22 | 23 | 24 | def options(argv=None): 25 | parser = argparse.ArgumentParser(description='PointNet-LK') 26 | 27 | # required. 28 | parser.add_argument('-o', '--outfile', required=True, type=str, 29 | metavar='BASENAME', help='output filename (prefix)') # the result: ${BASENAME}_model_best.pth 30 | parser.add_argument('-i', '--dataset-path', required=True, type=str, 31 | metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40' 32 | parser.add_argument('-c', '--categoryfile', required=True, type=str, 33 | metavar='PATH', help='path to the categories to be trained') # eg. './sampledata/modelnet40_half1.txt' 34 | 35 | # settings for input data 36 | parser.add_argument('--dataset-type', default='modelnet', choices=['modelnet', 'shapenet2'], 37 | metavar='DATASET', help='dataset type (default: modelnet)') 38 | parser.add_argument('--num-points', default=1024, type=int, 39 | metavar='N', help='points in point-cloud (default: 1024)') 40 | parser.add_argument('--mag', default=0.8, type=float, 41 | metavar='T', help='max. mag. of twist-vectors (perturbations) on training (default: 0.8)') 42 | 43 | # settings for PointNet 44 | parser.add_argument('--pointnet', default='tune', type=str, choices=['fixed', 'tune'], 45 | help='train pointnet (default: tune)') 46 | parser.add_argument('--transfer-from', default='', type=str, 47 | metavar='PATH', help='path to pointnet features file') 48 | parser.add_argument('--dim-k', default=1024, type=int, 49 | metavar='K', help='dim. of the feature vector (default: 1024)') 50 | parser.add_argument('--symfn', default='max', choices=['max', 'avg'], 51 | help='symmetric function (default: max)') 52 | 53 | # settings for LK 54 | parser.add_argument('--max-iter', default=10, type=int, 55 | metavar='N', help='max-iter on LK. (default: 10)') 56 | parser.add_argument('--delta', default=1.0e-2, type=float, 57 | metavar='D', help='step size for approx. Jacobian (default: 1.0e-2)') 58 | parser.add_argument('--learn-delta', dest='learn_delta', action='store_true', 59 | help='flag for training step size delta') 60 | 61 | # settings for on training 62 | parser.add_argument('-l', '--logfile', default='', type=str, 63 | metavar='LOGNAME', help='path to logfile (default: null (no logging))') 64 | parser.add_argument('-j', '--workers', default=4, type=int, 65 | metavar='N', help='number of data loading workers (default: 4)') 66 | parser.add_argument('-b', '--batch-size', default=32, type=int, 67 | metavar='N', help='mini-batch size (default: 32)') 68 | parser.add_argument('--epochs', default=200, type=int, 69 | metavar='N', help='number of total epochs to run') 70 | parser.add_argument('--start-epoch', default=0, type=int, 71 | metavar='N', help='manual epoch number (useful on restarts)') 72 | parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'], 73 | metavar='METHOD', help='name of an optimizer (default: Adam)') 74 | parser.add_argument('--resume', default='', type=str, 75 | metavar='PATH', help='path to latest checkpoint (default: null (no-use))') 76 | parser.add_argument('--pretrained', default='', type=str, 77 | metavar='PATH', help='path to pretrained model file (default: null (no-use))') 78 | parser.add_argument('--device', default='cuda:0', type=str, 79 | metavar='DEVICE', help='use CUDA if available') 80 | 81 | args = parser.parse_args(argv) 82 | return args 83 | 84 | def main(args): 85 | # dataset 86 | trainset, testset = get_datasets(args) 87 | 88 | # training 89 | act = Action(args) 90 | run(args, trainset, testset, act) 91 | 92 | 93 | def run(args, trainset, testset, action): 94 | if not torch.cuda.is_available(): 95 | args.device = 'cpu' 96 | args.device = torch.device(args.device) 97 | 98 | LOGGER.debug('Trainer (PID=%d), %s', os.getpid(), args) 99 | 100 | model = action.create_model() 101 | if args.pretrained: 102 | assert os.path.isfile(args.pretrained) 103 | model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) 104 | model.to(args.device) 105 | 106 | checkpoint = None 107 | if args.resume: 108 | assert os.path.isfile(args.resume) 109 | checkpoint = torch.load(args.resume) 110 | args.start_epoch = checkpoint['epoch'] 111 | model.load_state_dict(checkpoint['model']) 112 | 113 | # dataloader 114 | testloader = torch.utils.data.DataLoader( 115 | testset, 116 | batch_size=args.batch_size, shuffle=False, num_workers=args.workers) 117 | trainloader = torch.utils.data.DataLoader( 118 | trainset, 119 | batch_size=args.batch_size, shuffle=True, num_workers=args.workers) 120 | 121 | # optimizer 122 | min_loss = float('inf') 123 | learnable_params = filter(lambda p: p.requires_grad, model.parameters()) 124 | if args.optimizer == 'Adam': 125 | optimizer = torch.optim.Adam(learnable_params) 126 | else: 127 | optimizer = torch.optim.SGD(learnable_params, lr=0.1) 128 | 129 | if checkpoint is not None: 130 | min_loss = checkpoint['min_loss'] 131 | optimizer.load_state_dict(checkpoint['optimizer']) 132 | 133 | # training 134 | LOGGER.debug('train, begin') 135 | for epoch in range(args.start_epoch, args.epochs): 136 | #scheduler.step() 137 | 138 | running_loss, running_info = action.train_1(model, trainloader, optimizer, args.device) 139 | val_loss, val_info = action.eval_1(model, testloader, args.device) 140 | 141 | is_best = val_loss < min_loss 142 | min_loss = min(val_loss, min_loss) 143 | 144 | LOGGER.info('epoch, %04d, %f, %f, %f, %f', epoch + 1, running_loss, val_loss, running_info, val_info) 145 | snap = {'epoch': epoch + 1, 146 | 'model': model.state_dict(), 147 | 'min_loss': min_loss, 148 | 'optimizer' : optimizer.state_dict(),} 149 | if is_best: 150 | save_checkpoint(snap, args.outfile, 'snap_best') 151 | save_checkpoint(model.state_dict(), args.outfile, 'model_best') 152 | 153 | save_checkpoint(snap, args.outfile, 'snap_last') 154 | save_checkpoint(model.state_dict(), args.outfile, 'model_last') 155 | 156 | LOGGER.debug('train, end') 157 | 158 | def save_checkpoint(state, filename, suffix): 159 | torch.save(state, '{}_{}.pth'.format(filename, suffix)) 160 | 161 | 162 | class Action: 163 | def __init__(self, args): 164 | # PointNet 165 | self.pointnet = args.pointnet # tune or fixed 166 | self.transfer_from = args.transfer_from 167 | self.dim_k = args.dim_k 168 | self.sym_fn = None 169 | if args.symfn == 'max': 170 | self.sym_fn = ptlk.pointnet.symfn_max 171 | elif args.symfn == 'avg': 172 | self.sym_fn = ptlk.pointnet.symfn_avg 173 | # LK 174 | self.delta = args.delta 175 | self.learn_delta = args.learn_delta 176 | self.max_iter = args.max_iter 177 | self.xtol = 1.0e-7 178 | self.p0_zero_mean = True 179 | self.p1_zero_mean = True 180 | 181 | self._loss_type = 1 # see. self.compute_loss() 182 | 183 | def create_model(self): 184 | ptnet = self.create_pointnet_features() 185 | return self.create_from_pointnet_features(ptnet) 186 | 187 | def create_pointnet_features(self): 188 | ptnet = ptlk.pointnet.PointNet_features(self.dim_k, use_tnet=False, sym_fn=self.sym_fn) 189 | if self.transfer_from and os.path.isfile(self.transfer_from): 190 | ptnet.load_state_dict(torch.load(self.transfer_from, map_location='cpu')) 191 | if self.pointnet == 'tune': 192 | pass 193 | elif self.pointnet == 'fixed': 194 | for param in ptnet.parameters(): 195 | param.requires_grad_(False) 196 | return ptnet 197 | 198 | def create_from_pointnet_features(self, ptnet): 199 | return ptlk.pointlk.PointLK(ptnet, self.delta, self.learn_delta) 200 | 201 | def train_1(self, model, trainloader, optimizer, device): 202 | model.train() 203 | vloss = 0.0 204 | gloss = 0.0 205 | count = 0 206 | for i, data in enumerate(trainloader): 207 | loss, loss_g = self.compute_loss(model, data, device) 208 | 209 | # forward + backward + optimize 210 | optimizer.zero_grad() 211 | loss.backward() 212 | optimizer.step() 213 | 214 | vloss1 = loss.item() 215 | vloss += vloss1 216 | gloss1 = loss_g.item() 217 | gloss += gloss1 218 | count += 1 219 | 220 | ave_vloss = float(vloss)/count 221 | ave_gloss = float(gloss)/count 222 | return ave_vloss, ave_gloss 223 | 224 | def eval_1(self, model, testloader, device): 225 | model.eval() 226 | vloss = 0.0 227 | gloss = 0.0 228 | count = 0 229 | with torch.no_grad(): 230 | for i, data in enumerate(testloader): 231 | loss, loss_g = self.compute_loss(model, data, device) 232 | 233 | vloss1 = loss.item() 234 | vloss += vloss1 235 | gloss1 = loss_g.item() 236 | gloss += gloss1 237 | count += 1 238 | 239 | ave_vloss = float(vloss)/count 240 | ave_gloss = float(gloss)/count 241 | return ave_vloss, ave_gloss 242 | 243 | def compute_loss(self, model, data, device): 244 | p0, p1, igt = data 245 | p0 = p0.to(device) # template 246 | p1 = p1.to(device) # source 247 | igt = igt.to(device) # igt: p0 -> p1 248 | r = ptlk.pointlk.PointLK.do_forward(model, p0, p1, self.max_iter, self.xtol,\ 249 | self.p0_zero_mean, self.p1_zero_mean) 250 | #r = model(p0, p1, self.max_iter) 251 | est_g = model.g 252 | 253 | loss_g = ptlk.pointlk.PointLK.comp(est_g, igt) 254 | 255 | if self._loss_type == 0: 256 | loss_r = ptlk.pointlk.PointLK.rsq(r) 257 | loss = loss_r 258 | elif self._loss_type == 1: 259 | loss_r = ptlk.pointlk.PointLK.rsq(r) 260 | loss = loss_r + loss_g 261 | elif self._loss_type == 2: 262 | pr = model.prev_r 263 | if pr is not None: 264 | loss_r = ptlk.pointlk.PointLK.rsq(r - pr) 265 | else: 266 | loss_r = ptlk.pointlk.PointLK.rsq(r) 267 | loss = loss_r + loss_g 268 | else: 269 | loss = loss_g 270 | 271 | return loss, loss_g 272 | 273 | 274 | class ShapeNet2_transform_coordinate: 275 | def __init__(self): 276 | pass 277 | def __call__(self, mesh): 278 | return mesh.clone().rot_x() 279 | 280 | def get_datasets(args): 281 | 282 | cinfo = None 283 | if args.categoryfile: 284 | #categories = numpy.loadtxt(args.categoryfile, dtype=str, delimiter="\n").tolist() 285 | categories = [line.rstrip('\n') for line in open(args.categoryfile)] 286 | categories.sort() 287 | c_to_idx = {categories[i]: i for i in range(len(categories))} 288 | cinfo = (categories, c_to_idx) 289 | 290 | if args.dataset_type == 'modelnet': 291 | transform = torchvision.transforms.Compose([\ 292 | ptlk.data.transforms.Mesh2Points(),\ 293 | ptlk.data.transforms.OnUnitCube(),\ 294 | ptlk.data.transforms.Resampler(args.num_points),\ 295 | ]) 296 | 297 | traindata = ptlk.data.datasets.ModelNet(args.dataset_path, train=1, transform=transform, classinfo=cinfo) 298 | testdata = ptlk.data.datasets.ModelNet(args.dataset_path, train=0, transform=transform, classinfo=cinfo) 299 | 300 | mag_randomly = True 301 | trainset = ptlk.data.datasets.CADset4tracking(traindata,\ 302 | ptlk.data.transforms.RandomTransformSE3(args.mag, mag_randomly)) 303 | testset = ptlk.data.datasets.CADset4tracking(testdata,\ 304 | ptlk.data.transforms.RandomTransformSE3(args.mag, mag_randomly)) 305 | 306 | elif args.dataset_type == 'shapenet2': 307 | transform = torchvision.transforms.Compose([\ 308 | ShapeNet2_transform_coordinate(),\ 309 | ptlk.data.transforms.Mesh2Points(),\ 310 | ptlk.data.transforms.OnUnitCube(),\ 311 | ptlk.data.transforms.Resampler(args.num_points),\ 312 | ]) 313 | 314 | dataset = ptlk.data.datasets.ShapeNet2(args.dataset_path, transform=transform, classinfo=cinfo) 315 | traindata, testdata = dataset.split(0.8) 316 | 317 | mag_randomly = True 318 | trainset = ptlk.data.datasets.CADset4tracking(traindata,\ 319 | ptlk.data.transforms.RandomTransformSE3(args.mag, mag_randomly)) 320 | testset = ptlk.data.datasets.CADset4tracking(testdata,\ 321 | ptlk.data.transforms.RandomTransformSE3(args.mag, mag_randomly)) 322 | 323 | 324 | return trainset, testset 325 | 326 | 327 | if __name__ == '__main__': 328 | ARGS = options() 329 | 330 | logging.basicConfig( 331 | level=logging.DEBUG, 332 | format='%(levelname)s:%(name)s, %(asctime)s, %(message)s', 333 | filename=ARGS.logfile) 334 | LOGGER.debug('Training (PID=%d), %s', os.getpid(), ARGS) 335 | 336 | main(ARGS) 337 | LOGGER.debug('done (PID=%d)', os.getpid()) 338 | 339 | #EOF -------------------------------------------------------------------------------- /ptlk/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from . import sinc, so3, se3 3 | from . import pointnet, invmat, pointlk 4 | from . import data 5 | 6 | #EOF -------------------------------------------------------------------------------- /ptlk/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/invmat.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/invmat.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/pointlk.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/pointlk.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/pointnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/pointnet.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/se3.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/se3.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/sinc.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/sinc.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/__pycache__/so3.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/__pycache__/so3.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from . import mesh, transforms 3 | from . import globset, datasets 4 | 5 | #EOF -------------------------------------------------------------------------------- /ptlk/data/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/data/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/__pycache__/datasets.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/data/__pycache__/datasets.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/__pycache__/globset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/data/__pycache__/globset.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/__pycache__/mesh.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/data/__pycache__/mesh.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/__pycache__/transforms.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmgoforth/PointNetLK/fab02ef0fcd48cb7838c50a61b4d183a80eab65b/ptlk/data/__pycache__/transforms.cpython-36.pyc -------------------------------------------------------------------------------- /ptlk/data/datasets.py: -------------------------------------------------------------------------------- 1 | """ datasets """ 2 | 3 | import numpy 4 | import torch 5 | import torch.utils.data 6 | 7 | from . import globset 8 | from . import mesh 9 | from .. import so3 10 | from .. import se3 11 | 12 | 13 | class ModelNet(globset.Globset): 14 | """ [Princeton ModelNet](http://modelnet.cs.princeton.edu/) """ 15 | def __init__(self, dataset_path, train=1, transform=None, classinfo=None): 16 | loader = mesh.offread 17 | if train > 0: 18 | pattern = 'train/*.off' 19 | elif train == 0: 20 | pattern = 'test/*.off' 21 | else: 22 | pattern = ['train/*.off', 'test/*.off'] 23 | super().__init__(dataset_path, pattern, loader, transform, classinfo) 24 | 25 | class ShapeNet2(globset.Globset): 26 | """ [ShapeNet](https://www.shapenet.org/) v2 """ 27 | def __init__(self, dataset_path, transform=None, classinfo=None): 28 | loader = mesh.objread 29 | pattern = '*/models/model_normalized.obj' 30 | super().__init__(dataset_path, pattern, loader, transform, classinfo) 31 | 32 | 33 | class CADset4tracking(torch.utils.data.Dataset): 34 | def __init__(self, dataset, rigid_transform, source_modifier=None, template_modifier=None): 35 | self.dataset = dataset 36 | self.rigid_transform = rigid_transform 37 | self.source_modifier = source_modifier 38 | self.template_modifier = template_modifier 39 | 40 | def __len__(self): 41 | return len(self.dataset) 42 | 43 | def __getitem__(self, index): 44 | pm, _ = self.dataset[index] 45 | if self.source_modifier is not None: 46 | p_ = self.source_modifier(pm) 47 | p1 = self.rigid_transform(p_) 48 | else: 49 | p1 = self.rigid_transform(pm) 50 | igt = self.rigid_transform.igt 51 | 52 | if self.template_modifier is not None: 53 | p0 = self.template_modifier(pm) 54 | else: 55 | p0 = pm 56 | 57 | # p0: template, p1: source, igt: transform matrix from p0 to p1 58 | return p0, p1, igt 59 | 60 | 61 | class CADset4tracking_fixed_perturbation(torch.utils.data.Dataset): 62 | @staticmethod 63 | def generate_perturbations(batch_size, mag, randomly=False): 64 | if randomly: 65 | amp = torch.rand(batch_size, 1) * mag 66 | else: 67 | amp = mag 68 | x = torch.randn(batch_size, 6) 69 | x = x / x.norm(p=2, dim=1, keepdim=True) * amp 70 | return x.numpy() 71 | 72 | @staticmethod 73 | def generate_rotations(batch_size, mag, randomly=False): 74 | if randomly: 75 | amp = torch.rand(batch_size, 1) * mag 76 | else: 77 | amp = mag 78 | w = torch.randn(batch_size, 3) 79 | w = w / w.norm(p=2, dim=1, keepdim=True) * amp 80 | v = torch.zeros(batch_size, 3) 81 | x = torch.cat((w, v), dim=1) 82 | return x.numpy() 83 | 84 | def __init__(self, dataset, perturbation, source_modifier=None, template_modifier=None, 85 | fmt_trans=False): 86 | self.dataset = dataset 87 | self.perturbation = numpy.array(perturbation) # twist (len(dataset), 6) 88 | self.source_modifier = source_modifier 89 | self.template_modifier = template_modifier 90 | self.fmt_trans = fmt_trans # twist or (rotation and translation) 91 | 92 | def do_transform(self, p0, x): 93 | # p0: [N, 3] 94 | # x: [1, 6] 95 | if not self.fmt_trans: 96 | # x: twist-vector 97 | g = se3.exp(x).to(p0) # [1, 4, 4] 98 | p1 = se3.transform(g, p0) 99 | igt = g.squeeze(0) # igt: p0 -> p1 100 | else: 101 | # x: rotation and translation 102 | w = x[:, 0:3] 103 | q = x[:, 3:6] 104 | R = so3.exp(w).to(p0) # [1, 3, 3] 105 | g = torch.zeros(1, 4, 4) 106 | g[:, 3, 3] = 1 107 | g[:, 0:3, 0:3] = R # rotation 108 | g[:, 0:3, 3] = q # translation 109 | p1 = se3.transform(g, p0) 110 | igt = g.squeeze(0) # igt: p0 -> p1 111 | return p1, igt 112 | 113 | def __len__(self): 114 | return len(self.dataset) 115 | 116 | def __getitem__(self, index): 117 | twist = torch.from_numpy(numpy.array(self.perturbation[index])).contiguous().view(1, 6) 118 | pm, _ = self.dataset[index] 119 | x = twist.to(pm) 120 | if self.source_modifier is not None: 121 | p_ = self.source_modifier(pm) 122 | p1, igt = self.do_transform(p_, x) 123 | else: 124 | p1, igt = self.do_transform(pm, x) 125 | 126 | if self.template_modifier is not None: 127 | p0 = self.template_modifier(pm) 128 | else: 129 | p0 = pm 130 | 131 | # p0: template, p1: source, igt: transform matrix from p0 to p1 132 | return p0, p1, igt 133 | 134 | 135 | 136 | #EOF 137 | -------------------------------------------------------------------------------- /ptlk/data/globset.py: -------------------------------------------------------------------------------- 1 | """ glob. """ 2 | import os 3 | import glob 4 | import copy 5 | import six 6 | import numpy as np 7 | import torch 8 | import torch.utils.data 9 | 10 | 11 | def find_classes(root): 12 | """ find ${root}/${class}/* """ 13 | classes = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))] 14 | classes.sort() 15 | class_to_idx = {classes[i]: i for i in range(len(classes))} 16 | return classes, class_to_idx 17 | 18 | def classes_to_cinfo(classes): 19 | class_to_idx = {classes[i]: i for i in range(len(classes))} 20 | return classes, class_to_idx 21 | 22 | def glob_dataset(root, class_to_idx, ptns): 23 | """ glob ${root}/${class}/${ptns[i]} """ 24 | root = os.path.expanduser(root) 25 | samples = [] 26 | #class_size = [0 for i in range(len(class_to_idx))] 27 | for target in sorted(os.listdir(root)): 28 | d = os.path.join(root, target) 29 | if not os.path.isdir(d): 30 | continue 31 | 32 | target_idx = class_to_idx.get(target) 33 | if target_idx is None: 34 | continue 35 | 36 | #count = 0 37 | for i, ptn in enumerate(ptns): 38 | gptn = os.path.join(d, ptn) 39 | names = glob.glob(gptn) 40 | for path in sorted(names): 41 | item = (path, target_idx) 42 | samples.append(item) 43 | #count += 1 44 | #class_size[target_idx] = count 45 | 46 | return samples 47 | 48 | 49 | class Globset(torch.utils.data.Dataset): 50 | """ glob ${rootdir}/${classes}/${pattern} 51 | """ 52 | def __init__(self, rootdir, pattern, fileloader, transform=None, classinfo=None): 53 | super().__init__() 54 | 55 | if isinstance(pattern, six.string_types): 56 | pattern = [pattern] 57 | 58 | if classinfo is not None: 59 | classes, class_to_idx = classinfo 60 | else: 61 | classes, class_to_idx = find_classes(rootdir) 62 | 63 | samples = glob_dataset(rootdir, class_to_idx, pattern) 64 | if not samples: 65 | raise RuntimeError("Empty: rootdir={}, pattern(s)={}".format(rootdir, pattern)) 66 | 67 | self.rootdir = rootdir 68 | self.pattern = pattern 69 | self.fileloader = fileloader 70 | self.transform = transform 71 | 72 | self.classes = classes 73 | self.class_to_idx = class_to_idx 74 | self.samples = samples 75 | 76 | def __repr__(self): 77 | fmt_str = 'Dataset {}\n'.format(self.__class__.__name__) 78 | fmt_str += ' Number of datapoints: {}\n'.format(self.__len__()) 79 | fmt_str += ' Root Location: {}\n'.format(self.rootdir) 80 | fmt_str += ' File Patterns: {}\n'.format(self.pattern) 81 | fmt_str += ' File Loader: {}\n'.format(self.fileloader) 82 | tmp = ' Transforms (if any): ' 83 | fmt_str += '{0}{1}\n'.format(tmp, 84 | self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 85 | return fmt_str 86 | 87 | def __len__(self): 88 | return len(self.samples) 89 | 90 | def __getitem__(self, index): 91 | path, target = self.samples[index] 92 | sample = self.fileloader(path) 93 | if self.transform is not None: 94 | sample = self.transform(sample) 95 | 96 | return sample, target 97 | 98 | def num_classes(self): 99 | return len(self.classes) 100 | 101 | def class_name(self, cidx): 102 | return self.classes[cidx] 103 | 104 | def indices_in_class(self, cidx): 105 | targets = np.array(list(map(lambda s: s[1], self.samples))) 106 | return np.where(targets == cidx).tolist() 107 | 108 | def select_classes(self, cidxs): 109 | indices = [] 110 | for i in cidxs: 111 | idxs = self.indices_in_class(i) 112 | indices.extend(idxs) 113 | return indices 114 | 115 | def split(self, rate): 116 | """ dateset -> dataset1, dataset2. s.t. 117 | len(dataset1) = rate * len(dataset), 118 | len(dataset2) = (1-rate) * len(dataset) 119 | """ 120 | orig_size = len(self) 121 | select = np.zeros(orig_size, dtype=int) 122 | csize = np.zeros(len(self.classes), dtype=int) 123 | dsize = np.zeros(len(self.classes), dtype=int) 124 | 125 | for i in range(orig_size): 126 | _, target = self.samples[i] 127 | csize[target] += 1 128 | dsize = (csize * rate).astype(int) 129 | for i in range(orig_size): 130 | _, target = self.samples[i] 131 | if dsize[target] > 0: 132 | select[i] = 1 133 | dsize[target] -= 1 134 | 135 | dataset1 = copy.deepcopy(self) 136 | dataset2 = copy.deepcopy(self) 137 | 138 | samples1 = list(map(lambda i: dataset1.samples[i], np.where(select == 1)[0])) 139 | samples2 = list(map(lambda i: dataset2.samples[i], np.where(select == 0)[0])) 140 | 141 | dataset1.samples = samples1 142 | dataset2.samples = samples2 143 | return dataset1, dataset2 144 | 145 | 146 | 147 | #EOF 148 | -------------------------------------------------------------------------------- /ptlk/data/mesh.py: -------------------------------------------------------------------------------- 1 | """ 3-d mesh reader """ 2 | import os 3 | import copy 4 | import numpy 5 | from mpl_toolkits.mplot3d import Axes3D 6 | from mpl_toolkits.mplot3d.art3d import Poly3DCollection 7 | import matplotlib.pyplot 8 | 9 | 10 | class Mesh: 11 | def __init__(self): 12 | self._vertices = [] # array-like (N, D) 13 | self._faces = [] # array-like (M, K) 14 | self._edges = [] # array-like (L, 2) 15 | 16 | def clone(self): 17 | other = copy.deepcopy(self) 18 | return other 19 | 20 | def clear(self): 21 | for key in self.__dict__: 22 | self.__dict__[key] = [] 23 | 24 | def add_attr(self, name): 25 | self.__dict__[name] = [] 26 | 27 | @property 28 | def vertex_array(self): 29 | return numpy.array(self._vertices) 30 | 31 | @property 32 | def vertex_list(self): 33 | return list(map(tuple, self._vertices)) 34 | 35 | @staticmethod 36 | def faces2polygons(faces, vertices): 37 | p = list(map(lambda face: \ 38 | list(map(lambda vidx: vertices[vidx], face)), faces)) 39 | return p 40 | 41 | @property 42 | def polygon_list(self): 43 | p = Mesh.faces2polygons(self._faces, self._vertices) 44 | return p 45 | 46 | def plot(self, fig=None, ax=None, *args, **kwargs): 47 | p = self.polygon_list 48 | v = self.vertex_array 49 | if fig is None: 50 | fig = matplotlib.pyplot.gcf() 51 | if ax is None: 52 | ax = Axes3D(fig) 53 | if p: 54 | ax.add_collection3d(Poly3DCollection(p)) 55 | if v.shape: 56 | ax.scatter(v[:, 0], v[:, 1], v[:, 2], *args, **kwargs) 57 | ax.set_xlabel('X') 58 | ax.set_ylabel('Y') 59 | ax.set_zlabel('Z') 60 | return fig, ax 61 | 62 | def on_unit_sphere(self, zero_mean=False): 63 | # radius == 1 64 | v = self.vertex_array # (N, D) 65 | if zero_mean: 66 | a = numpy.mean(v[:, 0:3], axis=0, keepdims=True) # (1, 3) 67 | v[:, 0:3] = v[:, 0:3] - a 68 | n = numpy.linalg.norm(v[:, 0:3], axis=1) # (N,) 69 | m = numpy.max(n) # scalar 70 | v[:, 0:3] = v[:, 0:3] / m 71 | self._vertices = v 72 | return self 73 | 74 | def on_unit_cube(self, zero_mean=False): 75 | # volume == 1 76 | v = self.vertex_array # (N, D) 77 | if zero_mean: 78 | a = numpy.mean(v[:, 0:3], axis=0, keepdims=True) # (1, 3) 79 | v[:, 0:3] = v[:, 0:3] - a 80 | m = numpy.max(numpy.abs(v)) # scalar 81 | v[:, 0:3] = v[:, 0:3] / (m * 2) 82 | self._vertices = v 83 | return self 84 | 85 | def rot_x(self): 86 | # camera local (up: +Y, front: -Z) -> model local (up: +Z, front: +Y). 87 | v = self.vertex_array 88 | t = numpy.copy(v[:, 1]) 89 | v[:, 1] = -numpy.copy(v[:, 2]) 90 | v[:, 2] = t 91 | self._vertices = list(map(tuple, v)) 92 | return self 93 | 94 | def rot_zc(self): 95 | # R = [0, -1; 96 | # 1, 0] 97 | v = self.vertex_array 98 | x = numpy.copy(v[:, 0]) 99 | y = numpy.copy(v[:, 1]) 100 | v[:, 0] = -y 101 | v[:, 1] = x 102 | self._vertices = list(map(tuple, v)) 103 | return self 104 | 105 | def offread(filepath, points_only=True): 106 | """ read Geomview OFF file. """ 107 | with open(filepath, 'r') as fin: 108 | mesh, fixme = _load_off(fin, points_only) 109 | if fixme: 110 | _fix_modelnet_broken_off(filepath) 111 | return mesh 112 | 113 | def _load_off(fin, points_only): 114 | """ read Geomview OFF file. """ 115 | mesh = Mesh() 116 | 117 | fixme = False 118 | sig = fin.readline().strip() 119 | if sig == 'OFF': 120 | line = fin.readline().strip() 121 | num_verts, num_faces, num_edges = tuple([int(s) for s in line.split(' ')]) 122 | elif sig[0:3] == 'OFF': # ...broken data in ModelNet (missing '\n')... 123 | line = sig[3:] 124 | num_verts, num_faces, num_edges = tuple([int(s) for s in line.split(' ')]) 125 | fixme = True 126 | else: 127 | raise RuntimeError('unknown format') 128 | 129 | for v in range(num_verts): 130 | vp = tuple(float(s) for s in fin.readline().strip().split(' ')) 131 | mesh._vertices.append(vp) 132 | 133 | if points_only: 134 | return mesh, fixme 135 | 136 | for f in range(num_faces): 137 | fc = tuple([int(s) for s in fin.readline().strip().split(' ')][1:]) 138 | mesh._faces.append(fc) 139 | 140 | return mesh, fixme 141 | 142 | def _fix_modelnet_broken_off(filepath): 143 | oldfile = '{}.orig'.format(filepath) 144 | os.rename(filepath, oldfile) 145 | with open(oldfile, 'r') as fin: 146 | with open(filepath, 'w') as fout: 147 | sig = fin.readline().strip() 148 | line = sig[3:] 149 | print('OFF', file=fout) 150 | print(line, file=fout) 151 | for line in fin: 152 | print(line.strip(), file=fout) 153 | 154 | 155 | def objread(filepath, points_only=True): 156 | """Loads a Wavefront OBJ file. """ 157 | _vertices = [] 158 | _normals = [] 159 | _texcoords = [] 160 | _faces = [] 161 | _mtl_name = None 162 | 163 | material = None 164 | for line in open(filepath, "r"): 165 | if line.startswith('#'): continue 166 | values = line.split() 167 | if not values: continue 168 | if values[0] == 'v': 169 | v = tuple(map(float, values[1:4])) 170 | _vertices.append(v) 171 | elif values[0] == 'vn': 172 | v = tuple(map(float, values[1:4])) 173 | _normals.append(v) 174 | elif values[0] == 'vt': 175 | _texcoords.append(tuple(map(float, values[1:3]))) 176 | elif values[0] in ('usemtl', 'usemat'): 177 | material = values[1] 178 | elif values[0] == 'mtllib': 179 | _mtl_name = values[1] 180 | elif values[0] == 'f': 181 | face_ = [] 182 | texcoords_ = [] 183 | norms_ = [] 184 | for v in values[1:]: 185 | w = v.split('/') 186 | face_.append(int(w[0]) - 1) 187 | if len(w) >= 2 and len(w[1]) > 0: 188 | texcoords_.append(int(w[1]) - 1) 189 | else: 190 | texcoords_.append(-1) 191 | if len(w) >= 3 and len(w[2]) > 0: 192 | norms_.append(int(w[2]) - 1) 193 | else: 194 | norms_.append(-1) 195 | #_faces.append((face_, norms_, texcoords_, material)) 196 | _faces.append(face_) 197 | 198 | mesh = Mesh() 199 | mesh._vertices = _vertices 200 | if points_only: 201 | return mesh 202 | 203 | mesh._faces = _faces 204 | 205 | return mesh 206 | 207 | 208 | if __name__ == '__main__': 209 | def test1(): 210 | mesh = objread('model_normalized.obj', points_only=False) 211 | #mesh.on_unit_sphere() 212 | mesh.rot_x() 213 | mesh.plot(c='m') 214 | matplotlib.pyplot.show() 215 | test1() 216 | 217 | #EOF -------------------------------------------------------------------------------- /ptlk/data/transforms.py: -------------------------------------------------------------------------------- 1 | """ gives some transform methods for 3d points """ 2 | import math 3 | 4 | import numpy as np 5 | import torch 6 | import torch.utils.data 7 | 8 | from . import mesh 9 | from .. import so3 10 | from .. import se3 11 | 12 | 13 | class Mesh2Points: 14 | def __init__(self): 15 | pass 16 | 17 | def __call__(self, mesh): 18 | mesh = mesh.clone() 19 | v = mesh.vertex_array 20 | return torch.from_numpy(v).type(dtype=torch.float) 21 | 22 | class OnUnitSphere: 23 | def __init__(self, zero_mean=False): 24 | self.zero_mean = zero_mean 25 | 26 | def __call__(self, tensor): 27 | if self.zero_mean: 28 | m = tensor.mean(dim=0, keepdim=True) # [N, D] -> [1, D] 29 | v = tensor - m 30 | else: 31 | v = tensor 32 | nn = v.norm(p=2, dim=1) # [N, D] -> [N] 33 | nmax = torch.max(nn) 34 | return v / nmax 35 | 36 | class OnUnitCube: 37 | def __init__(self): 38 | pass 39 | 40 | def method1(self, tensor): 41 | m = tensor.mean(dim=0, keepdim=True) # [N, D] -> [1, D] 42 | v = tensor - m 43 | s = torch.max(v.abs()) 44 | v = v / s * 0.5 45 | return v 46 | 47 | def method2(self, tensor): 48 | c = torch.max(tensor, dim=0)[0] - torch.min(tensor, dim=0)[0] # [N, D] -> [D] 49 | s = torch.max(c) # -> scalar 50 | v = tensor / s 51 | return v - v.mean(dim=0, keepdim=True) 52 | 53 | def __call__(self, tensor): 54 | #return self.method1(tensor) 55 | return self.method2(tensor) 56 | 57 | 58 | class Resampler: 59 | """ [N, D] -> [M, D] """ 60 | def __init__(self, num): 61 | self.num = num 62 | 63 | def __call__(self, tensor): 64 | num_points, dim_p = tensor.size() 65 | out = torch.zeros(self.num, dim_p).to(tensor) 66 | 67 | selected = 0 68 | while selected < self.num: 69 | remainder = self.num - selected 70 | idx = torch.randperm(num_points) 71 | sel = min(remainder, num_points) 72 | val = tensor[idx[:sel]] 73 | out[selected:(selected + sel)] = val 74 | selected += sel 75 | return out 76 | 77 | class RandomTranslate: 78 | def __init__(self, mag=None, randomly=True): 79 | self.mag = 1.0 if mag is None else mag 80 | self.randomly = randomly 81 | self.igt = None 82 | 83 | def __call__(self, tensor): 84 | # tensor: [N, 3] 85 | amp = torch.rand(1) if self.randomly else 1.0 86 | t = torch.randn(1, 3).to(tensor) 87 | t = t / t.norm(p=2, dim=1, keepdim=True) * amp * self.mag 88 | 89 | g = torch.eye(4).to(tensor) 90 | g[0:3, 3] = t[0, :] 91 | self.igt = g # [4, 4] 92 | 93 | p1 = tensor + t 94 | return p1 95 | 96 | class RandomRotator: 97 | def __init__(self, mag=None, randomly=True): 98 | self.mag = math.pi if mag is None else mag 99 | self.randomly = randomly 100 | self.igt = None 101 | 102 | def __call__(self, tensor): 103 | # tensor: [N, 3] 104 | amp = torch.rand(1) if self.randomly else 1.0 105 | w = torch.randn(1, 3) 106 | w = w / w.norm(p=2, dim=1, keepdim=True) * amp * self.mag 107 | 108 | g = so3.exp(w).to(tensor) # [1, 3, 3] 109 | self.igt = g.squeeze(0) # [3, 3] 110 | 111 | p1 = so3.transform(g, tensor) # [1, 3, 3] x [N, 3] -> [N, 3] 112 | return p1 113 | 114 | class RandomRotatorZ: 115 | def __init__(self): 116 | self.mag = 2 * math.pi 117 | 118 | def __call__(self, tensor): 119 | # tensor: [N, 3] 120 | w = torch.Tensor([0, 0, 1]).view(1, 3) * torch.rand(1) * self.mag 121 | 122 | g = so3.exp(w).to(tensor) # [1, 3, 3] 123 | 124 | p1 = so3.transform(g, tensor) 125 | return p1 126 | 127 | class RandomJitter: 128 | """ generate perturbations """ 129 | def __init__(self, scale=0.01, clip=0.05): 130 | self.scale = scale 131 | self.clip = clip 132 | self.e = None 133 | 134 | def jitter(self, tensor): 135 | noise = torch.zeros_like(tensor).to(tensor) # [N, 3] 136 | noise.normal_(0, self.scale) 137 | noise.clamp_(-self.clip, self.clip) 138 | self.e = noise 139 | return tensor.add(noise) 140 | 141 | def __call__(self, tensor): 142 | return self.jitter(tensor) 143 | 144 | 145 | class RandomTransformSE3: 146 | """ rigid motion """ 147 | def __init__(self, mag=1, mag_randomly=False): 148 | self.mag = mag 149 | self.randomly = mag_randomly 150 | 151 | self.gt = None 152 | self.igt = None 153 | 154 | def generate_transform(self): 155 | # return: a twist-vector 156 | amp = self.mag 157 | if self.randomly: 158 | amp = torch.rand(1, 1) * self.mag 159 | x = torch.randn(1, 6) 160 | x = x / x.norm(p=2, dim=1, keepdim=True) * amp 161 | 162 | return x # [1, 6] 163 | 164 | def apply_transform(self, p0, x): 165 | # p0: [N, 3] 166 | # x: [1, 6] 167 | g = se3.exp(x).to(p0) # [1, 4, 4] 168 | gt = se3.exp(-x).to(p0) # [1, 4, 4] 169 | 170 | p1 = se3.transform(g, p0) 171 | self.gt = gt.squeeze(0) # gt: p1 -> p0 172 | self.igt = g.squeeze(0) # igt: p0 -> p1 173 | return p1 174 | 175 | def transform(self, tensor): 176 | x = self.generate_transform() 177 | return self.apply_transform(tensor, x) 178 | 179 | def __call__(self, tensor): 180 | return self.transform(tensor) 181 | 182 | 183 | 184 | #EOF -------------------------------------------------------------------------------- /ptlk/invmat.py: -------------------------------------------------------------------------------- 1 | """ inverse matrix """ 2 | 3 | import torch 4 | 5 | 6 | def batch_inverse(x): 7 | """ M(n) -> M(n); x -> x^-1 """ 8 | batch_size, h, w = x.size() 9 | assert h == w 10 | y = torch.zeros_like(x) 11 | for i in range(batch_size): 12 | y[i, :, :] = x[i, :, :].inverse() 13 | return y 14 | 15 | def batch_inverse_dx(y): 16 | """ backward """ 17 | # Let y(x) = x^-1. 18 | # compute dy 19 | # dy = dy(j,k) 20 | # = - y(j,m) * dx(m,n) * y(n,k) 21 | # = - y(j,m) * y(n,k) * dx(m,n) 22 | # therefore, 23 | # dy(j,k)/dx(m,n) = - y(j,m) * y(n,k) 24 | batch_size, h, w = y.size() 25 | assert h == w 26 | # compute dy(j,k,m,n) = dy(j,k)/dx(m,n) = - y(j,m) * y(n,k) 27 | # = - (y(j,:))' * y'(k,:) 28 | yl = y.repeat(1, 1, h).view(batch_size*h*h, h, 1) 29 | yr = y.transpose(1, 2).repeat(1, h, 1).view(batch_size*h*h, 1, h) 30 | dy = - yl.bmm(yr).view(batch_size, h, h, h, h) 31 | 32 | # compute dy(m,n,j,k) = dy(j,k)/dx(m,n) = - y(j,m) * y(n,k) 33 | # = - (y'(m,:))' * y(n,:) 34 | #yl = y.transpose(1, 2).repeat(1, 1, h).view(batch_size*h*h, h, 1) 35 | #yr = y.repeat(1, h, 1).view(batch_size*h*h, 1, h) 36 | #dy = - yl.bmm(yr).view(batch_size, h, h, h, h) 37 | 38 | return dy 39 | 40 | 41 | def batch_pinv_dx(x): 42 | """ returns y = (x'*x)^-1 * x' and dy/dx. """ 43 | # y = (x'*x)^-1 * x' 44 | # = s^-1 * x' 45 | # = b * x' 46 | # d{y(j,k)}/d{x(m,n)} 47 | # = d{b(j,i) * x(k,i)}/d{x(m,n)} 48 | # = d{b(j,i)}/d{x(m,n)} * x(k,i) + b(j,i) * d{x(k,i)}/d{x(m,n)} 49 | # d{b(j,i)}/d{x(m,n)} 50 | # = d{b(j,i)}/d{s(p,q)} * d{s(p,q)}/d{x(m,n)} 51 | # = -b(j,p)*b(q,i) * d{s(p,q)}/d{x(m,n)} 52 | # d{s(p,q)}/d{x(m,n)} 53 | # = d{x(t,p)*x(t,q)}/d{x(m,n)} 54 | # = d{x(t,p)}/d{x(m,n)} * x(t,q) + x(t,p) * d{x(t,q)}/d{x(m,n)} 55 | batch_size, h, w = x.size() 56 | xt = x.transpose(1, 2) 57 | s = xt.bmm(x) 58 | b = batch_inverse(s) 59 | y = b.bmm(xt) 60 | 61 | # dx/dx 62 | ex = torch.eye(h*w).to(x).unsqueeze(0).view(1, h, w, h, w) 63 | # ds/dx = dx(t,_)/dx * x(t,_) + x(t,_) * dx(t,_)/dx 64 | ex1 = ex.view(1, h, w*h*w) # [t, p*m*n] 65 | dx1 = x.transpose(1, 2).matmul(ex1).view(batch_size, w, w, h, w) # [q, p,m,n] 66 | ds_dx = dx1.transpose(1, 2) + dx1 # [p, q, m, n] 67 | # db/ds 68 | db_ds = batch_inverse_dx(b) # [j, i, p, q] 69 | # db/dx = db/d{s(p,q)} * d{s(p,q)}/dx 70 | db1 = db_ds.view(batch_size, w*w, w*w).bmm(ds_dx.view(batch_size, w*w, h*w)) 71 | db_dx = db1.view(batch_size, w, w, h, w) # [j, i, m, n] 72 | # dy/dx = db(_,i)/dx * x(_,i) + b(_,i) * dx(_,i)/dx 73 | dy1 = db_dx.transpose(1, 2).contiguous().view(batch_size, w, w*h*w) 74 | dy1 = x.matmul(dy1).view(batch_size, h, w, h, w) # [k, j, m, n] 75 | ext = ex.transpose(1, 2).contiguous().view(1, w, h*h*w) 76 | dy2 = b.matmul(ext).view(batch_size, w, h, h, w) # [j, k, m, n] 77 | dy_dx = dy1.transpose(1, 2) + dy2 78 | 79 | return y, dy_dx 80 | 81 | 82 | class InvMatrix(torch.autograd.Function): 83 | """ M(n) -> M(n); x -> x^-1. 84 | """ 85 | @staticmethod 86 | def forward(ctx, x): 87 | y = batch_inverse(x) 88 | ctx.save_for_backward(y) 89 | return y 90 | 91 | @staticmethod 92 | def backward(ctx, grad_output): 93 | y, = ctx.saved_tensors # v0.4 94 | #y, = ctx.saved_variables # v0.3.1 95 | batch_size, h, w = y.size() 96 | assert h == w 97 | 98 | # Let y(x) = x^-1 and assume any function f(y(x)). 99 | # compute df/dx(m,n)... 100 | # df/dx(m,n) = df/dy(j,k) * dy(j,k)/dx(m,n) 101 | # well, df/dy is 'grad_output' 102 | # and so we will return 'grad_input = df/dy(j,k) * dy(j,k)/dx(m,n)' 103 | 104 | dy = batch_inverse_dx(y) # dy(j,k,m,n) = dy(j,k)/dx(m,n) 105 | go = grad_output.contiguous().view(batch_size, 1, h*h) # [1, (j*k)] 106 | ym = dy.view(batch_size, h*h, h*h) # [(j*k), (m*n)] 107 | r = go.bmm(ym) # [1, (m*n)] 108 | grad_input = r.view(batch_size, h, h) # [m, n] 109 | 110 | return grad_input 111 | 112 | 113 | 114 | if __name__ == '__main__': 115 | def test(): 116 | x = torch.randn(2, 3, 2) 117 | x_val = x.requires_grad_() 118 | 119 | s_val = x_val.transpose(1, 2).bmm(x_val) 120 | s_inv = InvMatrix.apply(s_val) 121 | y_val = s_inv.bmm(x_val.transpose(1, 2)) 122 | y_val.sum().backward() 123 | t1 = x_val.grad 124 | 125 | y, dy_dx = batch_pinv_dx(x) 126 | t2 = dy_dx.sum(1).sum(1) 127 | 128 | print(t1) 129 | print(t2) 130 | print(t1 - t2) 131 | 132 | test() 133 | 134 | #EOF 135 | -------------------------------------------------------------------------------- /ptlk/pointlk.py: -------------------------------------------------------------------------------- 1 | """ PointLK ver. 2018.07.06. 2 | using approximated Jacobian by backward-difference. 3 | """ 4 | 5 | import numpy 6 | import torch 7 | 8 | from . import pointnet 9 | from . import se3, so3, invmat 10 | 11 | 12 | class PointLK(torch.nn.Module): 13 | def __init__(self, ptnet, delta=1.0e-2, learn_delta=False): 14 | super().__init__() 15 | self.ptnet = ptnet 16 | self.inverse = invmat.InvMatrix.apply 17 | self.exp = se3.Exp # [B, 6] -> [B, 4, 4] 18 | self.transform = se3.transform # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3] 19 | 20 | w1 = delta 21 | w2 = delta 22 | w3 = delta 23 | v1 = delta 24 | v2 = delta 25 | v3 = delta 26 | twist = torch.Tensor([w1, w2, w3, v1, v2, v3]) 27 | self.dt = torch.nn.Parameter(twist.view(1, 6), requires_grad=learn_delta) 28 | 29 | # results 30 | self.last_err = None 31 | self.g_series = None # for debug purpose 32 | self.prev_r = None 33 | self.g = None # estimation result 34 | self.itr = 0 35 | 36 | @staticmethod 37 | def rsq(r): 38 | # |r| should be 0 39 | z = torch.zeros_like(r) 40 | return torch.nn.functional.mse_loss(r, z, size_average=False) 41 | 42 | @staticmethod 43 | def comp(g, igt): 44 | """ |g*igt - I| (should be 0) """ 45 | assert g.size(0) == igt.size(0) 46 | assert g.size(1) == igt.size(1) and g.size(1) == 4 47 | assert g.size(2) == igt.size(2) and g.size(2) == 4 48 | A = g.matmul(igt) 49 | I = torch.eye(4).to(A).view(1, 4, 4).expand(A.size(0), 4, 4) 50 | return torch.nn.functional.mse_loss(A, I, size_average=True) * 16 51 | 52 | @staticmethod 53 | def do_forward(net, p0, p1, maxiter=10, xtol=1.0e-7, p0_zero_mean=True, p1_zero_mean=True): 54 | a0 = torch.eye(4).view(1, 4, 4).expand(p0.size(0), 4, 4).to(p0) # [B, 4, 4] 55 | a1 = torch.eye(4).view(1, 4, 4).expand(p1.size(0), 4, 4).to(p1) # [B, 4, 4] 56 | if p0_zero_mean: 57 | p0_m = p0.mean(dim=1) # [B, N, 3] -> [B, 3] 58 | a0[:, 0:3, 3] = p0_m 59 | q0 = p0 - p0_m.unsqueeze(1) 60 | else: 61 | q0 = p0 62 | 63 | if p1_zero_mean: 64 | #print(numpy.any(numpy.isnan(p1.numpy()))) 65 | p1_m = p1.mean(dim=1) # [B, N, 3] -> [B, 3] 66 | a1[:, 0:3, 3] = -p1_m 67 | q1 = p1 - p1_m.unsqueeze(1) 68 | else: 69 | q1 = p1 70 | 71 | r = net(q0, q1, maxiter=maxiter, xtol=xtol) 72 | 73 | if p0_zero_mean or p1_zero_mean: 74 | #output' = trans(p0_m) * output * trans(-p1_m) 75 | # = [I, p0_m;] * [R, t;] * [I, -p1_m;] 76 | # [0, 1 ] [0, 1 ] [0, 1 ] 77 | est_g = net.g 78 | if p0_zero_mean: 79 | est_g = a0.to(est_g).bmm(est_g) 80 | if p1_zero_mean: 81 | est_g = est_g.bmm(a1.to(est_g)) 82 | net.g = est_g 83 | 84 | est_gs = net.g_series # [M, B, 4, 4] 85 | if p0_zero_mean: 86 | est_gs = a0.unsqueeze(0).contiguous().to(est_gs).matmul(est_gs) 87 | if p1_zero_mean: 88 | est_gs = est_gs.matmul(a1.unsqueeze(0).contiguous().to(est_gs)) 89 | net.g_series = est_gs 90 | 91 | return r 92 | 93 | def forward(self, p0, p1, maxiter=10, xtol=1.0e-7): 94 | g0 = torch.eye(4).to(p0).view(1, 4, 4).expand(p0.size(0), 4, 4).contiguous() 95 | r, g, itr = self.iclk(g0, p0, p1, maxiter, xtol) 96 | 97 | self.g = g 98 | self.itr = itr 99 | return r 100 | 101 | def update(self, g, dx): 102 | # [B, 4, 4] x [B, 6] -> [B, 4, 4] 103 | dg = self.exp(dx) 104 | return dg.matmul(g) 105 | 106 | def approx_Jic(self, p0, f0, dt): 107 | # p0: [B, N, 3], Variable 108 | # f0: [B, K], corresponding feature vector 109 | # dt: [B, 6], Variable 110 | # Jk = (ptnet(p(-delta[k], p0)) - f0) / delta[k] 111 | 112 | batch_size = p0.size(0) 113 | num_points = p0.size(1) 114 | 115 | # compute transforms 116 | transf = torch.zeros(batch_size, 6, 4, 4).to(p0) 117 | for b in range(p0.size(0)): 118 | d = torch.diag(dt[b, :]) # [6, 6] 119 | D = self.exp(-d) # [6, 4, 4] 120 | transf[b, :, :, :] = D[:, :, :] 121 | transf = transf.unsqueeze(2).contiguous() # [B, 6, 1, 4, 4] 122 | p = self.transform(transf, p0.unsqueeze(1)) # x [B, 1, N, 3] -> [B, 6, N, 3] 123 | 124 | #f0 = self.ptnet(p0).unsqueeze(-1) # [B, K, 1] 125 | f0 = f0.unsqueeze(-1) # [B, K, 1] 126 | f = self.ptnet(p.view(-1, num_points, 3)).view(batch_size, 6, -1).transpose(1, 2) # [B, K, 6] 127 | 128 | df = f0 - f # [B, K, 6] 129 | J = df / dt.unsqueeze(1) 130 | 131 | return J 132 | 133 | def iclk(self, g0, p0, p1, maxiter, xtol): 134 | training = self.ptnet.training 135 | batch_size = p0.size(0) 136 | 137 | g = g0 138 | self.g_series = torch.zeros(maxiter+1, *g0.size(), dtype=g0.dtype) 139 | self.g_series[0] = g0.clone() 140 | 141 | if training: 142 | # first, update BatchNorm modules 143 | f0 = self.ptnet(p0) 144 | f1 = self.ptnet(p1) 145 | self.ptnet.eval() # and fix them. 146 | 147 | # re-calc. with current modules 148 | f0 = self.ptnet(p0) # [B, N, 3] -> [B, K] 149 | 150 | # approx. J by finite difference 151 | dt = self.dt.to(p0).expand(batch_size, 6) 152 | J = self.approx_Jic(p0, f0, dt) 153 | 154 | self.last_err = None 155 | itr = -1 156 | # compute pinv(J) to solve J*x = -r 157 | try: 158 | Jt = J.transpose(1, 2) # [B, 6, K] 159 | H = Jt.bmm(J) # [B, 6, 6] 160 | B = self.inverse(H) 161 | pinv = B.bmm(Jt) # [B, 6, K] 162 | except RuntimeError as err: 163 | # singular...? 164 | self.last_err = err 165 | #print(err) 166 | # Perhaps we can use MP-inverse, but,... 167 | # probably, self.dt is way too small... 168 | f1 = self.ptnet(p1) # [B, N, 3] -> [B, K] 169 | r = f1 - f0 170 | self.ptnet.train(training) 171 | return r, g, itr 172 | 173 | itr = 0 174 | r = None 175 | for itr in range(maxiter): 176 | self.prev_r = r 177 | p = self.transform(g.unsqueeze(1), p1) # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3] 178 | f = self.ptnet(p) # [B, N, 3] -> [B, K] 179 | r = f - f0 180 | 181 | dx = -pinv.bmm(r.unsqueeze(-1)).view(batch_size, 6) 182 | 183 | # DEBUG. 184 | #norm_r = r.norm(p=2, dim=1) 185 | #print('itr,{},|r|,{}'.format(itr+1, ','.join(map(str, norm_r.data.tolist())))) 186 | 187 | check = dx.norm(p=2, dim=1, keepdim=True).max() 188 | if float(check) < xtol: 189 | if itr == 0: 190 | self.last_err = 0 # no update. 191 | break 192 | 193 | g = self.update(g, dx) 194 | self.g_series[itr+1] = g.clone() 195 | 196 | rep = len(range(itr, maxiter)) 197 | self.g_series[(itr+1):] = g.clone().unsqueeze(0).repeat(rep, 1, 1, 1) 198 | 199 | self.ptnet.train(training) 200 | return r, g, (itr+1) 201 | 202 | 203 | 204 | #EOF -------------------------------------------------------------------------------- /ptlk/pointnet.py: -------------------------------------------------------------------------------- 1 | """ PointNet 2 | References. 3 | .. [1] Charles R. Qi, Hao Su, Kaichun Mo and Leonidas J. Guibas, 4 | "PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation", 5 | CVPR (2017) 6 | """ 7 | 8 | import torch 9 | 10 | def flatten(x): 11 | return x.view(x.size(0), -1) 12 | 13 | class Flatten(torch.nn.Module): 14 | def __init__(self): 15 | super().__init__() 16 | def forward(self, x): 17 | return flatten(x) 18 | 19 | 20 | def mlp_layers(nch_input, nch_layers, b_shared=True, bn_momentum=0.1, dropout=0.0): 21 | """ [B, Cin, N] -> [B, Cout, N] or 22 | [B, Cin] -> [B, Cout] 23 | """ 24 | layers = [] 25 | last = nch_input 26 | for i, outp in enumerate(nch_layers): 27 | if b_shared: 28 | weights = torch.nn.Conv1d(last, outp, 1) 29 | else: 30 | weights = torch.nn.Linear(last, outp) 31 | layers.append(weights) 32 | layers.append(torch.nn.BatchNorm1d(outp, momentum=bn_momentum)) 33 | layers.append(torch.nn.ReLU()) 34 | if b_shared == False and dropout > 0.0: 35 | layers.append(torch.nn.Dropout(dropout)) 36 | last = outp 37 | return layers 38 | 39 | class MLPNet(torch.nn.Module): 40 | """ Multi-layer perception. 41 | [B, Cin, N] -> [B, Cout, N] or 42 | [B, Cin] -> [B, Cout] 43 | """ 44 | def __init__(self, nch_input, nch_layers, b_shared=True, bn_momentum=0.1, dropout=0.0): 45 | super().__init__() 46 | list_layers = mlp_layers(nch_input, nch_layers, b_shared, bn_momentum, dropout) 47 | self.layers = torch.nn.Sequential(*list_layers) 48 | 49 | def forward(self, inp): 50 | out = self.layers(inp) 51 | return out 52 | 53 | 54 | def symfn_max(x): 55 | # [B, K, N] -> [B, K, 1] 56 | a = torch.nn.functional.max_pool1d(x, x.size(-1)) 57 | #a, _ = torch.max(x, dim=-1, keepdim=True) 58 | return a 59 | 60 | def symfn_avg(x): 61 | a = torch.nn.functional.avg_pool1d(x, x.size(-1)) 62 | #a = torch.sum(x, dim=-1, keepdim=True) / x.size(-1) 63 | return a 64 | 65 | class TNet(torch.nn.Module): 66 | """ [B, K, N] -> [B, K, K] 67 | """ 68 | def __init__(self, K): 69 | super().__init__() 70 | # [B, K, N] -> [B, K*K] 71 | self.mlp1 = torch.nn.Sequential(*mlp_layers(K, [64, 128, 1024], b_shared=True)) 72 | self.mlp2 = torch.nn.Sequential(*mlp_layers(1024, [512, 256], b_shared=False)) 73 | self.lin = torch.nn.Linear(256, K*K) 74 | 75 | for param in self.mlp1.parameters(): 76 | torch.nn.init.constant_(param, 0.0) 77 | for param in self.mlp2.parameters(): 78 | torch.nn.init.constant_(param, 0.0) 79 | for param in self.lin.parameters(): 80 | torch.nn.init.constant_(param, 0.0) 81 | 82 | def forward(self, inp): 83 | K = inp.size(1) 84 | N = inp.size(2) 85 | eye = torch.eye(K).unsqueeze(0).to(inp) # [1, K, K] 86 | 87 | x = self.mlp1(inp) 88 | x = flatten(torch.nn.functional.max_pool1d(x, N)) 89 | x = self.mlp2(x) 90 | x = self.lin(x) 91 | 92 | x = x.view(-1, K, K) 93 | x = x + eye 94 | return x 95 | 96 | 97 | class PointNet_features(torch.nn.Module): 98 | def __init__(self, dim_k=1024, use_tnet=False, sym_fn=symfn_max, scale=1): 99 | super().__init__() 100 | mlp_h1 = [int(64/scale), int(64/scale)] 101 | mlp_h2 = [int(64/scale), int(128/scale), int(dim_k/scale)] 102 | 103 | self.h1 = MLPNet(3, mlp_h1, b_shared=True).layers 104 | self.h2 = MLPNet(mlp_h1[-1], mlp_h2, b_shared=True).layers 105 | #self.sy = torch.nn.Sequential(torch.nn.MaxPool1d(num_points), Flatten()) 106 | self.sy = sym_fn 107 | 108 | self.tnet1 = TNet(3) if use_tnet else None 109 | self.tnet2 = TNet(mlp_h1[-1]) if use_tnet else None 110 | 111 | self.t_out_t2 = None 112 | self.t_out_h1 = None 113 | 114 | def forward(self, points): 115 | """ points -> features 116 | [B, N, 3] -> [B, K] 117 | """ 118 | x = points.transpose(1, 2) # [B, 3, N] 119 | if self.tnet1: 120 | t1 = self.tnet1(x) 121 | x = t1.bmm(x) 122 | 123 | x = self.h1(x) 124 | if self.tnet2: 125 | t2 = self.tnet2(x) 126 | self.t_out_t2 = t2 127 | x = t2.bmm(x) 128 | self.t_out_h1 = x # local features 129 | 130 | x = self.h2(x) 131 | #x = flatten(torch.nn.functional.max_pool1d(x, x.size(-1))) 132 | x = flatten(self.sy(x)) 133 | 134 | #if self.ret_global: 135 | # pass 136 | #else: 137 | # # local + global 138 | # l0 = self.t_out_h1 # [B, 64, N] 139 | # g0 = x # [B, K] 140 | # x = torch.cat((l0, g0.unsqueeze(2).repeat(1, 1, num_points)), dim=1) 141 | 142 | return x 143 | 144 | class PointNet_classifier(torch.nn.Module): 145 | def __init__(self, num_c, ptfeat, dim_k): 146 | super().__init__() 147 | self.features = ptfeat 148 | list_layers = mlp_layers(dim_k, [512, 256], b_shared=False, bn_momentum=0.1, dropout=0.0) 149 | #list_layers = mlp_layers(dim_k, [512, 256], b_shared=False, bn_momentum=0.1, dropout=0.3) 150 | list_layers.append(torch.nn.Linear(256, num_c)) 151 | self.classifier = torch.nn.Sequential(*list_layers) 152 | 153 | def forward(self, points): 154 | feat = self.features(points) 155 | out = self.classifier(feat) 156 | return out 157 | 158 | def loss(self, out, target, w=0.001): 159 | loss_c = torch.nn.functional.nll_loss( 160 | torch.nn.functional.log_softmax(out, dim=1), target, size_average=False) 161 | 162 | t2 = self.features.t_out_t2 163 | if (t2 is None) or (w == 0): 164 | return loss_c 165 | 166 | batch = t2.size(0) 167 | K = t2.size(1) # [B, K, K] 168 | I = torch.eye(K).repeat(batch, 1, 1).to(t2) 169 | A = t2.bmm(t2.transpose(1, 2)) 170 | loss_m = torch.nn.functional.mse_loss(A, I, size_average=False) 171 | loss = loss_c + w * loss_m 172 | return loss 173 | 174 | 175 | #EOF -------------------------------------------------------------------------------- /ptlk/se3.py: -------------------------------------------------------------------------------- 1 | """ 3-d rigid body transfomation group and corresponding Lie algebra. """ 2 | import torch 3 | from .sinc import sinc1, sinc2, sinc3 4 | from . import so3 5 | 6 | def twist_prod(x, y): 7 | x_ = x.view(-1, 6) 8 | y_ = y.view(-1, 6) 9 | 10 | xw, xv = x_[:, 0:3], x_[:, 3:6] 11 | yw, yv = y_[:, 0:3], y_[:, 3:6] 12 | 13 | zw = so3.cross_prod(xw, yw) 14 | zv = so3.cross_prod(xw, yv) + so3.cross_prod(xv, yw) 15 | 16 | z = torch.cat((zw, zv), dim=1) 17 | 18 | return z.view_as(x) 19 | 20 | def liebracket(x, y): 21 | return twist_prod(x, y) 22 | 23 | 24 | def mat(x): 25 | # size: [*, 6] -> [*, 4, 4] 26 | x_ = x.view(-1, 6) 27 | w1, w2, w3 = x_[:, 0], x_[:, 1], x_[:, 2] 28 | v1, v2, v3 = x_[:, 3], x_[:, 4], x_[:, 5] 29 | O = torch.zeros_like(w1) 30 | 31 | X = torch.stack(( 32 | torch.stack(( O, -w3, w2, v1), dim=1), 33 | torch.stack(( w3, O, -w1, v2), dim=1), 34 | torch.stack((-w2, w1, O, v3), dim=1), 35 | torch.stack(( O, O, O, O), dim=1)), dim=1) 36 | return X.view(*(x.size()[0:-1]), 4, 4) 37 | 38 | def vec(X): 39 | X_ = X.view(-1, 4, 4) 40 | w1, w2, w3 = X_[:, 2, 1], X_[:, 0, 2], X_[:, 1, 0] 41 | v1, v2, v3 = X_[:, 0, 3], X_[:, 1, 3], X_[:, 2, 3] 42 | x = torch.stack((w1, w2, w3, v1, v2, v3), dim=1) 43 | return x.view(*X.size()[0:-2], 6) 44 | 45 | def genvec(): 46 | return torch.eye(6) 47 | 48 | def genmat(): 49 | return mat(genvec()) 50 | 51 | def exp(x): 52 | x_ = x.view(-1, 6) 53 | w, v = x_[:, 0:3], x_[:, 3:6] 54 | t = w.norm(p=2, dim=1).view(-1, 1, 1) 55 | W = so3.mat(w) 56 | S = W.bmm(W) 57 | I = torch.eye(3).to(w) 58 | 59 | # Rodrigues' rotation formula. 60 | #R = cos(t)*eye(3) + sinc1(t)*W + sinc2(t)*(w*w'); 61 | # = eye(3) + sinc1(t)*W + sinc2(t)*S 62 | R = I + sinc1(t)*W + sinc2(t)*S 63 | 64 | #V = sinc1(t)*eye(3) + sinc2(t)*W + sinc3(t)*(w*w') 65 | # = eye(3) + sinc2(t)*W + sinc3(t)*S 66 | V = I + sinc2(t)*W + sinc3(t)*S 67 | 68 | p = V.bmm(v.contiguous().view(-1, 3, 1)) 69 | 70 | z = torch.Tensor([0, 0, 0, 1]).view(1, 1, 4).repeat(x_.size(0), 1, 1).to(x) 71 | Rp = torch.cat((R, p), dim=2) 72 | g = torch.cat((Rp, z), dim=1) 73 | 74 | return g.view(*(x.size()[0:-1]), 4, 4) 75 | 76 | def inverse(g): 77 | g_ = g.view(-1, 4, 4) 78 | R = g_[:, 0:3, 0:3] 79 | p = g_[:, 0:3, 3] 80 | Q = R.transpose(1, 2) 81 | q = -Q.matmul(p.unsqueeze(-1)) 82 | 83 | z = torch.Tensor([0, 0, 0, 1]).view(1, 1, 4).repeat(g_.size(0), 1, 1).to(g) 84 | Qq = torch.cat((Q, q), dim=2) 85 | ig = torch.cat((Qq, z), dim=1) 86 | 87 | return ig.view(*(g.size()[0:-2]), 4, 4) 88 | 89 | 90 | def log(g): 91 | g_ = g.view(-1, 4, 4) 92 | R = g_[:, 0:3, 0:3] 93 | p = g_[:, 0:3, 3] 94 | 95 | w = so3.log(R) 96 | H = so3.inv_vecs_Xg_ig(w) 97 | v = H.bmm(p.contiguous().view(-1, 3, 1)).view(-1, 3) 98 | 99 | x = torch.cat((w, v), dim=1) 100 | return x.view(*(g.size()[0:-2]), 6) 101 | 102 | def transform(g, a): 103 | # g : SE(3), * x 4 x 4 104 | # a : R^3, * x 3[x N] 105 | g_ = g.view(-1, 4, 4) 106 | R = g_[:, 0:3, 0:3].contiguous().view(*(g.size()[0:-2]), 3, 3) 107 | p = g_[:, 0:3, 3].contiguous().view(*(g.size()[0:-2]), 3) 108 | if len(g.size()) == len(a.size()): 109 | b = R.matmul(a) + p.unsqueeze(-1) 110 | else: 111 | b = R.matmul(a.unsqueeze(-1)).squeeze(-1) + p 112 | return b 113 | 114 | def group_prod(g, h): 115 | # g, h : SE(3) 116 | g1 = g.matmul(h) 117 | return g1 118 | 119 | 120 | class ExpMap(torch.autograd.Function): 121 | """ Exp: se(3) -> SE(3) 122 | """ 123 | @staticmethod 124 | def forward(ctx, x): 125 | """ Exp: R^6 -> M(4), 126 | size: [B, 6] -> [B, 4, 4], 127 | or [B, 1, 6] -> [B, 1, 4, 4] 128 | """ 129 | ctx.save_for_backward(x) 130 | g = exp(x) 131 | return g 132 | 133 | @staticmethod 134 | def backward(ctx, grad_output): 135 | x, = ctx.saved_tensors 136 | g = exp(x) 137 | gen_k = genmat().to(x) 138 | 139 | # Let z = f(g) = f(exp(x)) 140 | # dz = df/dgij * dgij/dxk * dxk 141 | # = df/dgij * (d/dxk)[exp(x)]_ij * dxk 142 | # = df/dgij * [gen_k*g]_ij * dxk 143 | 144 | dg = gen_k.matmul(g.view(-1, 1, 4, 4)) 145 | # (k, i, j) 146 | dg = dg.to(grad_output) 147 | 148 | go = grad_output.contiguous().view(-1, 1, 4, 4) 149 | dd = go * dg 150 | grad_input = dd.sum(-1).sum(-1) 151 | 152 | return grad_input 153 | 154 | Exp = ExpMap.apply 155 | 156 | 157 | #EOF 158 | -------------------------------------------------------------------------------- /ptlk/sinc.py: -------------------------------------------------------------------------------- 1 | """ sinc(t) := sin(t) / t """ 2 | import torch 3 | from torch import sin, cos 4 | 5 | def sinc1(t): 6 | """ sinc1: t -> sin(t)/t """ 7 | e = 0.01 8 | r = torch.zeros_like(t) 9 | a = torch.abs(t) 10 | 11 | s = a < e 12 | c = (s == 0) 13 | t2 = t[s] ** 2 14 | r[s] = 1 - t2/6*(1 - t2/20*(1 - t2/42)) # Taylor series O(t^8) 15 | r[c] = sin(t[c]) / t[c] 16 | 17 | return r 18 | 19 | def sinc1_dt(t): 20 | """ d/dt(sinc1) """ 21 | e = 0.01 22 | r = torch.zeros_like(t) 23 | a = torch.abs(t) 24 | 25 | s = a < e 26 | c = (s == 0) 27 | t2 = t ** 2 28 | r[s] = -t[s]/3*(1 - t2[s]/10*(1 - t2[s]/28*(1 - t2[s]/54))) # Taylor series O(t^8) 29 | r[c] = cos(t[c])/t[c] - sin(t[c])/t2[c] 30 | 31 | return r 32 | 33 | def sinc1_dt_rt(t): 34 | """ d/dt(sinc1) / t """ 35 | e = 0.01 36 | r = torch.zeros_like(t) 37 | a = torch.abs(t) 38 | 39 | s = a < e 40 | c = (s == 0) 41 | t2 = t ** 2 42 | r[s] = -1/3*(1 - t2[s]/10*(1 - t2[s]/28*(1 - t2[s]/54))) # Taylor series O(t^8) 43 | r[c] = (cos(t[c]) / t[c] - sin(t[c]) / t2[c]) / t[c] 44 | 45 | return r 46 | 47 | 48 | def rsinc1(t): 49 | """ rsinc1: t -> t/sinc1(t) """ 50 | e = 0.01 51 | r = torch.zeros_like(t) 52 | a = torch.abs(t) 53 | 54 | s = a < e 55 | c = (s == 0) 56 | t2 = t[s] ** 2 57 | r[s] = (((31*t2)/42 + 7)*t2/60 + 1)*t2/6 + 1 # Taylor series O(t^8) 58 | r[c] = t[c] / sin(t[c]) 59 | 60 | return r 61 | 62 | def rsinc1_dt(t): 63 | """ d/dt(rsinc1) """ 64 | e = 0.01 65 | r = torch.zeros_like(t) 66 | a = torch.abs(t) 67 | 68 | s = a < e 69 | c = (s == 0) 70 | t2 = t[s] ** 2 71 | r[s] = ((((127*t2)/30 + 31)*t2/28 + 7)*t2/30 + 1)*t[s]/3 # Taylor series O(t^8) 72 | r[c] = 1/sin(t[c]) - (t[c]*cos(t[c]))/(sin(t[c])*sin(t[c])) 73 | 74 | return r 75 | 76 | def rsinc1_dt_csc(t): 77 | """ d/dt(rsinc1) / sin(t) """ 78 | e = 0.01 79 | r = torch.zeros_like(t) 80 | a = torch.abs(t) 81 | 82 | s = a < e 83 | c = (s == 0) 84 | t2 = t[s] ** 2 85 | r[s] = t2*(t2*((4*t2)/675 + 2/63) + 2/15) + 1/3 # Taylor series O(t^8) 86 | r[c] = (1/sin(t[c]) - (t[c]*cos(t[c]))/(sin(t[c])*sin(t[c]))) / sin(t[c]) 87 | 88 | return r 89 | 90 | 91 | def sinc2(t): 92 | """ sinc2: t -> (1 - cos(t)) / (t**2) """ 93 | e = 0.01 94 | r = torch.zeros_like(t) 95 | a = torch.abs(t) 96 | 97 | s = a < e 98 | c = (s == 0) 99 | t2 = t ** 2 100 | r[s] = 1/2*(1-t2[s]/12*(1-t2[s]/30*(1-t2[s]/56))) # Taylor series O(t^8) 101 | r[c] = (1-cos(t[c]))/t2[c] 102 | 103 | return r 104 | 105 | def sinc2_dt(t): 106 | """ d/dt(sinc2) """ 107 | e = 0.01 108 | r = torch.zeros_like(t) 109 | a = torch.abs(t) 110 | 111 | s = a < e 112 | c = (s == 0) 113 | t2 = t ** 2 114 | r[s] = -t[s]/12*(1 - t2[s]/5*(1.0/3 - t2[s]/56*(1.0/2 - t2[s]/135))) # Taylor series O(t^8) 115 | r[c] = sin(t[c])/t2[c] - 2*(1-cos(t[c]))/(t2[c]*t[c]) 116 | 117 | return r 118 | 119 | 120 | def sinc3(t): 121 | """ sinc3: t -> (t - sin(t)) / (t**3) """ 122 | e = 0.01 123 | r = torch.zeros_like(t) 124 | a = torch.abs(t) 125 | 126 | s = a < e 127 | c = (s == 0) 128 | t2 = t[s] ** 2 129 | r[s] = 1/6*(1-t2/20*(1-t2/42*(1-t2/72))) # Taylor series O(t^8) 130 | r[c] = (t[c]-sin(t[c]))/(t[c]**3) 131 | 132 | return r 133 | 134 | def sinc3_dt(t): 135 | """ d/dt(sinc3) """ 136 | e = 0.01 137 | r = torch.zeros_like(t) 138 | a = torch.abs(t) 139 | 140 | s = a < e 141 | c = (s == 0) 142 | t2 = t[s] ** 2 143 | r[s] = -t[s]/60*(1 - t2/21*(1 - t2/24*(1.0/2 - t2/165))) # Taylor series O(t^8) 144 | r[c] = (3*sin(t[c]) - t[c]*(cos(t[c]) + 2))/(t[c]**4) 145 | 146 | return r 147 | 148 | 149 | def sinc4(t): 150 | """ sinc4: t -> 1/t^2 * (1/2 - sinc2(t)) 151 | = 1/t^2 * (1/2 - (1 - cos(t))/t^2) 152 | """ 153 | e = 0.01 154 | r = torch.zeros_like(t) 155 | a = torch.abs(t) 156 | 157 | s = a < e 158 | c = (s == 0) 159 | t2 = t ** 2 160 | r[s] = 1/24*(1-t2/30*(1-t2/56*(1-t2/90))) # Taylor series O(t^8) 161 | r[c] = (0.5 - (1 - cos(t))/t2) / t2 162 | 163 | 164 | class Sinc1_autograd(torch.autograd.Function): 165 | @staticmethod 166 | def forward(ctx, theta): 167 | ctx.save_for_backward(theta) 168 | return sinc1(theta) 169 | 170 | @staticmethod 171 | def backward(ctx, grad_output): 172 | theta, = ctx.saved_tensors 173 | grad_theta = None 174 | if ctx.needs_input_grad[0]: 175 | grad_theta = grad_output * sinc1_dt(theta).to(grad_output) 176 | return grad_theta 177 | 178 | Sinc1 = Sinc1_autograd.apply 179 | 180 | class RSinc1_autograd(torch.autograd.Function): 181 | @staticmethod 182 | def forward(ctx, theta): 183 | ctx.save_for_backward(theta) 184 | return rsinc1(theta) 185 | 186 | @staticmethod 187 | def backward(ctx, grad_output): 188 | theta, = ctx.saved_tensors 189 | grad_theta = None 190 | if ctx.needs_input_grad[0]: 191 | grad_theta = grad_output * rsinc1_dt(theta).to(grad_output) 192 | return grad_theta 193 | 194 | RSinc1 = RSinc1_autograd.apply 195 | 196 | class Sinc2_autograd(torch.autograd.Function): 197 | @staticmethod 198 | def forward(ctx, theta): 199 | ctx.save_for_backward(theta) 200 | return sinc2(theta) 201 | 202 | @staticmethod 203 | def backward(ctx, grad_output): 204 | theta, = ctx.saved_tensors 205 | grad_theta = None 206 | if ctx.needs_input_grad[0]: 207 | grad_theta = grad_output * sinc2_dt(theta).to(grad_output) 208 | return grad_theta 209 | 210 | Sinc2 = Sinc2_autograd.apply 211 | 212 | class Sinc3_autograd(torch.autograd.Function): 213 | @staticmethod 214 | def forward(ctx, theta): 215 | ctx.save_for_backward(theta) 216 | return sinc3(theta) 217 | 218 | @staticmethod 219 | def backward(ctx, grad_output): 220 | theta, = ctx.saved_tensors 221 | grad_theta = None 222 | if ctx.needs_input_grad[0]: 223 | grad_theta = grad_output * sinc3_dt(theta).to(grad_output) 224 | return grad_theta 225 | 226 | Sinc3 = Sinc3_autograd.apply 227 | 228 | 229 | #EOF 230 | -------------------------------------------------------------------------------- /ptlk/so3.py: -------------------------------------------------------------------------------- 1 | """ 3-d rotation group and corresponding Lie algebra """ 2 | import torch 3 | from . import sinc 4 | from .sinc import sinc1, sinc2, sinc3 5 | 6 | 7 | def cross_prod(x, y): 8 | z = torch.cross(x.view(-1, 3), y.view(-1, 3), dim=1).view_as(x) 9 | return z 10 | 11 | def liebracket(x, y): 12 | return cross_prod(x, y) 13 | 14 | def mat(x): 15 | # size: [*, 3] -> [*, 3, 3] 16 | x_ = x.view(-1, 3) 17 | x1, x2, x3 = x_[:, 0], x_[:, 1], x_[:, 2] 18 | O = torch.zeros_like(x1) 19 | 20 | X = torch.stack(( 21 | torch.stack((O, -x3, x2), dim=1), 22 | torch.stack((x3, O, -x1), dim=1), 23 | torch.stack((-x2, x1, O), dim=1)), dim=1) 24 | return X.view(*(x.size()[0:-1]), 3, 3) 25 | 26 | def vec(X): 27 | X_ = X.view(-1, 3, 3) 28 | x1, x2, x3 = X_[:, 2, 1], X_[:, 0, 2], X_[:, 1, 0] 29 | x = torch.stack((x1, x2, x3), dim=1) 30 | return x.view(*X.size()[0:-2], 3) 31 | 32 | def genvec(): 33 | return torch.eye(3) 34 | 35 | def genmat(): 36 | return mat(genvec()) 37 | 38 | def RodriguesRotation(x): 39 | # for autograd 40 | w = x.view(-1, 3) 41 | t = w.norm(p=2, dim=1).view(-1, 1, 1) 42 | W = mat(w) 43 | S = W.bmm(W) 44 | I = torch.eye(3).to(w) 45 | 46 | # Rodrigues' rotation formula. 47 | #R = cos(t)*eye(3) + sinc1(t)*W + sinc2(t)*(w*w'); 48 | #R = eye(3) + sinc1(t)*W + sinc2(t)*S 49 | 50 | R = I + sinc.Sinc1(t)*W + sinc.Sinc2(t)*S 51 | 52 | return R.view(*(x.size()[0:-1]), 3, 3) 53 | 54 | def exp(x): 55 | w = x.view(-1, 3) 56 | t = w.norm(p=2, dim=1).view(-1, 1, 1) 57 | W = mat(w) 58 | S = W.bmm(W) 59 | I = torch.eye(3).to(w) 60 | 61 | # Rodrigues' rotation formula. 62 | #R = cos(t)*eye(3) + sinc1(t)*W + sinc2(t)*(w*w'); 63 | #R = eye(3) + sinc1(t)*W + sinc2(t)*S 64 | 65 | R = I + sinc1(t)*W + sinc2(t)*S 66 | 67 | return R.view(*(x.size()[0:-1]), 3, 3) 68 | 69 | def inverse(g): 70 | R = g.view(-1, 3, 3) 71 | Rt = R.transpose(1, 2) 72 | return Rt.view_as(g) 73 | 74 | def btrace(X): 75 | # batch-trace: [B, N, N] -> [B] 76 | n = X.size(-1) 77 | X_ = X.view(-1, n, n) 78 | tr = torch.zeros(X_.size(0)).to(X) 79 | for i in range(tr.size(0)): 80 | m = X_[i, :, :] 81 | tr[i] = torch.trace(m) 82 | return tr.view(*(X.size()[0:-2])) 83 | 84 | def log(g): 85 | eps = 1.0e-7 86 | R = g.view(-1, 3, 3) 87 | tr = btrace(R) 88 | c = (tr - 1) / 2 89 | t = torch.acos(c) 90 | sc = sinc1(t) 91 | idx0 = (torch.abs(sc) <= eps) 92 | idx1 = (torch.abs(sc) > eps) 93 | sc = sc.view(-1, 1, 1) 94 | 95 | X = torch.zeros_like(R) 96 | if idx1.any(): 97 | X[idx1] = (R[idx1] - R[idx1].transpose(1, 2)) / (2*sc[idx1]) 98 | 99 | if idx0.any(): 100 | # t[idx0] == math.pi 101 | t2 = t[idx0] ** 2 102 | A = (R[idx0] + torch.eye(3).type_as(R).unsqueeze(0)) * t2.view(-1, 1, 1) / 2 103 | aw1 = torch.sqrt(A[:, 0, 0]) 104 | aw2 = torch.sqrt(A[:, 1, 1]) 105 | aw3 = torch.sqrt(A[:, 2, 2]) 106 | sgn_3 = torch.sign(A[:, 0, 2]) 107 | sgn_3[sgn_3 == 0] = 1 108 | sgn_23 = torch.sign(A[:, 1, 2]) 109 | sgn_23[sgn_23 == 0] = 1 110 | sgn_2 = sgn_23 * sgn_3 111 | w1 = aw1 112 | w2 = aw2 * sgn_2 113 | w3 = aw3 * sgn_3 114 | w = torch.stack((w1, w2, w3), dim=-1) 115 | W = mat(w) 116 | X[idx0] = W 117 | 118 | x = vec(X.view_as(g)) 119 | return x 120 | 121 | def transform(g, a): 122 | # g in SO(3): * x 3 x 3 123 | # a in R^3: * x 3[x N] 124 | if len(g.size()) == len(a.size()): 125 | b = g.matmul(a) 126 | else: 127 | b = g.matmul(a.unsqueeze(-1)).squeeze(-1) 128 | return b 129 | 130 | def group_prod(g, h): 131 | # g, h : SO(3) 132 | g1 = g.matmul(h) 133 | return g1 134 | 135 | 136 | 137 | def vecs_Xg_ig(x): 138 | """ Vi = vec(dg/dxi * inv(g)), where g = exp(x) 139 | (== [Ad(exp(x))] * vecs_ig_Xg(x)) 140 | """ 141 | t = x.view(-1, 3).norm(p=2, dim=1).view(-1, 1, 1) 142 | X = mat(x) 143 | S = X.bmm(X) 144 | #B = x.view(-1,3,1).bmm(x.view(-1,1,3)) # B = x*x' 145 | I = torch.eye(3).to(X) 146 | 147 | #V = sinc1(t)*eye(3) + sinc2(t)*X + sinc3(t)*B 148 | #V = eye(3) + sinc2(t)*X + sinc3(t)*S 149 | 150 | V = I + sinc2(t)*X + sinc3(t)*S 151 | 152 | return V.view(*(x.size()[0:-1]), 3, 3) 153 | 154 | def inv_vecs_Xg_ig(x): 155 | """ H = inv(vecs_Xg_ig(x)) """ 156 | t = x.view(-1, 3).norm(p=2, dim=1).view(-1, 1, 1) 157 | X = mat(x) 158 | S = X.bmm(X) 159 | I = torch.eye(3).to(x) 160 | 161 | e = 0.01 162 | eta = torch.zeros_like(t) 163 | s = (t < e) 164 | c = (s == 0) 165 | t2 = t[s] ** 2 166 | eta[s] = ((t2/40 + 1)*t2/42 + 1)*t2/720 + 1/12 # O(t**8) 167 | eta[c] = (1 - (t[c]/2) / torch.tan(t[c]/2)) / (t[c]**2) 168 | 169 | H = I - 1/2*X + eta*S 170 | return H.view(*(x.size()[0:-1]), 3, 3) 171 | 172 | 173 | class ExpMap(torch.autograd.Function): 174 | """ Exp: so(3) -> SO(3) 175 | """ 176 | @staticmethod 177 | def forward(ctx, x): 178 | """ Exp: R^3 -> M(3), 179 | size: [B, 3] -> [B, 3, 3], 180 | or [B, 1, 3] -> [B, 1, 3, 3] 181 | """ 182 | ctx.save_for_backward(x) 183 | g = exp(x) 184 | return g 185 | 186 | @staticmethod 187 | def backward(ctx, grad_output): 188 | x, = ctx.saved_tensors 189 | g = exp(x) 190 | gen_k = genmat().to(x) 191 | #gen_1 = gen_k[0, :, :] 192 | #gen_2 = gen_k[1, :, :] 193 | #gen_3 = gen_k[2, :, :] 194 | 195 | # Let z = f(g) = f(exp(x)) 196 | # dz = df/dgij * dgij/dxk * dxk 197 | # = df/dgij * (d/dxk)[exp(x)]_ij * dxk 198 | # = df/dgij * [gen_k*g]_ij * dxk 199 | 200 | dg = gen_k.matmul(g.view(-1, 1, 3, 3)) 201 | # (k, i, j) 202 | dg = dg.to(grad_output) 203 | 204 | go = grad_output.contiguous().view(-1, 1, 3, 3) 205 | dd = go * dg 206 | grad_input = dd.sum(-1).sum(-1) 207 | 208 | return grad_input 209 | 210 | Exp = ExpMap.apply 211 | 212 | 213 | #EOF 214 | --------------------------------------------------------------------------------