├── FCGF_APR ├── LICENSE ├── README.md ├── config.py ├── config │ ├── file_LoKITTI_50.npy │ ├── file_LoNUSCENES_50.npy │ ├── test_kitti.txt │ ├── train_kitti.txt │ └── val_kitti.txt ├── lib │ ├── __init__.py │ ├── complement_data_loader.py │ ├── complement_trainer.py │ ├── data_loaders.py │ ├── eval.py │ ├── metrics.py │ ├── timer.py │ ├── trainer.py │ └── transforms.py ├── model │ ├── __init__.py │ ├── common.py │ ├── mlp.py │ ├── residual_block.py │ ├── resunet.py │ └── simpleunet.py ├── requirements.txt ├── scripts │ ├── benchmark_util.py │ ├── test_apr.py │ ├── test_apr_kitti.sh │ ├── test_apr_nuscenes.sh │ ├── test_fcgf.py │ ├── test_fcgf_kitti.sh │ ├── test_fcgf_nuscenes.sh │ ├── train_apr_kitti.sh │ ├── train_apr_nuscenes.sh │ ├── train_fcgf_kitti.sh │ └── train_fcgf_nuscenes.sh ├── train.py └── util │ ├── __init__.py │ ├── file.py │ ├── misc.py │ ├── pointcloud.py │ ├── trajectory.py │ ├── transform_estimation.py │ └── visualization.py ├── LICENSE ├── Predator_APR ├── LICENSE ├── README.md ├── common │ ├── colors.py │ ├── math │ │ ├── random.py │ │ ├── se3.py │ │ └── so3.py │ ├── math_torch │ │ └── se3.py │ ├── misc.py │ └── torch.py ├── configs │ ├── kitti │ │ ├── file_LoKITTI_40.npy │ │ ├── file_LoKITTI_50.npy │ │ ├── test_kitti.txt │ │ ├── train_kitti.txt │ │ └── val_kitti.txt │ ├── models.py │ ├── nuscenes │ │ ├── file_LoNUSCENES_40.npy │ │ └── file_LoNUSCENES_50.npy │ ├── test │ │ ├── kitti.yaml │ │ └── nuscenes.yaml │ └── train │ │ ├── kitti.yaml │ │ └── nuscenes.yaml ├── cpp_wrappers │ ├── compile_wrappers.sh │ ├── cpp_neighbors │ │ ├── build.bat │ │ ├── neighbors │ │ │ ├── neighbors.cpp │ │ │ └── neighbors.h │ │ ├── setup.py │ │ └── wrapper.cpp │ ├── cpp_subsampling │ │ ├── build.bat │ │ ├── grid_subsampling │ │ │ ├── grid_subsampling.cpp │ │ │ └── grid_subsampling.h │ │ ├── setup.py │ │ └── wrapper.cpp │ └── cpp_utils │ │ ├── cloud │ │ ├── cloud.cpp │ │ └── cloud.h │ │ └── nanoflann │ │ └── nanoflann.hpp ├── datasets │ ├── __init__.py │ ├── dataloader.py │ ├── indoor.py │ ├── kitti.py │ ├── modelnet.py │ ├── nuscenes.py │ └── transforms.py ├── kernels │ ├── dispositions │ │ └── k_015_center_3D.ply │ └── kernel_points.py ├── lib │ ├── __init__.py │ ├── benchmark.py │ ├── benchmark_utils.py │ ├── loss.py │ ├── ply.py │ ├── tester.py │ ├── timer.py │ ├── trainer.py │ └── utils.py ├── main.py ├── models │ ├── __init__.py │ ├── architectures.py │ ├── blocks.py │ ├── gcn.py │ └── mlp.py ├── requirements.txt └── scripts │ └── cal_overlap.py ├── README.md └── resources ├── 1683269100777.png ├── 1683269374754.png └── export_kitti_minimal.py /FCGF_APR/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of 4 | this software and associated documentation files (the "Software"), to deal in 5 | the Software without restriction, including without limitation the rights to 6 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 7 | of the Software, and to permit persons to whom the Software is furnished to do 8 | so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in all 11 | copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | SOFTWARE. 20 | -------------------------------------------------------------------------------- /FCGF_APR/README.md: -------------------------------------------------------------------------------- 1 | # APR: Online Distant Point Cloud Registration Through Aggregated Point Cloud Reconstruction, Implemented with Fully Convolutional Geometric Features ([FCGF](https://github.com/chrischoy/FCGF)) 2 | 3 | For many driving safety applications, it is of great importance to accurately register LiDAR point clouds generated on distant moving vehicles. However, such point clouds have extremely different point density and sensor perspective on the same object, making registration on such point clouds very hard. In this paper, we propose a novel feature extraction framework, called APR, for online distant point cloud registration. Specifically, APR leverages an autoencoder design, where the autoencoder reconstructs a denser aggregated point cloud with several frames instead of the original single input point cloud. Our design forces the encoder to extract features with rich local geometry information based on one single input point cloud. Such features are then used for online distant point cloud registration. We conduct extensive experiments against state-of-the-art (SOTA) feature extractors on KITTI and nuScenes datasets. Results show that APR outperforms all other extractors by a large margin, increasing average registration recall of SOTA extractors by 7.1% on LoKITTI and 4.6% on LoNuScenes. 4 | 5 | This repository is an implementation of APR using FCGF as the feature extractor. 6 | 7 | ## Requirements 8 | 9 | - Ubuntu 14.04 or higher 10 | - CUDA 11.1 or higher 11 | - Python v3.7 or higher 12 | - Pytorch v1.6 or higher 13 | - [MinkowskiEngine](https://github.com/stanfordvl/MinkowskiEngine) v0.5 or higher 14 | 15 | ## Training & Testing 16 | 17 | First, use FCGF_APR as the working directory: 18 | 19 | ``` 20 | cd ./FCGF_APR 21 | conda activate apr 22 | ``` 23 | 24 | If you haven't downloaded any of the datasets (KITTI and nuScenes) according to our specification, please refer to the README.md in the [parent directory](../README.md). 25 | 26 | ### Setting the distance between two LiDARs (registration difficulty) 27 | 28 | As the major focus of this paper, we divide the registration datasets into different slices according to the distance $d$ between two LiDARs. Greater $d$ leads to a smaller overlap and more divergent point density, resulting in a higher registration difficulty. We denote range of $d$ with the parameter `--pair_min_dist` and `--pair_max_dist`, which can be found in `./scripts/train_{$method}_{$dataset}.sh`. For example, setting 29 | 30 | ``` 31 | --pair_min_dist 5 \ 32 | --pair_max_dist 20 \ 33 | ``` 34 | 35 | will set $d\in [5m,20m]$. In other words, for every pair of point clouds, the ground-truth euclidean distance betwen two corresponding LiDAR positions (i.e., the origins of the two specified point clouds) obeys a uniform distribution between 5m and 20m. 36 | 37 | ### Training suggestions 38 | 39 | For cases where you want `--pair_max_dist` to be larger than 20, we recommend following the two-stage training paradigm as pointed out in Section 5 of our paper: 40 | 41 | 1. Pretrain a model with the following distance parameters: `--pair_min_dist 5 --pair_max_dist 20`. Record the pretrained model path that is printed at the beginning of the training. It shoud be some path like this: `./outputs/Experiments/PairComplementKittiDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout32/YYYY-MM-DD_HH-MM-SS` 42 | 2. Finetune a new model on `--pair_min_dist 5 --pair_max_dist {$YOUR_SPECIFIC_DISTANCE}`, and paste the pretrained model path to `--resume "{$PRETRAINED_PATH}/chechpoint.pth"` and `--resume_dir "{$PRETRAINED_PATH}"`. Do not forget to set `--finetune_restart true`. 43 | 44 | Emperically, the pretraining strategy helps a lot in model convergence especially when the distance is large; Otherwise the model just diverges. 45 | 46 | ### Launch the training 47 | 48 | Notes: 49 | 50 | 1. Remember to set `--use_old_pose` to true when using the nuScenes dataset. 51 | 2. The symmetric APR setup can be enabled by setting `--symmetric` to True. 52 | 53 | To train FCGF-APR on either dataset, run either of the following command inside conda environment `apr`: 54 | 55 | ``` 56 | ./scripts/train_apr_kitti.sh 57 | ./scripts/train_apr_nuscenes.sh 58 | ``` 59 | 60 | The baseline method FCGF can be trained similarly with our dataset: 61 | 62 | ``` 63 | ./scripts/train_fcgf_kitti.sh 64 | ./scripts/train_fcgf_nuscenes.sh 65 | ``` 66 | 67 | ### Testing 68 | 69 | To test FCGF-APR on either dataset, set `OUT_DIR` to the specific model path before running the corresponding script inside conda environment `apr`: 70 | 71 | ``` 72 | ./scripts/test_apr_kitti.sh 73 | ./scripts/test_apr_nuscenes.sh 74 | ``` 75 | 76 | The baseline method FCGF can be tested similarly: 77 | 78 | ``` 79 | ./scripts/test_fcgf_kitti.sh 80 | ./scripts/test_fcgf_nuscenes.sh 81 | ``` 82 | 83 | ## Pre-trained models 84 | 85 | We provide our [model](https://drive.google.com/file/d/1mLqiahQMgYMRyB4XKhp-HJdy5yavL2fj/view?usp=sharing) pretrained on FCGF+APR, with different point cloud distance. Extract the file into the 'outputs' directory and constitute './outputs/some_model' into the 'OUT_DIR' section in the test scripts to reproduce the results showed in our paper. 86 | -------------------------------------------------------------------------------- /FCGF_APR/config.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | arg_lists = [] 4 | parser = argparse.ArgumentParser() 5 | 6 | 7 | def add_argument_group(name): 8 | arg = parser.add_argument_group(name) 9 | arg_lists.append(arg) 10 | return arg 11 | 12 | 13 | def str2bool(v): 14 | return v.lower() in ('true', '1') 15 | 16 | 17 | logging_arg = add_argument_group('Logging') 18 | logging_arg.add_argument('--out_dir', type=str, default='outputs') 19 | 20 | trainer_arg = add_argument_group('Trainer') 21 | trainer_arg.add_argument('--trainer', type=str, default='HardestContrastiveLossTrainer') 22 | trainer_arg.add_argument('--save_freq_epoch', type=int, default=1) 23 | trainer_arg.add_argument('--batch_size', type=int, default=4) 24 | trainer_arg.add_argument('--val_batch_size', type=int, default=1) 25 | 26 | # Hard negative mining 27 | trainer_arg.add_argument('--use_hard_negative', type=str2bool, default=True) 28 | trainer_arg.add_argument('--hard_negative_sample_ratio', type=int, default=0.05) 29 | trainer_arg.add_argument('--hard_negative_max_num', type=int, default=3000) 30 | trainer_arg.add_argument('--num_pos_per_batch', type=int, default=1024) 31 | trainer_arg.add_argument('--num_hn_samples_per_batch', type=int, default=256) 32 | 33 | # Metric learning loss 34 | trainer_arg.add_argument('--neg_thresh', type=float, default=1.4) 35 | trainer_arg.add_argument('--pos_thresh', type=float, default=0.1) 36 | trainer_arg.add_argument('--neg_weight', type=float, default=1) 37 | 38 | # Data augmentation 39 | trainer_arg.add_argument('--use_random_scale', type=str2bool, default=False) 40 | trainer_arg.add_argument('--min_scale', type=float, default=0.8) 41 | trainer_arg.add_argument('--max_scale', type=float, default=1.2) 42 | trainer_arg.add_argument('--use_random_rotation', type=str2bool, default=True) 43 | trainer_arg.add_argument('--rotation_range', type=float, default=360) 44 | 45 | # Data loader configs 46 | trainer_arg.add_argument('--train_phase', type=str, default="train") 47 | trainer_arg.add_argument('--val_phase', type=str, default="val") 48 | trainer_arg.add_argument('--test_phase', type=str, default="test") 49 | 50 | trainer_arg.add_argument('--stat_freq', type=int, default=40) 51 | trainer_arg.add_argument('--test_valid', type=str2bool, default=True) 52 | trainer_arg.add_argument('--val_max_iter', type=int, default=400) 53 | trainer_arg.add_argument('--val_epoch_freq', type=int, default=1) 54 | trainer_arg.add_argument( 55 | '--positive_pair_search_voxel_size_multiplier', type=float, default=1.5) 56 | 57 | trainer_arg.add_argument('--hit_ratio_thresh', type=float, default=0.1) 58 | trainer_arg.add_argument('--min_sample_frame_dist', type=float, default=10.0) 59 | trainer_arg.add_argument('--complement_pair_dist', type=float, default=10.0) 60 | trainer_arg.add_argument('--num_complement_one_side', type=int, default=5) 61 | 62 | # Triplets 63 | trainer_arg.add_argument('--triplet_num_pos', type=int, default=256) 64 | trainer_arg.add_argument('--triplet_num_hn', type=int, default=512) 65 | trainer_arg.add_argument('--triplet_num_rand', type=int, default=1024) 66 | 67 | # dNetwork specific configurations 68 | net_arg = add_argument_group('Network') 69 | net_arg.add_argument('--model', type=str, default='ResUNetFatBN') 70 | net_arg.add_argument('--encoder_model', type=str, default='ResUNetFatBN') 71 | net_arg.add_argument('--generator_model', type=str, default='ResUNetBN2C') 72 | net_arg.add_argument('--model_n_out', type=int, default=32, help='Feature dimension') 73 | net_arg.add_argument('--conv1_kernel_size', type=int, default=5) 74 | net_arg.add_argument('--normalize_feature', type=str2bool, default=True) 75 | net_arg.add_argument('--dist_type', type=str, default='L2') 76 | net_arg.add_argument('--best_val_metric', type=str, default='feat_match_ratio') 77 | net_arg.add_argument('--point_generation_ratio', type=int, default=6) 78 | net_arg.add_argument('--regularization_strength', type=float, default=0.1) 79 | net_arg.add_argument('--regularization_type', type=str, default='L2') 80 | 81 | # Optimizer arguments 82 | opt_arg = add_argument_group('Optimizer') 83 | opt_arg.add_argument('--optimizer', type=str, default='SGD') 84 | opt_arg.add_argument('--max_epoch', type=int, default=100) 85 | opt_arg.add_argument('--lr', type=float, default=1e-1) 86 | opt_arg.add_argument('--loss_ratio', type=float, default=1e-5) 87 | opt_arg.add_argument('--momentum', type=float, default=0.8) 88 | opt_arg.add_argument('--sgd_momentum', type=float, default=0.9) 89 | opt_arg.add_argument('--sgd_dampening', type=float, default=0.1) 90 | opt_arg.add_argument('--adam_beta1', type=float, default=0.9) 91 | opt_arg.add_argument('--adam_beta2', type=float, default=0.999) 92 | opt_arg.add_argument('--weight_decay', type=float, default=1e-4) 93 | opt_arg.add_argument('--iter_size', type=int, default=1, help='accumulate gradient') 94 | opt_arg.add_argument('--bn_momentum', type=float, default=0.05) 95 | opt_arg.add_argument('--exp_gamma', type=float, default=0.99) 96 | opt_arg.add_argument('--scheduler', type=str, default='ExpLR') 97 | opt_arg.add_argument( 98 | '--icp_cache_path', type=str, default="/home/chrischoy/datasets/FCGF/kitti/icp/") 99 | 100 | misc_arg = add_argument_group('Misc') 101 | misc_arg.add_argument('--use_gpu', type=str2bool, default=True) 102 | misc_arg.add_argument('--weights', type=str, default=None) 103 | misc_arg.add_argument('--weights_dir', type=str, default=None) 104 | misc_arg.add_argument('--resume', type=str, default=None) 105 | misc_arg.add_argument('--resume_dir', type=str, default=None) 106 | misc_arg.add_argument('--train_num_thread', type=int, default=2) 107 | misc_arg.add_argument('--val_num_thread', type=int, default=1) 108 | misc_arg.add_argument('--test_num_thread', type=int, default=2) 109 | misc_arg.add_argument('--fast_validation', type=str2bool, default=False) 110 | misc_arg.add_argument( 111 | '--nn_max_n', 112 | type=int, 113 | default=500, 114 | help='The maximum number of features to find nearest neighbors in batch') 115 | 116 | # Dataset specific configurations 117 | data_arg = add_argument_group('Data') 118 | data_arg.add_argument('--dataset', type=str, default='ThreeDMatchPairDataset') 119 | data_arg.add_argument('--voxel_size', type=float, default=0.025) 120 | data_arg.add_argument( 121 | '--threed_match_dir', type=str, default="/home/chrischoy/datasets/FCGF/threedmatch") 122 | data_arg.add_argument( 123 | '--kitti_root', type=str, default="/home/chrischoy/datasets/FCGF/kitti/") 124 | data_arg.add_argument( 125 | '--kitti_max_time_diff', 126 | type=int, 127 | default=3, 128 | help='max time difference between pairs (non inclusive)') 129 | data_arg.add_argument('--kitti_date', type=str, default='2011_09_26') 130 | data_arg.add_argument('--pair_min_dist', type=int, default=-1) 131 | data_arg.add_argument('--pair_max_dist', type=int, default=-1) 132 | data_arg.add_argument('--mutate_neighbour_percentage', type=float, default=0.) 133 | data_arg.add_argument('--LoKITTI', type=str2bool, default=False) 134 | 135 | # Debug configurations, placed in order to find the best hyperparameters 136 | debug_arg = add_argument_group('Debug') 137 | debug_arg.add_argument('--use_old_pose', type=str2bool, default=True) 138 | debug_arg.add_argument('--debug_force_icp_recalculation', type=str2bool, default=False) 139 | debug_arg.add_argument('--debug_use_old_complement', type=str2bool, default=False) 140 | debug_arg.add_argument('--finetune_restart', type=str2bool, default=False) 141 | debug_arg.add_argument('--symmetric', type=str2bool, default=False) 142 | 143 | def get_config(): 144 | args = parser.parse_args() 145 | return args 146 | -------------------------------------------------------------------------------- /FCGF_APR/config/file_LoKITTI_50.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/FCGF_APR/config/file_LoKITTI_50.npy -------------------------------------------------------------------------------- /FCGF_APR/config/file_LoNUSCENES_50.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/FCGF_APR/config/file_LoNUSCENES_50.npy -------------------------------------------------------------------------------- /FCGF_APR/config/test_kitti.txt: -------------------------------------------------------------------------------- 1 | 8 2 | 9 3 | 10 4 | -------------------------------------------------------------------------------- /FCGF_APR/config/train_kitti.txt: -------------------------------------------------------------------------------- 1 | 0 2 | 1 3 | 2 4 | 3 5 | 4 6 | 5 7 | -------------------------------------------------------------------------------- /FCGF_APR/config/val_kitti.txt: -------------------------------------------------------------------------------- 1 | 6 2 | 7 3 | -------------------------------------------------------------------------------- /FCGF_APR/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/FCGF_APR/lib/__init__.py -------------------------------------------------------------------------------- /FCGF_APR/lib/eval.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import open3d as o3d 4 | 5 | from lib.metrics import pdist 6 | from scipy.spatial import cKDTree 7 | 8 | 9 | def find_nn_cpu(feat0, feat1, return_distance=False): 10 | feat1tree = cKDTree(feat1) 11 | dists, nn_inds = feat1tree.query(feat0, k=1, n_jobs=-1) 12 | if return_distance: 13 | return nn_inds, dists 14 | else: 15 | return nn_inds 16 | 17 | 18 | def find_nn_gpu(F0, F1, nn_max_n=-1, return_distance=False, dist_type='SquareL2'): 19 | # Too much memory if F0 or F1 large. Divide the F0 20 | if nn_max_n > 1: 21 | N = len(F0) 22 | C = int(np.ceil(N / nn_max_n)) 23 | stride = nn_max_n 24 | dists, inds = [], [] 25 | for i in range(C): 26 | dist = pdist(F0[i * stride:(i + 1) * stride], F1, dist_type=dist_type) 27 | min_dist, ind = dist.min(dim=1) 28 | dists.append(min_dist.detach().unsqueeze(1).cpu()) 29 | inds.append(ind.cpu()) 30 | 31 | if C * stride < N: 32 | dist = pdist(F0[C * stride:], F1, dist_type=dist_type) 33 | min_dist, ind = dist.min(dim=1) 34 | dists.append(min_dist.detach().unsqueeze(1).cpu()) 35 | inds.append(ind.cpu()) 36 | 37 | dists = torch.cat(dists) 38 | inds = torch.cat(inds) 39 | assert len(inds) == N 40 | else: 41 | dist = pdist(F0, F1, dist_type=dist_type) 42 | min_dist, inds = dist.min(dim=1) 43 | dists = min_dist.detach().unsqueeze(1).cpu() 44 | inds = inds.cpu() 45 | if return_distance: 46 | return inds, dists 47 | else: 48 | return inds 49 | -------------------------------------------------------------------------------- /FCGF_APR/lib/metrics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch 4 | import torch.functional as F 5 | 6 | 7 | def eval_metrics(output, target): 8 | output = (F.sigmoid(output) > 0.5).cpu().data.numpy() 9 | target = target.cpu().data.numpy() 10 | return np.linalg.norm(output - target) 11 | 12 | 13 | def corr_dist(est, gth, xyz0, xyz1, weight=None, max_dist=1): 14 | xyz0_est = xyz0 @ est[:3, :3].t() + est[:3, 3] 15 | xyz0_gth = xyz0 @ gth[:3, :3].t() + gth[:3, 3] 16 | dists = torch.clamp(torch.sqrt(((xyz0_est - xyz0_gth).pow(2)).sum(1)), max=max_dist) 17 | if weight is not None: 18 | dists = weight * dists 19 | return dists.mean() 20 | 21 | 22 | def pdist(A, B, dist_type='L2'): 23 | if dist_type == 'L2': 24 | D2 = torch.sum((A.unsqueeze(1) - B.unsqueeze(0)).pow(2), 2) 25 | return torch.sqrt(D2 + 1e-7) 26 | elif dist_type == 'SquareL2': 27 | return torch.sum((A.unsqueeze(1) - B.unsqueeze(0)).pow(2), 2) 28 | else: 29 | raise NotImplementedError('Not implemented') 30 | 31 | 32 | def get_loss_fn(loss): 33 | if loss == 'corr_dist': 34 | return corr_dist 35 | else: 36 | raise ValueError(f'Loss {loss}, not defined') 37 | -------------------------------------------------------------------------------- /FCGF_APR/lib/timer.py: -------------------------------------------------------------------------------- 1 | import math 2 | import time 3 | 4 | 5 | class AverageMeter(object): 6 | """Computes and stores the average and current value""" 7 | 8 | def __init__(self): 9 | self.reset() 10 | 11 | def reset(self): 12 | self.val = 0 13 | self.avg = 0 14 | self.sum = 0.0 15 | self.sq_sum = 0.0 16 | self.count = 0 17 | 18 | def update(self, val, n=1): 19 | self.val = val 20 | self.sum += val * n 21 | self.count += n 22 | self.avg = self.sum / self.count 23 | self.sq_sum += val**2 * n 24 | self.var = self.sq_sum / self.count - self.avg**2 25 | 26 | 27 | class Timer(object): 28 | """A simple timer.""" 29 | 30 | def __init__(self, binary_fn=None, init_val=0): 31 | self.total_time = 0. 32 | self.calls = 0 33 | self.start_time = 0. 34 | self.diff = 0. 35 | self.binary_fn = binary_fn 36 | self.tmp = init_val 37 | 38 | def reset(self): 39 | self.total_time = 0 40 | self.calls = 0 41 | self.start_time = 0 42 | self.diff = 0 43 | 44 | @property 45 | def avg(self): 46 | return self.total_time / self.calls 47 | 48 | def tic(self): 49 | # using time.time instead of time.clock because time time.clock 50 | # does not normalize for multithreading 51 | self.start_time = time.time() 52 | 53 | def toc(self, average=True, accumulate=False): 54 | self.diff = time.time() - self.start_time 55 | self.total_time += self.diff 56 | if not accumulate: 57 | self.calls += 1 58 | if self.binary_fn: 59 | self.tmp = self.binary_fn(self.tmp, self.diff) 60 | if average and not accumulate: 61 | return self.avg 62 | else: 63 | return self.diff 64 | 65 | def incCount(self): 66 | self.calls += 1 67 | 68 | 69 | class MinTimer(Timer): 70 | 71 | def __init__(self): 72 | Timer.__init__(self, binary_fn=lambda x, y: min(x, y), init_val=math.inf) 73 | 74 | @property 75 | def min(self): 76 | return self.tmp 77 | -------------------------------------------------------------------------------- /FCGF_APR/lib/transforms.py: -------------------------------------------------------------------------------- 1 | import random 2 | import numpy as np 3 | 4 | import torch 5 | 6 | 7 | class Compose: 8 | 9 | def __init__(self, transforms): 10 | self.transforms = transforms 11 | 12 | def __call__(self, coords, feats): 13 | for transform in self.transforms: 14 | coords, feats = transform(coords, feats) 15 | return coords, feats 16 | 17 | 18 | class Jitter: 19 | 20 | def __init__(self, mu=0, sigma=0.01): 21 | self.mu = mu 22 | self.sigma = sigma 23 | 24 | def __call__(self, coords, feats): 25 | if random.random() < 0.95: 26 | if isinstance(feats, np.ndarray): 27 | feats += np.random.normal(self.mu, self.sigma, (feats.shape[0], feats.shape[1])) 28 | else: 29 | feats += (torch.randn_like(feats) * self.sigma) + self.mu 30 | return coords, feats 31 | 32 | 33 | class ChromaticShift: 34 | 35 | def __init__(self, mu=0, sigma=0.1): 36 | self.mu = mu 37 | self.sigma = sigma 38 | 39 | def __call__(self, coords, feats): 40 | if random.random() < 0.95: 41 | feats[:, :3] += np.random.normal(self.mu, self.sigma, (1, 3)) 42 | return coords, feats 43 | -------------------------------------------------------------------------------- /FCGF_APR/model/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import model.simpleunet as simpleunets 3 | import model.resunet as resunets 4 | import model.mlp as mlp 5 | 6 | MODELS = [] 7 | 8 | 9 | def add_models(module): 10 | MODELS.extend([getattr(module, a) for a in dir(module) if 'Net' in a or 'MLP' in a]) 11 | 12 | 13 | add_models(simpleunets) 14 | add_models(resunets) 15 | add_models(mlp) 16 | 17 | 18 | def load_model(name): 19 | '''Creates and returns an instance of the model given its class name. 20 | ''' 21 | # Find the model class from its name 22 | all_models = MODELS 23 | mdict = {model.__name__: model for model in all_models} 24 | if name not in mdict: 25 | logging.info(f'Invalid model index. You put {name}. Options are:') 26 | # Display a list of valid model names 27 | for model in all_models: 28 | logging.info('\t* {}'.format(model.__name__)) 29 | return None 30 | NetClass = mdict[name] 31 | 32 | return NetClass 33 | -------------------------------------------------------------------------------- /FCGF_APR/model/common.py: -------------------------------------------------------------------------------- 1 | import MinkowskiEngine as ME 2 | 3 | 4 | def get_norm(norm_type, num_feats, bn_momentum=0.05, D=-1): 5 | if norm_type == 'BN': 6 | return ME.MinkowskiBatchNorm(num_feats, momentum=bn_momentum) 7 | elif norm_type == 'IN': 8 | return ME.MinkowskiInstanceNorm(num_feats, dimension=D) 9 | else: 10 | raise ValueError(f'Type {norm_type}, not defined') 11 | -------------------------------------------------------------------------------- /FCGF_APR/model/mlp.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | # import MinkowskiEngine as ME 3 | # import MinkowskiEngine.MinkowskiFunctional as MEF 4 | # from model.common import get_norm 5 | 6 | class GenerativeMLP(nn.Module): 7 | CHANNELS = [None, 512, 128, None] 8 | 9 | def __init__(self, 10 | in_channel=125, 11 | out_points=6, 12 | bn_momentum=0.1): 13 | super().__init__() 14 | CHANNELS = self.CHANNELS 15 | self.mlp = nn.Sequential( 16 | nn.Linear(in_channel, CHANNELS[1]), 17 | nn.ReLU(), 18 | nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 19 | nn.Linear(CHANNELS[1], CHANNELS[2]), 20 | nn.ReLU(), 21 | nn.BatchNorm1d(CHANNELS[2], momentum=bn_momentum), 22 | nn.Linear(CHANNELS[2], out_points*3), 23 | nn.ReLU() 24 | ) 25 | 26 | def forward(self, x): 27 | y = self.mlp(x) 28 | # print(y) 29 | return y 30 | 31 | 32 | # class GenerativeMLP_99(GenerativeMLP): 33 | # CHANNELS = [None, 512, 512, None] 34 | 35 | 36 | class GenerativeMLP_98(GenerativeMLP): 37 | CHANNELS = [None, 512, 256, None] 38 | 39 | 40 | class GenerativeMLP_54(GenerativeMLP): 41 | CHANNELS = [None, 32, 16, None] 42 | 43 | 44 | class GenerativeMLP_4(nn.Module): 45 | CHANNELS = [None, 16, None] 46 | 47 | def __init__(self, 48 | in_channel=125, 49 | out_points=6, 50 | bn_momentum=0.1): 51 | super().__init__() 52 | CHANNELS = self.CHANNELS 53 | self.mlp = nn.Sequential( 54 | nn.Linear(in_channel, CHANNELS[1]), 55 | nn.ReLU(), 56 | nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 57 | nn.Linear(CHANNELS[1], out_points*3), 58 | nn.ReLU() 59 | ) 60 | 61 | def forward(self, x): 62 | y = self.mlp(x) 63 | # print(y) 64 | return y 65 | 66 | 67 | class GenerativeMLP_11_10_9(nn.Module): 68 | CHANNELS = [None, 2048, 1024, 512, None] 69 | 70 | def __init__(self, 71 | in_channel=125, 72 | out_points=6, 73 | bn_momentum=0.1): 74 | super().__init__() 75 | CHANNELS = self.CHANNELS 76 | self.mlp = nn.Sequential( 77 | nn.Linear(in_channel, CHANNELS[1]), 78 | nn.ReLU(), 79 | nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 80 | nn.Linear(CHANNELS[1], CHANNELS[2]), 81 | nn.ReLU(), 82 | nn.BatchNorm1d(CHANNELS[2], momentum=bn_momentum), 83 | nn.Linear(CHANNELS[2], CHANNELS[3]), 84 | nn.ReLU(), 85 | nn.BatchNorm1d(CHANNELS[3], momentum=bn_momentum), 86 | nn.Linear(CHANNELS[3], out_points*3), 87 | nn.ReLU() 88 | ) 89 | 90 | def forward(self, x): 91 | y = self.mlp(x) 92 | # print(y) 93 | return y -------------------------------------------------------------------------------- /FCGF_APR/model/residual_block.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | from model.common import get_norm 4 | 5 | import MinkowskiEngine as ME 6 | import MinkowskiEngine.MinkowskiFunctional as MEF 7 | 8 | 9 | class BasicBlockBase(nn.Module): 10 | expansion = 1 11 | NORM_TYPE = 'BN' 12 | 13 | def __init__(self, 14 | inplanes, 15 | planes, 16 | stride=1, 17 | dilation=1, 18 | downsample=None, 19 | bn_momentum=0.1, 20 | D=3): 21 | super(BasicBlockBase, self).__init__() 22 | 23 | self.conv1 = ME.MinkowskiConvolution( 24 | inplanes, planes, kernel_size=3, stride=stride, dimension=D) 25 | self.norm1 = get_norm(self.NORM_TYPE, planes, bn_momentum=bn_momentum, D=D) 26 | self.conv2 = ME.MinkowskiConvolution( 27 | planes, 28 | planes, 29 | kernel_size=3, 30 | stride=1, 31 | dilation=dilation, 32 | bias=False, 33 | dimension=D) 34 | self.norm2 = get_norm(self.NORM_TYPE, planes, bn_momentum=bn_momentum, D=D) 35 | self.downsample = downsample 36 | 37 | def forward(self, x): 38 | residual = x 39 | 40 | out = self.conv1(x) 41 | out = self.norm1(out) 42 | out = MEF.relu(out) 43 | 44 | out = self.conv2(out) 45 | out = self.norm2(out) 46 | 47 | if self.downsample is not None: 48 | residual = self.downsample(x) 49 | 50 | out += residual 51 | out = MEF.relu(out) 52 | 53 | return out 54 | 55 | 56 | class BasicBlockBN(BasicBlockBase): 57 | NORM_TYPE = 'BN' 58 | 59 | 60 | class BasicBlockIN(BasicBlockBase): 61 | NORM_TYPE = 'IN' 62 | 63 | 64 | def get_block(norm_type, 65 | inplanes, 66 | planes, 67 | stride=1, 68 | dilation=1, 69 | downsample=None, 70 | bn_momentum=0.1, 71 | D=3): 72 | if norm_type == 'BN': 73 | return BasicBlockBN(inplanes, planes, stride, dilation, downsample, bn_momentum, D) 74 | elif norm_type == 'IN': 75 | return BasicBlockIN(inplanes, planes, stride, dilation, downsample, bn_momentum, D) 76 | else: 77 | raise ValueError(f'Type {norm_type}, not defined') 78 | -------------------------------------------------------------------------------- /FCGF_APR/model/resunet.py: -------------------------------------------------------------------------------- 1 | # -*- coding: future_fstrings -*- 2 | import torch 3 | import MinkowskiEngine as ME 4 | import MinkowskiEngine.MinkowskiFunctional as MEF 5 | from model.common import get_norm 6 | 7 | from model.residual_block import get_block 8 | 9 | 10 | class ResUNet2(ME.MinkowskiNetwork): 11 | NORM_TYPE = None 12 | BLOCK_NORM_TYPE = 'BN' 13 | CHANNELS = [None, 32, 64, 128, 256] 14 | TR_CHANNELS = [None, 32, 64, 64, 128] 15 | 16 | # To use the model, must call initialize_coords before forward pass. 17 | # Once data is processed, call clear to reset the model before calling initialize_coords 18 | def __init__(self, 19 | in_channels=3, 20 | out_channels=32, 21 | bn_momentum=0.1, 22 | normalize_feature=None, 23 | conv1_kernel_size=None, 24 | D=3): 25 | ME.MinkowskiNetwork.__init__(self, D) 26 | NORM_TYPE = self.NORM_TYPE 27 | BLOCK_NORM_TYPE = self.BLOCK_NORM_TYPE 28 | CHANNELS = self.CHANNELS 29 | TR_CHANNELS = self.TR_CHANNELS 30 | self.normalize_feature = normalize_feature 31 | self.conv1 = ME.MinkowskiConvolution( 32 | in_channels=in_channels, 33 | out_channels=CHANNELS[1], 34 | kernel_size=conv1_kernel_size, 35 | stride=1, 36 | dilation=1, 37 | bias=False, 38 | dimension=D) 39 | self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) 40 | 41 | self.block1 = get_block( 42 | BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, D=D) 43 | 44 | self.conv2 = ME.MinkowskiConvolution( 45 | in_channels=CHANNELS[1], 46 | out_channels=CHANNELS[2], 47 | kernel_size=3, 48 | stride=2, 49 | dilation=1, 50 | bias=False, 51 | dimension=D) 52 | self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) 53 | 54 | self.block2 = get_block( 55 | BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, D=D) 56 | 57 | self.conv3 = ME.MinkowskiConvolution( 58 | in_channels=CHANNELS[2], 59 | out_channels=CHANNELS[3], 60 | kernel_size=3, 61 | stride=2, 62 | dilation=1, 63 | bias=False, 64 | dimension=D) 65 | self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) 66 | 67 | self.block3 = get_block( 68 | BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, D=D) 69 | 70 | self.conv4 = ME.MinkowskiConvolution( 71 | in_channels=CHANNELS[3], 72 | out_channels=CHANNELS[4], 73 | kernel_size=3, 74 | stride=2, 75 | dilation=1, 76 | bias=False, 77 | dimension=D) 78 | self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D) 79 | 80 | self.block4 = get_block( 81 | BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, D=D) 82 | 83 | self.conv4_tr = ME.MinkowskiConvolutionTranspose( 84 | in_channels=CHANNELS[4], 85 | out_channels=TR_CHANNELS[4], 86 | kernel_size=3, 87 | stride=2, 88 | dilation=1, 89 | bias=False, 90 | dimension=D) 91 | self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) 92 | 93 | self.block4_tr = get_block( 94 | BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) 95 | 96 | self.conv3_tr = ME.MinkowskiConvolutionTranspose( 97 | in_channels=CHANNELS[3] + TR_CHANNELS[4], 98 | out_channels=TR_CHANNELS[3], 99 | kernel_size=3, 100 | stride=2, 101 | dilation=1, 102 | bias=False, 103 | dimension=D) 104 | self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) 105 | 106 | self.block3_tr = get_block( 107 | BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) 108 | 109 | self.conv2_tr = ME.MinkowskiConvolutionTranspose( 110 | in_channels=CHANNELS[2] + TR_CHANNELS[3], 111 | out_channels=TR_CHANNELS[2], 112 | kernel_size=3, 113 | stride=2, 114 | dilation=1, 115 | bias=False, 116 | dimension=D) 117 | self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) 118 | 119 | self.block2_tr = get_block( 120 | BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) 121 | 122 | self.conv1_tr = ME.MinkowskiConvolution( 123 | in_channels=CHANNELS[1] + TR_CHANNELS[2], 124 | out_channels=TR_CHANNELS[1], 125 | kernel_size=1, 126 | stride=1, 127 | dilation=1, 128 | bias=False, 129 | dimension=D) 130 | 131 | # self.block1_tr = BasicBlockBN(TR_CHANNELS[1], TR_CHANNELS[1], bn_momentum=bn_momentum, D=D) 132 | 133 | self.final = ME.MinkowskiConvolution( 134 | in_channels=TR_CHANNELS[1], 135 | out_channels=out_channels, 136 | kernel_size=1, 137 | stride=1, 138 | dilation=1, 139 | bias=True, 140 | dimension=D) 141 | 142 | def forward(self, x): 143 | out_s1 = self.conv1(x) 144 | out_s1 = self.norm1(out_s1) 145 | out_s1 = self.block1(out_s1) 146 | out = MEF.relu(out_s1) 147 | 148 | out_s2 = self.conv2(out) 149 | out_s2 = self.norm2(out_s2) 150 | out_s2 = self.block2(out_s2) 151 | out = MEF.relu(out_s2) 152 | 153 | out_s4 = self.conv3(out) 154 | out_s4 = self.norm3(out_s4) 155 | out_s4 = self.block3(out_s4) 156 | out = MEF.relu(out_s4) 157 | 158 | out_s8 = self.conv4(out) 159 | out_s8 = self.norm4(out_s8) 160 | out_s8 = self.block4(out_s8) 161 | out = MEF.relu(out_s8) 162 | 163 | out = self.conv4_tr(out) 164 | out = self.norm4_tr(out) 165 | out = self.block4_tr(out) 166 | out_s4_tr = MEF.relu(out) 167 | 168 | out = ME.cat(out_s4_tr, out_s4) 169 | 170 | out = self.conv3_tr(out) 171 | out = self.norm3_tr(out) 172 | out = self.block3_tr(out) 173 | out_s2_tr = MEF.relu(out) 174 | 175 | out = ME.cat(out_s2_tr, out_s2) 176 | 177 | out = self.conv2_tr(out) 178 | out = self.norm2_tr(out) 179 | out = self.block2_tr(out) 180 | out_s1_tr = MEF.relu(out) 181 | 182 | out = ME.cat(out_s1_tr, out_s1) 183 | out = self.conv1_tr(out) 184 | out = MEF.relu(out) 185 | out = self.final(out) 186 | 187 | if self.normalize_feature: 188 | return ME.SparseTensor( 189 | out.F / torch.norm(out.F, p=2, dim=1, keepdim=True), 190 | coordinate_map_key=out.coordinate_map_key, 191 | coordinate_manager=out.coordinate_manager) 192 | else: 193 | return out 194 | 195 | 196 | class ResUNetBN2(ResUNet2): 197 | NORM_TYPE = 'BN' 198 | 199 | 200 | class ResUNetBN2B(ResUNet2): 201 | NORM_TYPE = 'BN' 202 | CHANNELS = [None, 32, 64, 128, 256] 203 | TR_CHANNELS = [None, 64, 64, 64, 64] 204 | 205 | 206 | class ResUNetBN2C(ResUNet2): 207 | NORM_TYPE = 'BN' 208 | CHANNELS = [None, 32, 64, 128, 256] 209 | TR_CHANNELS = [None, 64, 64, 64, 128] 210 | 211 | 212 | class ResUNetBN2D(ResUNet2): 213 | NORM_TYPE = 'BN' 214 | CHANNELS = [None, 32, 64, 128, 256] 215 | TR_CHANNELS = [None, 64, 64, 128, 128] 216 | 217 | 218 | class ResUNetBN2E(ResUNet2): 219 | NORM_TYPE = 'BN' 220 | CHANNELS = [None, 128, 128, 128, 256] 221 | TR_CHANNELS = [None, 64, 128, 128, 128] 222 | 223 | 224 | class ResUNetFatBN(ResUNet2): 225 | NORM_TYPE = 'BN' 226 | CHANNELS = [None, 32, 64, 128, 256] 227 | TR_CHANNELS = [None, 128, 128, 128, 256] 228 | 229 | 230 | class ResUNetIN2(ResUNet2): 231 | NORM_TYPE = 'BN' 232 | BLOCK_NORM_TYPE = 'IN' 233 | 234 | 235 | class ResUNetIN2B(ResUNetBN2B): 236 | NORM_TYPE = 'BN' 237 | BLOCK_NORM_TYPE = 'IN' 238 | 239 | 240 | class ResUNetIN2C(ResUNetBN2C): 241 | NORM_TYPE = 'BN' 242 | BLOCK_NORM_TYPE = 'IN' 243 | 244 | 245 | class ResUNetIN2D(ResUNetBN2D): 246 | NORM_TYPE = 'BN' 247 | BLOCK_NORM_TYPE = 'IN' 248 | 249 | 250 | class ResUNetIN2E(ResUNetBN2E): 251 | NORM_TYPE = 'BN' 252 | BLOCK_NORM_TYPE = 'IN' 253 | -------------------------------------------------------------------------------- /FCGF_APR/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | # pytorch # for anaconda, please refer to pytorch.org for installation 3 | scipy 4 | matplotlib 5 | open3d 6 | # to visualize it, you need tensorflow, but it doesn't have to be in the same virual environment :) 7 | tensorboardX 8 | # MinkowskiEngine 9 | # Or follow the installation instruction on github.com/StanfordVL/MinkowskiEngine 10 | future-fstrings 11 | easydict 12 | joblib 13 | scikit-learn 14 | chamferdist 15 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/benchmark_util.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import os 3 | import logging 4 | import numpy as np 5 | 6 | from util.trajectory import CameraPose 7 | from util.pointcloud import compute_overlap_ratio, \ 8 | make_open3d_point_cloud, make_open3d_feature_from_numpy 9 | 10 | 11 | def run_ransac(xyz0, xyz1, feat0, feat1, voxel_size): 12 | distance_threshold = voxel_size * 1.5 13 | result_ransac = o3d.registration.registration_ransac_based_on_feature_matching( 14 | xyz0, xyz1, feat0, feat1, distance_threshold, 15 | o3d.registration.TransformationEstimationPointToPoint(False), 4, [ 16 | o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), 17 | o3d.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold) 18 | ], o3d.registration.RANSACConvergenceCriteria(4000000, 500)) 19 | return result_ransac.transformation 20 | 21 | 22 | def gather_results(results): 23 | traj = [] 24 | for r in results: 25 | success = r[0] 26 | if success: 27 | traj.append(CameraPose([r[1], r[2], r[3]], r[4])) 28 | return traj 29 | 30 | 31 | def gen_matching_pair(pts_num): 32 | matching_pairs = [] 33 | for i in range(pts_num): 34 | for j in range(i + 1, pts_num): 35 | matching_pairs.append([i, j, pts_num]) 36 | return matching_pairs 37 | 38 | 39 | def read_data(feature_path, name): 40 | data = np.load(os.path.join(feature_path, name + ".npz")) 41 | xyz = make_open3d_point_cloud(data['xyz']) 42 | feat = make_open3d_feature_from_numpy(data['feature']) 43 | return data['points'], xyz, feat 44 | 45 | 46 | def do_single_pair_matching(feature_path, set_name, m, voxel_size): 47 | i, j, s = m 48 | name_i = "%s_%03d" % (set_name, i) 49 | name_j = "%s_%03d" % (set_name, j) 50 | logging.info("matching %s %s" % (name_i, name_j)) 51 | points_i, xyz_i, feat_i = read_data(feature_path, name_i) 52 | points_j, xyz_j, feat_j = read_data(feature_path, name_j) 53 | if len(xyz_i.points) < len(xyz_j.points): 54 | trans = run_ransac(xyz_i, xyz_j, feat_i, feat_j, voxel_size) 55 | else: 56 | trans = run_ransac(xyz_j, xyz_i, feat_j, feat_i, voxel_size) 57 | trans = np.linalg.inv(trans) 58 | ratio = compute_overlap_ratio(xyz_i, xyz_j, trans, voxel_size) 59 | logging.info(f"{ratio}") 60 | if ratio > 0.3: 61 | return [True, i, j, s, np.linalg.inv(trans)] 62 | else: 63 | return [False, i, j, s, np.identity(4)] 64 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_apr.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d # prevent loading error 2 | 3 | import sys 4 | import logging 5 | import json 6 | import argparse 7 | import numpy as np 8 | from easydict import EasyDict as edict 9 | 10 | import torch 11 | from model import load_model 12 | 13 | from lib.complement_data_loader import make_data_loader 14 | 15 | from util.pointcloud import make_open3d_point_cloud, make_open3d_feature 16 | from lib.eval import find_nn_gpu 17 | from lib.timer import AverageMeter, Timer 18 | 19 | from chamferdist import ChamferDistance 20 | 21 | import MinkowskiEngine as ME 22 | import matplotlib.pyplot as plt 23 | # import os 24 | # os.environ["CUDA_VISIBLE_DEVICES"] = "0" 25 | 26 | ch = logging.StreamHandler(sys.stdout) 27 | logging.getLogger().setLevel(logging.INFO) 28 | logging.basicConfig( 29 | format='%(asctime)s %(message)s', datefmt='%m/%d %H:%M:%S', handlers=[ch]) 30 | 31 | 32 | def chamfer_distance(array1, array2, device): 33 | cd_dist = ChamferDistance() 34 | n1 = len(array1) 35 | n2 = len(array2) 36 | array1 = torch.unsqueeze(array1, dim=0).to(device) 37 | array2 = torch.unsqueeze(array2, dim=0).to(device) 38 | forward_cd_dist = cd_dist(array1, array2) 39 | backward_cd_dist = cd_dist(array2, array1) 40 | return forward_cd_dist / n1 + backward_cd_dist / n2 41 | 42 | 43 | def find_corr(xyz0, xyz1, F0, F1, subsample_size=-1): 44 | subsample = len(F0) > subsample_size 45 | if subsample_size > 0 and subsample: 46 | N0 = min(len(F0), subsample_size) 47 | N1 = min(len(F1), subsample_size) 48 | inds0 = np.random.choice(len(F0), N0, replace=False) 49 | inds1 = np.random.choice(len(F1), N1, replace=False) 50 | F0, F1 = F0[inds0], F1[inds1] 51 | 52 | # Compute the nn 53 | nn_inds = find_nn_gpu(F0, F1, nn_max_n=500) 54 | if subsample_size > 0 and subsample: 55 | return xyz0[inds0], xyz1[inds1[nn_inds]] 56 | else: 57 | return xyz0, xyz1[nn_inds] 58 | 59 | def apply_transform(pts, trans): 60 | R = trans[:3, :3] 61 | T = trans[:3, 3] 62 | return pts @ R.t() + T 63 | 64 | def evaluate_nn_dist(xyz0, xyz1, T_gth): 65 | xyz0 = apply_transform(xyz0, T_gth) 66 | dist = np.sqrt(((xyz0 - xyz1)**2).sum(1) + 1e-6) 67 | return dist.tolist() 68 | 69 | def main(config): 70 | test_loader = make_data_loader( 71 | config, config.test_phase, 1, num_threads=config.test_num_thread, shuffle=False) 72 | 73 | num_feats = 1 74 | 75 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 76 | 77 | Model = load_model(config.encoder_model) 78 | model = Model( 79 | num_feats, 80 | config.model_n_out, 81 | bn_momentum=config.bn_momentum, 82 | conv1_kernel_size=config.conv1_kernel_size, 83 | normalize_feature=config.normalize_feature) 84 | checkpoint = torch.load(config.save_dir + '/checkpoint.pth') 85 | model.load_state_dict(checkpoint['encoder_state_dict']) 86 | model = model.to(device) 87 | model.eval() 88 | 89 | success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter(), AverageMeter() 90 | data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer() 91 | 92 | test_iter = test_loader.__iter__() 93 | N = len(test_iter) 94 | n_gpu_failures = 0 95 | 96 | # downsample_voxel_size = 2 * config.voxel_size 97 | dists_success = [] 98 | dists_fail = [] 99 | dists_nn = [] 100 | list_rte = [] 101 | list_rre = [] 102 | trans_gt = [] 103 | T_gt = [] 104 | T_est = [] 105 | 106 | rte_thresh = 2 107 | rre_thresh = 5 108 | 109 | print(f"rre thresh: {rre_thresh}; rte_thresh: {rte_thresh}") 110 | 111 | for i in range(len(test_iter)): 112 | data_timer.tic() 113 | try: 114 | data_dict = test_iter.next() 115 | except ValueError: 116 | n_gpu_failures += 1 117 | logging.info(f"# Erroneous GPU Pair {n_gpu_failures}") 118 | continue 119 | data_timer.toc() 120 | xyz0, xyz1 = data_dict['pcd0'][0], data_dict['pcd1'][0] 121 | T_gth = data_dict['T_gt'] 122 | T_gt.append(T_gth) 123 | dist_gth = np.sqrt(np.sum((T_gth[:3, 3].cpu().numpy())**2)) 124 | trans_gt.append(dist_gth) 125 | xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy() 126 | 127 | pcd0 = make_open3d_point_cloud(xyz0np) 128 | pcd1 = make_open3d_point_cloud(xyz1np) 129 | 130 | with torch.no_grad(): 131 | feat_timer.tic() 132 | sinput0 = ME.SparseTensor( 133 | data_dict['sinput0_F'].to(device), coordinates=data_dict['sinput0_C'].to(device)) 134 | enc0 = model(sinput0) 135 | F0 = enc0.F.detach() 136 | sinput1 = ME.SparseTensor( 137 | data_dict['sinput1_F'].to(device), coordinates=data_dict['sinput1_C'].to(device)) 138 | enc1 = model(sinput1) 139 | F1 = enc1.F.detach() 140 | feat_timer.toc() 141 | 142 | xyz0_corr, xyz1_corr = find_corr(xyz0, xyz1, F0, F1, subsample_size=5000) 143 | dists_nn.append(evaluate_nn_dist(xyz0_corr, xyz1_corr, T_gth)) 144 | 145 | feat0 = make_open3d_feature(F0, 32, F0.shape[0]) 146 | feat1 = make_open3d_feature(F1, 32, F1.shape[0]) 147 | 148 | reg_timer.tic() 149 | distance_threshold = config.voxel_size * 1.0 150 | ransac_result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( 151 | pcd0, pcd1, feat0, feat1, False, distance_threshold, 152 | o3d.pipelines.registration.TransformationEstimationPointToPoint(False), 4, [ 153 | o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), 154 | o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold) 155 | ], o3d.pipelines.registration.RANSACConvergenceCriteria(4000000, 10000)) 156 | T_ransac = torch.from_numpy(ransac_result.transformation.astype(np.float32)) 157 | reg_timer.toc() 158 | 159 | T_est.append(T_ransac) 160 | 161 | # Translation error 162 | rte = np.linalg.norm(T_ransac[:3, 3] - T_gth[:3, 3]) 163 | rre = np.arccos((np.trace(T_ransac[:3, :3].t() @ T_gth[:3, :3]) - 1) / 2) 164 | 165 | # Check if the ransac was successful. successful if rte < 2m and rre < 5◦ 166 | # http://openaccess.thecvf.com/content_ECCV_2018/papers/Zi_Jian_Yew_3DFeat-Net_Weakly_Supervised_ECCV_2018_paper.pdf 167 | 168 | # rte_thresh = 2 169 | # rre_thresh = 5 170 | 171 | if rte < rte_thresh: 172 | rte_meter.update(rte) 173 | 174 | if not np.isnan(rre) and rre < np.pi / 180 * rre_thresh: 175 | rre_meter.update(rre * 180 / np.pi) 176 | 177 | if rte < rte_thresh and not np.isnan(rre) and rre < np.pi / 180 * rre_thresh: 178 | success_meter.update(1) 179 | dists_success.append(dist_gth) 180 | else: 181 | success_meter.update(0) 182 | dists_fail.append(dist_gth) 183 | logging.info(f"Failed with RTE: {rte}, RRE: {rre * 180 / np.pi}") 184 | 185 | list_rte.append(rte) 186 | list_rre.append(rre) 187 | 188 | if i % 10 == 0: 189 | logging.info( 190 | f"{i} / {N}: Data time: {data_timer.avg}, Feat time: {feat_timer.avg}," + 191 | f" Reg time: {reg_timer.avg}, RTE: {rte_meter.avg}," + 192 | f" RRE: {rre_meter.avg}, Success: {success_meter.sum} / {success_meter.count}" 193 | + f" ({success_meter.avg * 100} %)") 194 | data_timer.reset() 195 | feat_timer.reset() 196 | reg_timer.reset() 197 | 198 | print(f"rre thresh: {rre_thresh}; rte_thresh: {rte_thresh}") 199 | 200 | # dists_nn = np.array(dists_nn) # dists_nn: N_pairs * N_num_point_one_frame 201 | # hit_ratios = [] 202 | # for seed_tao_1 in range(1, 11): 203 | # tao_1 = seed_tao_1 * 0.1 204 | # hit_ratios.append(np.mean((dists_nn < tao_1).astype(float), axis=1).tolist()) 205 | 206 | # fmrs = [] 207 | # # hit_ratios: 10 * N_pairs 208 | # for seed_tao_2 in range(1, 11): 209 | # tao_2 = seed_tao_2 * 0.01 210 | # fmrs.append(np.mean((np.array(hit_ratios) > tao_2).astype(float), axis=1).tolist()) 211 | 212 | # fmrs: 10*10, dim0: seed_tao_2; dim2: seed_tao_1 213 | # np.save('FMRs_apr.npy', fmrs) 214 | 215 | logging.info( 216 | f"RTE: {rte_meter.avg}, var: {rte_meter.var}," + 217 | f" RRE: {rre_meter.avg}, var: {rre_meter.var}, Success: {success_meter.sum} " + 218 | f"/ {success_meter.count} ({success_meter.avg * 100} %)") 219 | 220 | 221 | def str2bool(v): 222 | return v.lower() in ('true', '1') 223 | 224 | 225 | if __name__ == '__main__': 226 | parser = argparse.ArgumentParser() 227 | parser.add_argument('--save_dir', default=None, type=str) 228 | parser.add_argument('--test_phase', default='test', type=str) 229 | parser.add_argument('--LoKITTI', default=False, type=str2bool) 230 | parser.add_argument('--LoNUSCENES', default=False, type=str2bool) 231 | parser.add_argument('--test_num_thread', default=5, type=int) 232 | parser.add_argument('--pair_min_dist', default=None, type=int) 233 | parser.add_argument('--pair_max_dist', default=None, type=int) 234 | parser.add_argument('--downsample_single', default=1.0, type=float) 235 | parser.add_argument('--kitti_root', type=str, default="/data/kitti/") 236 | parser.add_argument('--dataset', type=str, default="PairComplementKittiDataset") 237 | args = parser.parse_args() 238 | 239 | config = json.load(open(args.save_dir + '/config.json', 'r')) 240 | config = edict(config) 241 | config.save_dir = args.save_dir 242 | config.test_phase = 'test' 243 | config.kitti_root = args.kitti_root 244 | config.kitti_odometry_root = args.kitti_root + '/dataset' 245 | config.test_num_thread = args.test_num_thread 246 | config.LoKITTI = args.LoKITTI 247 | config.LoNUSCENES = args.LoNUSCENES 248 | config.debug_use_old_complement = True 249 | config.debug_need_complement = False 250 | config.phase = 'test' 251 | config.dataset = args.dataset 252 | if config.dataset == "PairComplementNuscenesDataset": 253 | config.use_old_pose = True 254 | else: 255 | config.use_old_pose = False 256 | 257 | if args.pair_min_dist is not None and args.pair_max_dist is not None: 258 | config.pair_min_dist = args.pair_min_dist 259 | config.pair_max_dist = args.pair_max_dist 260 | config.downsample_single = args.downsample_single 261 | 262 | main(config) 263 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_apr_kitti.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=0 6 | 7 | export KITTI_PATH="/mnt/disk/KITTIOdometry_Full" 8 | export OUT_DIR=${OUT_DIR:-./outputs/Experiments/PairComplementNuscenesDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-01-29_00-28-55} 9 | export PYTHONUNBUFFERED="True" 10 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 11 | export VERSION=$(git rev-parse HEAD) 12 | 13 | echo $OUT_DIR 14 | mkdir -m 755 -p $OUT_DIR 15 | LOG=${OUT_DIR}/log_${TIME}.txt 16 | 17 | echo "Host: " $(hostname) | tee -a $LOG 18 | echo "Conda " $(which conda) | tee -a $LOG 19 | echo $(pwd) | tee -a $LOG 20 | echo "Version: " $VERSION | tee -a $LOG 21 | # echo "Git diff" | tee -a $LOG 22 | # echo "" | tee -a $LOG 23 | # git diff | tee -a $LOG 24 | echo "" | tee -a $LOG 25 | nvidia-smi | tee -a $LOG 26 | 27 | 28 | # Test 29 | python -m scripts.test_rcar_kitti \ 30 | --kitti_root ${KITTI_PATH} \ 31 | --LoKITTI false \ 32 | --pair_min_dist 5 \ 33 | --pair_max_dist 20 \ 34 | --save_dir ${OUT_DIR} | tee -a $LOG 35 | # --downsample_single 1.0 \ 36 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_apr_nuscenes.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=0 6 | 7 | export KITTI_PATH="/mnt/disk/NUSCENES/nusc_kitti" 8 | export OUT_DIR=${OUT_DIR:-./outputs/Experiments/PairComplementKittiDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-04-19_20-14-42} 9 | export PYTHONUNBUFFERED="True" 10 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 11 | export VERSION=$(git rev-parse HEAD) 12 | 13 | echo $OUT_DIR 14 | mkdir -m 755 -p $OUT_DIR 15 | LOG=${OUT_DIR}/log_${TIME}.txt 16 | 17 | echo "Host: " $(hostname) | tee -a $LOG 18 | echo "Conda " $(which conda) | tee -a $LOG 19 | echo $(pwd) | tee -a $LOG 20 | echo "Version: " $VERSION | tee -a $LOG 21 | # echo "Git diff" | tee -a $LOG 22 | # echo "" | tee -a $LOG 23 | # git diff | tee -a $LOG 24 | echo "" | tee -a $LOG 25 | nvidia-smi | tee -a $LOG 26 | 27 | 28 | # Test 29 | python -m scripts.test_rcar_kitti \ 30 | --kitti_root ${KITTI_PATH} \ 31 | --LoNUSCENES false \ 32 | --pair_min_dist 5 \ 33 | --pair_max_dist 20 \ 34 | --dataset PairComplementNuscenesDataset \ 35 | --save_dir ${OUT_DIR} | tee -a $LOG 36 | # --downsample_single 1.0 \ 37 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_fcgf.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d # prevent loading error 2 | 3 | import sys 4 | import logging 5 | import json 6 | import argparse 7 | import numpy as np 8 | from easydict import EasyDict as edict 9 | 10 | import torch 11 | from model import load_model 12 | 13 | from lib.complement_data_loader import make_data_loader 14 | from lib.eval import find_nn_gpu 15 | 16 | from util.pointcloud import make_open3d_point_cloud, make_open3d_feature 17 | from lib.timer import AverageMeter, Timer 18 | 19 | import MinkowskiEngine as ME 20 | # import os 21 | # os.environ["CUDA_VISIBLE_DEVICES"] = "3" 22 | 23 | ch = logging.StreamHandler(sys.stdout) 24 | logging.getLogger().setLevel(logging.INFO) 25 | logging.basicConfig( 26 | format='%(asctime)s %(message)s', datefmt='%m/%d %H:%M:%S', handlers=[ch]) 27 | 28 | def find_corr(xyz0, xyz1, F0, F1, subsample_size=-1): 29 | subsample = len(F0) > subsample_size 30 | if subsample_size > 0 and subsample: 31 | N0 = min(len(F0), subsample_size) 32 | N1 = min(len(F1), subsample_size) 33 | inds0 = np.random.choice(len(F0), N0, replace=False) 34 | inds1 = np.random.choice(len(F1), N1, replace=False) 35 | F0, F1 = F0[inds0], F1[inds1] 36 | 37 | # Compute the nn 38 | nn_inds = find_nn_gpu(F0, F1, nn_max_n=500) 39 | if subsample_size > 0 and subsample: 40 | return xyz0[inds0], xyz1[inds1[nn_inds]] 41 | else: 42 | return xyz0, xyz1[nn_inds] 43 | 44 | def apply_transform(pts, trans): 45 | R = trans[:3, :3] 46 | T = trans[:3, 3] 47 | return pts @ R.t() + T 48 | 49 | def evaluate_nn_dist(xyz0, xyz1, T_gth): 50 | xyz0 = apply_transform(xyz0, T_gth) 51 | dist = np.sqrt(((xyz0 - xyz1)**2).sum(1) + 1e-6) 52 | return dist.tolist() 53 | 54 | def random_sample(pcd, feats, N): 55 | """ 56 | Do random sampling to get exact N points and associated features 57 | pcd: [N,3] 58 | feats: [N,C] 59 | """ 60 | if(isinstance(pcd,torch.Tensor)): 61 | n1 = pcd.size(0) 62 | elif(isinstance(pcd, np.ndarray)): 63 | n1 = pcd.shape[0] 64 | 65 | if n1 == N: 66 | return pcd, feats 67 | 68 | if n1 > N: 69 | choice = np.random.permutation(n1)[:N] 70 | else: 71 | choice = np.random.choice(n1, N) 72 | 73 | return pcd[choice], feats[choice] 74 | 75 | def main(config): 76 | test_loader = make_data_loader( 77 | config, config.test_phase, 1, num_threads=config.test_num_thread, shuffle=False) 78 | 79 | num_feats = 1 80 | 81 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 82 | 83 | Model = load_model(config.model) 84 | model = Model( 85 | num_feats, 86 | config.model_n_out, 87 | bn_momentum=config.bn_momentum, 88 | conv1_kernel_size=config.conv1_kernel_size, 89 | normalize_feature=config.normalize_feature) 90 | checkpoint = torch.load(config.save_dir + '/best_val_checkpoint.pth') 91 | model.load_state_dict(checkpoint['state_dict']) 92 | model = model.to(device) 93 | model.eval() 94 | 95 | success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter(), AverageMeter() 96 | data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer() 97 | 98 | 99 | test_iter = test_loader.__iter__() 100 | N = len(test_iter) 101 | n_gpu_failures = 0 102 | 103 | dists_success = [] 104 | dists_fail = [] 105 | dists_nn = [] 106 | list_rte = [] 107 | list_rre = [] 108 | trans_gt = [] 109 | T_gt = [] 110 | T_est = [] 111 | 112 | rte_thresh = 2 113 | rre_thresh = 5 114 | print(f"rre thresh: {rre_thresh}; rte_thresh: {rte_thresh}") 115 | 116 | for i in range(len(test_iter)): 117 | data_timer.tic() 118 | 119 | data_dict = test_iter.next() 120 | data_timer.toc() 121 | xyz0, xyz1 = data_dict['pcd0'][0], data_dict['pcd1'][0] 122 | T_gth = data_dict['T_gt'] 123 | T_gt.append(T_gth) 124 | dist_gth = np.sqrt(np.sum((T_gth[:3, 3].cpu().numpy())**2)) 125 | trans_gt.append(dist_gth) 126 | xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy() 127 | 128 | pcd0 = make_open3d_point_cloud(xyz0np) 129 | pcd1 = make_open3d_point_cloud(xyz1np) 130 | 131 | with torch.no_grad(): 132 | feat_timer.tic() 133 | sinput0 = ME.SparseTensor( 134 | data_dict['sinput0_F'].to(device), coordinates=data_dict['sinput0_C'].to(device)) 135 | enc0 = model(sinput0) 136 | F0 = enc0.F.detach() 137 | sinput1 = ME.SparseTensor( 138 | data_dict['sinput1_F'].to(device), coordinates=data_dict['sinput1_C'].to(device)) 139 | enc1 = model(sinput1) 140 | F1 = enc1.F.detach() 141 | feat_timer.toc() 142 | 143 | xyz0_corr, xyz1_corr = find_corr(xyz0, xyz1, F0, F1, subsample_size=5000) 144 | dists_nn.append(evaluate_nn_dist(xyz0_corr, xyz1_corr, T_gth)) 145 | 146 | n_points = 5000 147 | ######################################## 148 | # run random sampling or probabilistic sampling 149 | xyz0np, F0 = random_sample(xyz0np, F0, n_points) 150 | xyz1np, F1 = random_sample(xyz1np, F1, n_points) 151 | 152 | pcd0 = make_open3d_point_cloud(xyz0np) 153 | pcd1 = make_open3d_point_cloud(xyz1np) 154 | 155 | feat0 = make_open3d_feature(F0, 32, F0.shape[0]) 156 | feat1 = make_open3d_feature(F1, 32, F1.shape[0]) 157 | 158 | reg_timer.tic() 159 | distance_threshold = config.voxel_size * 1.0 160 | ransac_result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( 161 | pcd0, pcd1, feat0, feat1, False, distance_threshold, 162 | o3d.pipelines.registration.TransformationEstimationPointToPoint(False), 4, [ 163 | o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), 164 | o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold) 165 | ], o3d.pipelines.registration.RANSACConvergenceCriteria(4000000, 10000)) 166 | T_ransac = torch.from_numpy(ransac_result.transformation.astype(np.float32)) 167 | reg_timer.toc() 168 | 169 | T_est.append(T_ransac) 170 | 171 | # Translation error 172 | rte = np.linalg.norm(T_ransac[:3, 3] - T_gth[:3, 3]) 173 | rre = np.arccos((np.trace(T_ransac[:3, :3].t() @ T_gth[:3, :3]) - 1) / 2) 174 | 175 | # Check if the ransac was successful. successful if rte < 2m and rre < 5◦ 176 | # http://openaccess.thecvf.com/content_ECCV_2018/papers/Zi_Jian_Yew_3DFeat-Net_Weakly_Supervised_ECCV_2018_paper.pdf 177 | # rte_thresh = 0.6 178 | # rre_thresh = 1.5 179 | 180 | if rte < rte_thresh: 181 | rte_meter.update(rte) 182 | 183 | if not np.isnan(rre) and rre < np.pi / 180 * rre_thresh: 184 | rre_meter.update(rre * 180 / np.pi) 185 | 186 | if rte < rte_thresh and not np.isnan(rre) and rre < np.pi / 180 * rre_thresh: 187 | success_meter.update(1) 188 | dists_success.append(dist_gth) 189 | else: 190 | success_meter.update(0) 191 | dists_fail.append(dist_gth) 192 | logging.info(f"Failed with RTE: {rte}, RRE: {rre * 180 / np.pi}") 193 | 194 | list_rte.append(rte) 195 | list_rre.append(rre) 196 | 197 | if i % 10 == 0: 198 | logging.info( 199 | f"{i} / {N}: Data time: {data_timer.avg}, Feat time: {feat_timer.avg}," + 200 | f" Reg time: {reg_timer.avg}, RTE: {rte_meter.avg}," + 201 | f" RRE: {rre_meter.avg}, Success: {success_meter.sum} / {success_meter.count}" 202 | + f" ({success_meter.avg * 100} %)") 203 | data_timer.reset() 204 | feat_timer.reset() 205 | reg_timer.reset() 206 | 207 | print(f"rre thresh: {rre_thresh}; rte_thresh: {rte_thresh}") 208 | 209 | np.savez(f"results_fcgf_LoNuScenes.npz", rres=list_rre, rtes=list_rte, dists=trans_gt, T_gt=T_gt, T_est=T_est) 210 | 211 | # dists_nn = np.array(dists_nn, dtype=object) # dists_nn: N_pairs * N_num_point_one_frame 212 | # hit_ratios = [] 213 | # for seed_tao_1 in range(1, 11): 214 | # tao_1 = seed_tao_1 * 0.1 215 | # hit_ratios.append(np.mean((dists_nn < tao_1).astype(float), axis=1).tolist()) 216 | 217 | # fmrs = [] 218 | # # hit_ratios: 10 * N_pairs 219 | # for seed_tao_2 in range(1, 11): 220 | # tao_2 = seed_tao_2 * 0.01 221 | # fmrs.append(np.mean((np.array(hit_ratios) > tao_2).astype(float), axis=1).tolist()) 222 | 223 | # fmrs: 10*10, dim0: seed_tao_2; dim2: seed_tao_1 224 | # np.save('FMRs_fcgf.npy', fmrs) 225 | 226 | logging.info( 227 | f"RTE: {rte_meter.avg}, var: {rte_meter.var}," + 228 | f" RRE: {rre_meter.avg}, var: {rre_meter.var}, Success: {success_meter.sum} " + 229 | f"/ {success_meter.count} ({success_meter.avg * 100} %)") 230 | 231 | 232 | def str2bool(v): 233 | return v.lower() in ('true', '1') 234 | 235 | 236 | if __name__ == '__main__': 237 | parser = argparse.ArgumentParser() 238 | parser.add_argument('--save_dir', default=None, type=str) 239 | parser.add_argument('--test_phase', default='test', type=str) 240 | parser.add_argument('--LoKITTI', default=False, type=str2bool) 241 | parser.add_argument('--LoNUSCENES', default=False, type=str2bool) 242 | parser.add_argument('--test_num_thread', default=5, type=int) 243 | parser.add_argument('--pair_min_dist', default=None, type=int) 244 | parser.add_argument('--pair_max_dist', default=None, type=int) 245 | parser.add_argument('--downsample_single', default=1.0, type=float) 246 | parser.add_argument('--kitti_root', type=str, default="/data/kitti/") 247 | parser.add_argument('--dataset', type=str, default="PairComplementKittiDataset") 248 | args = parser.parse_args() 249 | 250 | config = json.load(open(args.save_dir + '/config.json', 'r')) 251 | config = edict(config) 252 | config.save_dir = args.save_dir 253 | config.test_phase = args.test_phase 254 | config.kitti_root = args.kitti_root 255 | config.kitti_odometry_root = args.kitti_root + '/dataset' 256 | config.test_num_thread = args.test_num_thread 257 | config.LoKITTI = args.LoKITTI 258 | config.LoNUSCENES = args.LoNUSCENES 259 | config.debug_use_old_complement = True 260 | config.debug_need_complement = False 261 | config.phase = 'test' 262 | config.dataset = args.dataset 263 | if config.dataset == "PairComplementNuscenesDataset": 264 | config.use_old_pose = True 265 | else: 266 | config.use_old_pose = False 267 | # config.debug_force_icp_recalculation = True 268 | 269 | if args.pair_min_dist is not None and args.pair_max_dist is not None: 270 | config.pair_min_dist = args.pair_min_dist 271 | config.pair_max_dist = args.pair_max_dist 272 | config.downsample_single = args.downsample_single 273 | 274 | main(config) 275 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_fcgf_kitti.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=0 6 | 7 | export KITTI_PATH="/mnt/disk/KITTIOdometry_Full" 8 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 9 | export VERSION=$(git rev-parse HEAD) 10 | export OUT_DIR=${OUT_DIR:-./outputs/Experiments/PairComplementKittiDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout32/2023-02-27_17-03-23} 11 | export PYTHONUNBUFFERED="True" 12 | 13 | echo $OUT_DIR 14 | 15 | mkdir -m 755 -p $OUT_DIR 16 | 17 | LOG=${OUT_DIR}/log_${TIME}.txt 18 | 19 | echo "Host: " $(hostname) | tee -a $LOG 20 | echo "Conda " $(which conda) | tee -a $LOG 21 | echo $(pwd) | tee -a $LOG 22 | echo "Version: " $VERSION | tee -a $LOG 23 | # echo "Git diff" | tee -a $LOG 24 | # echo "" | tee -a $LOG 25 | # git diff | tee -a $LOG 26 | echo "" | tee -a $LOG 27 | nvidia-smi | tee -a $LOG 28 | 29 | 30 | # Test 31 | python -m scripts.test_kitti \ 32 | --kitti_root ${KITTI_PATH} \ 33 | --LoKITTI false \ 34 | --pair_min_dist 40 \ 35 | --pair_max_dist 50 \ 36 | --save_dir ${OUT_DIR} | tee -a $LOG 37 | # --downsample_single 1.0 \ 38 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/test_fcgf_nuscenes.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=0 6 | 7 | export KITTI_PATH="/mnt/disk/NUSCENES/nusc_kitti" 8 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 9 | export VERSION=$(git rev-parse HEAD) 10 | export OUT_DIR=${OUT_DIR:-./outputs/Experiments/PairComplementKittiDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-01-11_15-00-27} 11 | export PYTHONUNBUFFERED="True" 12 | 13 | echo $OUT_DIR 14 | 15 | mkdir -m 755 -p $OUT_DIR 16 | 17 | LOG=${OUT_DIR}/log_${TIME}.txt 18 | 19 | echo "Host: " $(hostname) | tee -a $LOG 20 | echo "Conda " $(which conda) | tee -a $LOG 21 | echo $(pwd) | tee -a $LOG 22 | echo "Version: " $VERSION | tee -a $LOG 23 | # echo "Git diff" | tee -a $LOG 24 | # echo "" | tee -a $LOG 25 | # git diff | tee -a $LOG 26 | echo "" | tee -a $LOG 27 | nvidia-smi | tee -a $LOG 28 | 29 | 30 | # Test 31 | python -m scripts.test_kitti \ 32 | --kitti_root ${KITTI_PATH} \ 33 | --LoNUSCENES false \ 34 | --pair_min_dist 5 \ 35 | --pair_max_dist 20 \ 36 | --dataset PairComplementNuscenesDataset \ 37 | --save_dir ${OUT_DIR} | tee -a $LOG 38 | # --downsample_single 1.0 \ 39 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/train_apr_kitti.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=3 6 | 7 | export KITTI_PATH="/mnt/disk/KITTIOdometry_Full" 8 | 9 | export DATA_ROOT="./outputs/Experiments" 10 | export DATASET=${DATASET:-PairComplementKittiDataset} 11 | export TRAINER=${TRAINER:-GenerativePairTrainer} 12 | export ENCODER_MODEL=${MODEL:-ResUNetFatBN} 13 | export MODEL_N_OUT=${MODEL_N_OUT:-128} 14 | export GENERATOR_MODEL=${MODEL:-GenerativeMLP_98} 15 | export OPTIMIZER=${OPTIMIZER:-SGD} 16 | export LR=${LR:-1e-1} 17 | export LOSS_RATIO=${LOSS_RATIO:-2e-3} 18 | export WEIGHT_DECAY=${WEIGHT_DECAY:-1e-4} 19 | export MAX_EPOCH=${MAX_EPOCH:-200} 20 | export BATCH_SIZE=${BATCH_SIZE:-4} 21 | export VAL_BATCH_SIZE=${VAL_BATCH_SIZE:-1} 22 | export ITER_SIZE=${ITER_SIZE:-1} 23 | export BEST_VAL_METRIC=${BEST_VAL_METRIC:-feat_match_ratio} 24 | export VOXEL_SIZE=${VOXEL_SIZE:-0.3} 25 | export POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER=${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER:-1.5} 26 | export CONV1_KERNEL_SIZE=${CONV1_KERNEL_SIZE:-5} 27 | export EXP_GAMMA=${EXP_GAMMA:-0.99} 28 | export RANDOM_SCALE=${RANDOM_SCALE:-True} 29 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 30 | export VERSION=$(git rev-parse HEAD) 31 | 32 | export OUT_DIR=${DATA_ROOT}/${DATASET}-v${VOXEL_SIZE}/${TRAINER}/${MODEL}/${OPTIMIZER}-lr${LR}-e${MAX_EPOCH}-b${BATCH_SIZE}i${ITER_SIZE}-modelnout${MODEL_N_OUT}${PATH_POSTFIX}/${TIME} 33 | 34 | export PYTHONUNBUFFERED="True" 35 | 36 | # export RESUME_FILE="./outputs/Experiments/PairComplementKittiDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-12-18_12-23-19/checkpoint.pth" 37 | # export RESUME_DIR="./outputs/Experiments/PairComplementKittiDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-12-18_12-23-19" 38 | 39 | echo $OUT_DIR 40 | echo "Using GPU No:" 41 | echo $CUDA_VISIBLE_DEVICES 42 | 43 | mkdir -m 755 -p $OUT_DIR 44 | 45 | LOG=${OUT_DIR}/log_${TIME}.txt 46 | 47 | echo "Host: " $(hostname) | tee -a $LOG 48 | echo "Conda " $(which conda) | tee -a $LOG 49 | echo $(pwd) | tee -a $LOG 50 | echo "Version: " $VERSION | tee -a $LOG 51 | # echo "Git diff" | tee -a $LOG 52 | # echo "" | tee -a $LOG 53 | # git diff | tee -a $LOG 54 | echo "" | tee -a $LOG 55 | nvidia-smi | tee -a $LOG 56 | 57 | # Training With Resume 58 | python train.py \ 59 | --dataset ${DATASET} \ 60 | --trainer ${TRAINER} \ 61 | --encoder_model ${ENCODER_MODEL} \ 62 | --generator_model ${GENERATOR_MODEL} \ 63 | --model_n_out ${MODEL_N_OUT} \ 64 | --conv1_kernel_size ${CONV1_KERNEL_SIZE} \ 65 | --optimizer ${OPTIMIZER} \ 66 | --lr ${LR} \ 67 | --loss_ratio ${LOSS_RATIO} \ 68 | --batch_size ${BATCH_SIZE} \ 69 | --val_batch_size ${VAL_BATCH_SIZE} \ 70 | --iter_size ${ITER_SIZE} \ 71 | --max_epoch ${MAX_EPOCH} \ 72 | --voxel_size ${VOXEL_SIZE} \ 73 | --out_dir ${OUT_DIR} \ 74 | --use_random_scale ${RANDOM_SCALE} \ 75 | --use_random_rotation true \ 76 | --positive_pair_search_voxel_size_multiplier ${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER} \ 77 | --weight_decay ${WEIGHT_DECAY} \ 78 | --kitti_root ${KITTI_PATH} \ 79 | --hit_ratio_thresh 0.3 \ 80 | --exp_gamma ${EXP_GAMMA} \ 81 | --complement_pair_dist 10 \ 82 | --num_complement_one_side 3 \ 83 | --point_generation_ratio 4 \ 84 | --regularization_strength 0.01 \ 85 | --regularization_type L2 \ 86 | --best_val_metric ${BEST_VAL_METRIC} \ 87 | --debug_use_old_complement false \ 88 | --use_old_pose false \ 89 | --pair_min_dist 5 \ 90 | --pair_max_dist 20 \ 91 | --symmetric False \ 92 | --mutate_neighbour_percentage 0 \ 93 | # --resume ${RESUME_FILE} \ 94 | # --resume_dir ${RESUME_DIR} \ 95 | # --finetune_restart false \ 96 | $MISC_ARGS 2>&1 | tee -a $LOG 97 | 98 | # Test 99 | # python -m scripts.test_rcar_kitti \ 100 | # --kitti_root ${KITTI_PATH} \ 101 | # --save_dir ${OUT_DIR} | tee -a $LOG 102 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/train_apr_nuscenes.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=1 6 | 7 | export KITTI_PATH="/mnt/disk/NUSCENES/nusc_kitti" 8 | 9 | export DATA_ROOT="./outputs/Experiments" 10 | export DATASET=${DATASET:-PairComplementNuscenesDataset} 11 | export TRAINER=${TRAINER:-GenerativePairTrainer} 12 | export ENCODER_MODEL=${MODEL:-ResUNetFatBN} 13 | export MODEL_N_OUT=${MODEL_N_OUT:-128} 14 | export GENERATOR_MODEL=${MODEL:-ResUNetFatBN} 15 | export OPTIMIZER=${OPTIMIZER:-SGD} 16 | export LR=${LR:-1e-1} 17 | export LOSS_RATIO=${LOSS_RATIO:-2e-3} 18 | export WEIGHT_DECAY=${WEIGHT_DECAY:-1e-4} 19 | export MAX_EPOCH=${MAX_EPOCH:-200} 20 | export BATCH_SIZE=${BATCH_SIZE:-4} 21 | export VAL_BATCH_SIZE=${VAL_BATCH_SIZE:-1} 22 | export ITER_SIZE=${ITER_SIZE:-1} 23 | export BEST_VAL_METRIC=${BEST_VAL_METRIC:-feat_match_ratio} 24 | export VOXEL_SIZE=${VOXEL_SIZE:-0.3} 25 | export POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER=${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER:-1.5} 26 | export CONV1_KERNEL_SIZE=${CONV1_KERNEL_SIZE:-5} 27 | export EXP_GAMMA=${EXP_GAMMA:-0.99} 28 | export RANDOM_SCALE=${RANDOM_SCALE:-True} 29 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 30 | export VERSION=$(git rev-parse HEAD) 31 | 32 | export OUT_DIR=${DATA_ROOT}/${DATASET}-v${VOXEL_SIZE}/${TRAINER}/${MODEL}/${OPTIMIZER}-lr${LR}-e${MAX_EPOCH}-b${BATCH_SIZE}i${ITER_SIZE}-modelnout${MODEL_N_OUT}${PATH_POSTFIX}/${TIME} 33 | 34 | export PYTHONUNBUFFERED="True" 35 | 36 | # export RESUME_FILE="./outputs/Experiments/PairComplementNuscenesDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-12-22_13-00-33/checkpoint.pth" 37 | # export RESUME_DIR="./outputs/Experiments/PairComplementNuscenesDataset-v0.3/GenerativePairTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-12-22_13-00-33" 38 | 39 | echo $OUT_DIR 40 | echo "Using GPU No:" 41 | echo $CUDA_VISIBLE_DEVICES 42 | 43 | mkdir -m 755 -p $OUT_DIR 44 | 45 | LOG=${OUT_DIR}/log_${TIME}.txt 46 | 47 | echo "Host: " $(hostname) | tee -a $LOG 48 | echo "Conda " $(which conda) | tee -a $LOG 49 | echo $(pwd) | tee -a $LOG 50 | echo "Version: " $VERSION | tee -a $LOG 51 | # echo "Git diff" | tee -a $LOG 52 | # echo "" | tee -a $LOG 53 | # git diff | tee -a $LOG 54 | echo "" | tee -a $LOG 55 | nvidia-smi | tee -a $LOG 56 | 57 | # Training With Resume 58 | python train.py \ 59 | --dataset ${DATASET} \ 60 | --trainer ${TRAINER} \ 61 | --encoder_model ${ENCODER_MODEL} \ 62 | --generator_model ${GENERATOR_MODEL} \ 63 | --model_n_out ${MODEL_N_OUT} \ 64 | --conv1_kernel_size ${CONV1_KERNEL_SIZE} \ 65 | --optimizer ${OPTIMIZER} \ 66 | --lr ${LR} \ 67 | --loss_ratio ${LOSS_RATIO} \ 68 | --batch_size ${BATCH_SIZE} \ 69 | --val_batch_size ${VAL_BATCH_SIZE} \ 70 | --iter_size ${ITER_SIZE} \ 71 | --max_epoch ${MAX_EPOCH} \ 72 | --voxel_size ${VOXEL_SIZE} \ 73 | --out_dir ${OUT_DIR} \ 74 | --use_random_scale ${RANDOM_SCALE} \ 75 | --use_random_rotation true \ 76 | --positive_pair_search_voxel_size_multiplier ${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER} \ 77 | --weight_decay ${WEIGHT_DECAY} \ 78 | --kitti_root ${KITTI_PATH} \ 79 | --hit_ratio_thresh 0.3 \ 80 | --exp_gamma ${EXP_GAMMA} \ 81 | --complement_pair_dist 10 \ 82 | --num_complement_one_side 3 \ 83 | --point_generation_ratio 4 \ 84 | --regularization_strength 0.01 \ 85 | --regularization_type L2 \ 86 | --best_val_metric ${BEST_VAL_METRIC} \ 87 | --debug_use_old_complement false \ 88 | --use_old_pose true \ 89 | --pair_min_dist 5 \ 90 | --pair_max_dist 20 \ 91 | --symmetric True \ 92 | --mutate_neighbour_percentage 0.9 \ 93 | # --resume ${RESUME_FILE} \ 94 | # --resume_dir ${RESUME_DIR} \ 95 | # --finetune_restart true \ 96 | $MISC_ARGS 2>&1 | tee -a $LOG 97 | 98 | # Test 99 | # python -m scripts.test_rcar_kitti \ 100 | # --kitti_root ${KITTI_PATH} \ 101 | # --save_dir ${OUT_DIR} | tee -a $LOG 102 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/train_fcgf_kitti.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=1 6 | 7 | export KITTI_PATH="/mnt/disk/KITTIOdometry_Full" 8 | 9 | export DATA_ROOT="./outputs/Experiments" 10 | export DATASET=${DATASET:-PairComplementKittiDataset} 11 | export TRAINER=${TRAINER:-HardestContrastiveLossTrainer} 12 | export ENCODER_MODEL=${MODEL:-ResUNetFatBN} 13 | export MODEL_N_OUT=${MODEL_N_OUT:-32} 14 | export OPTIMIZER=${OPTIMIZER:-SGD} 15 | export LR=${LR:-1e-1} 16 | export WEIGHT_DECAY=${WEIGHT_DECAY:-1e-4} 17 | export MAX_EPOCH=${MAX_EPOCH:-200} 18 | export BATCH_SIZE=${BATCH_SIZE:-4} 19 | export VAL_BATCH_SIZE=${VAL_BATCH_SIZE:-1} 20 | export ITER_SIZE=${ITER_SIZE:-1} 21 | export BEST_VAL_METRIC=${BEST_VAL_METRIC:-feat_match_ratio} 22 | export VOXEL_SIZE=${VOXEL_SIZE:-0.3} 23 | export POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER=${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER:-1.5} 24 | export CONV1_KERNEL_SIZE=${CONV1_KERNEL_SIZE:-5} 25 | export EXP_GAMMA=${EXP_GAMMA:-0.99} 26 | export RANDOM_SCALE=${RANDOM_SCALE:-True} 27 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 28 | export VERSION=$(git rev-parse HEAD) 29 | 30 | export OUT_DIR=${DATA_ROOT}/${DATASET}-v${VOXEL_SIZE}/${TRAINER}/${MODEL}/${OPTIMIZER}-lr${LR}-e${MAX_EPOCH}-b${BATCH_SIZE}i${ITER_SIZE}-modelnout${MODEL_N_OUT}${PATH_POSTFIX}/${TIME} 31 | 32 | export PYTHONUNBUFFERED="True" 33 | 34 | # export RESUME_FILE="./outputs/Experiments/PairComplementKittiDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout32/2023-02-13_19-48-37/checkpoint.pth" 35 | # export RESUME_DIR="./outputs/Experiments/PairComplementKittiDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout32/2023-02-13_19-48-37" 36 | 37 | echo $OUT_DIR 38 | echo "Using GPU No:" 39 | echo $CUDA_VISIBLE_DEVICES 40 | 41 | mkdir -m 755 -p $OUT_DIR 42 | 43 | LOG=${OUT_DIR}/log_${TIME}.txt 44 | 45 | echo "Host: " $(hostname) | tee -a $LOG 46 | echo "Conda " $(which conda) | tee -a $LOG 47 | echo $(pwd) | tee -a $LOG 48 | echo "Version: " $VERSION | tee -a $LOG 49 | # echo "Git diff" | tee -a $LOG 50 | # echo "" | tee -a $LOG 51 | # git diff | tee -a $LOG 52 | echo "" | tee -a $LOG 53 | nvidia-smi | tee -a $LOG 54 | 55 | # Training With Resume 56 | python train.py \ 57 | --dataset ${DATASET} \ 58 | --trainer ${TRAINER} \ 59 | --model ${ENCODER_MODEL} \ 60 | --model_n_out ${MODEL_N_OUT} \ 61 | --conv1_kernel_size ${CONV1_KERNEL_SIZE} \ 62 | --optimizer ${OPTIMIZER} \ 63 | --lr ${LR} \ 64 | --batch_size ${BATCH_SIZE} \ 65 | --val_batch_size ${VAL_BATCH_SIZE} \ 66 | --iter_size ${ITER_SIZE} \ 67 | --max_epoch ${MAX_EPOCH} \ 68 | --voxel_size ${VOXEL_SIZE} \ 69 | --out_dir ${OUT_DIR} \ 70 | --use_random_scale ${RANDOM_SCALE} \ 71 | --weight_decay ${WEIGHT_DECAY} \ 72 | --positive_pair_search_voxel_size_multiplier ${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER} \ 73 | --kitti_root ${KITTI_PATH} \ 74 | --hit_ratio_thresh 0.3 \ 75 | --exp_gamma ${EXP_GAMMA} \ 76 | --complement_pair_dist 10 \ 77 | --num_complement_one_side 3 \ 78 | --point_generation_ratio 4 \ 79 | --regularization_strength 0.01 \ 80 | --regularization_type L2 \ 81 | --best_val_metric ${BEST_VAL_METRIC} \ 82 | --debug_use_old_complement false \ 83 | --use_old_pose false \ 84 | --pair_min_dist 5 \ 85 | --pair_max_dist 20 \ 86 | # --resume ${RESUME_FILE} \ 87 | # --resume_dir ${RESUME_DIR} \ 88 | # --finetune_restart true \ 89 | $MISC_ARGS 2>&1 | tee -a $LOG 90 | 91 | # Test 92 | # python -m scripts.test_kitti \ 93 | # --kitti_root ${KITTI_PATH} \ 94 | # --save_dir ${OUT_DIR} | tee -a $LOG 95 | -------------------------------------------------------------------------------- /FCGF_APR/scripts/train_fcgf_nuscenes.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | export PATH_POSTFIX=$1 3 | export MISC_ARGS=$2 4 | 5 | export CUDA_VISIBLE_DEVICES=2 6 | 7 | export KITTI_PATH="/mnt/disk/NUSCENES/nusc_kitti" 8 | 9 | export DATA_ROOT="./outputs/Experiments" 10 | export DATASET=${DATASET:-PairComplementNuscenesDataset} 11 | export TRAINER=${TRAINER:-HardestContrastiveLossTrainer} 12 | export ENCODER_MODEL=${MODEL:-ResUNetBN2C} 13 | export MODEL_N_OUT=${MODEL_N_OUT:-128} 14 | export OPTIMIZER=${OPTIMIZER:-SGD} 15 | export LR=${LR:-1e-1} 16 | export WEIGHT_DECAY=${WEIGHT_DECAY:-1e-4} 17 | export MAX_EPOCH=${MAX_EPOCH:-200} 18 | export BATCH_SIZE=${BATCH_SIZE:-4} 19 | export VAL_BATCH_SIZE=${VAL_BATCH_SIZE:-1} 20 | export ITER_SIZE=${ITER_SIZE:-1} 21 | export BEST_VAL_METRIC=${BEST_VAL_METRIC:-feat_match_ratio} 22 | export VOXEL_SIZE=${VOXEL_SIZE:-0.3} 23 | export POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER=${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER:-1.5} 24 | export CONV1_KERNEL_SIZE=${CONV1_KERNEL_SIZE:-5} 25 | export EXP_GAMMA=${EXP_GAMMA:-0.99} 26 | export RANDOM_SCALE=${RANDOM_SCALE:-True} 27 | export TIME=$(date +"%Y-%m-%d_%H-%M-%S") 28 | export VERSION=$(git rev-parse HEAD) 29 | 30 | export OUT_DIR=${DATA_ROOT}/${DATASET}-v${VOXEL_SIZE}/${TRAINER}/${MODEL}/${OPTIMIZER}-lr${LR}-e${MAX_EPOCH}-b${BATCH_SIZE}i${ITER_SIZE}-modelnout${MODEL_N_OUT}${PATH_POSTFIX}/${TIME} 31 | 32 | # export OUT_DIR=${OUT_DIR:-./outputs/Experiments/PairComplementNuscenesDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-02-15_19-52-13} 33 | export PYTHONUNBUFFERED="True" 34 | 35 | # export RESUME_FILE="./outputs/Experiments/PairComplementNuscenesDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-01-28_13-39-04/checkpoint.pth" 36 | # export RESUME_DIR="./outputs/Experiments/PairComplementNuscenesDataset-v0.3/HardestContrastiveLossTrainer//SGD-lr1e-1-e200-b4i1-modelnout128/2022-01-28_13-39-04" 37 | 38 | echo $OUT_DIR 39 | echo "Using GPU No:" 40 | echo $CUDA_VISIBLE_DEVICES 41 | 42 | mkdir -m 755 -p $OUT_DIR 43 | 44 | LOG=${OUT_DIR}/log_${TIME}.txt 45 | 46 | echo "Host: " $(hostname) | tee -a $LOG 47 | echo "Conda " $(which conda) | tee -a $LOG 48 | echo $(pwd) | tee -a $LOG 49 | echo "Version: " $VERSION | tee -a $LOG 50 | # echo "Git diff" | tee -a $LOG 51 | # echo "" | tee -a $LOG 52 | # git diff | tee -a $LOG 53 | echo "" | tee -a $LOG 54 | nvidia-smi | tee -a $LOG 55 | 56 | # Training With Resume 57 | python train.py \ 58 | --dataset ${DATASET} \ 59 | --trainer ${TRAINER} \ 60 | --model ${ENCODER_MODEL} \ 61 | --model_n_out ${MODEL_N_OUT} \ 62 | --conv1_kernel_size ${CONV1_KERNEL_SIZE} \ 63 | --optimizer ${OPTIMIZER} \ 64 | --lr ${LR} \ 65 | --batch_size ${BATCH_SIZE} \ 66 | --val_batch_size ${VAL_BATCH_SIZE} \ 67 | --iter_size ${ITER_SIZE} \ 68 | --max_epoch ${MAX_EPOCH} \ 69 | --voxel_size ${VOXEL_SIZE} \ 70 | --out_dir ${OUT_DIR} \ 71 | --use_random_scale ${RANDOM_SCALE} \ 72 | --weight_decay ${WEIGHT_DECAY} \ 73 | --positive_pair_search_voxel_size_multiplier ${POSITIVE_PAIR_SEARCH_VOXEL_SIZE_MULTIPLIER} \ 74 | --kitti_root ${KITTI_PATH} \ 75 | --hit_ratio_thresh 0.3 \ 76 | --exp_gamma ${EXP_GAMMA} \ 77 | --complement_pair_dist 10 \ 78 | --num_complement_one_side 4 \ 79 | --point_generation_ratio 4 \ 80 | --regularization_strength 0.01 \ 81 | --regularization_type L2 \ 82 | --best_val_metric ${BEST_VAL_METRIC} \ 83 | --debug_use_old_complement false \ 84 | --use_old_pose true \ 85 | --pair_min_dist 5 \ 86 | --pair_max_dist 20 \ 87 | # --resume ${RESUME_FILE} \ 88 | # --resume_dir ${RESUME_DIR} \ 89 | # --finetune_restart true \ 90 | $MISC_ARGS 2>&1 | tee -a $LOG 91 | 92 | # Test 93 | # python -m scripts.test_kitti \ 94 | # --kitti_root ${KITTI_PATH} \ 95 | # --save_dir ${OUT_DIR} | tee -a $LOG 96 | -------------------------------------------------------------------------------- /FCGF_APR/train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: future_fstrings -*- 2 | import open3d as o3d # prevent loading error 3 | 4 | import sys 5 | import json 6 | import logging 7 | import torch 8 | from easydict import EasyDict as edict 9 | 10 | from lib.complement_data_loader import make_data_loader 11 | from config import get_config 12 | 13 | from lib.trainer import ContrastiveLossTrainer, HardestContrastiveLossTrainer, \ 14 | TripletLossTrainer, HardestTripletLossTrainer 15 | from lib.complement_trainer import GenerativePairTrainer 16 | 17 | 18 | ch = logging.StreamHandler(sys.stdout) 19 | logging.getLogger().setLevel(logging.INFO) 20 | logging.basicConfig( 21 | format='%(asctime)s %(message)s', datefmt='%m/%d %H:%M:%S', handlers=[ch]) 22 | 23 | torch.manual_seed(0) 24 | torch.cuda.manual_seed(0) 25 | 26 | logging.basicConfig(level=logging.INFO, format="") 27 | 28 | 29 | def get_trainer(trainer): 30 | if trainer == 'ContrastiveLossTrainer': 31 | return ContrastiveLossTrainer 32 | elif trainer == 'HardestContrastiveLossTrainer': 33 | return HardestContrastiveLossTrainer 34 | elif trainer == 'TripletLossTrainer': 35 | return TripletLossTrainer 36 | elif trainer == 'HardestTripletLossTrainer': 37 | return HardestTripletLossTrainer 38 | elif trainer == 'GenerativePairTrainer': 39 | return GenerativePairTrainer 40 | else: 41 | raise ValueError(f'Trainer {trainer} not found') 42 | 43 | 44 | def main(config, resume=False): 45 | train_loader = make_data_loader( 46 | config, 47 | config.train_phase, 48 | config.batch_size, 49 | num_threads=config.train_num_thread) 50 | 51 | if config.test_valid: 52 | val_loader = make_data_loader( 53 | config, 54 | config.val_phase, 55 | config.val_batch_size, 56 | num_threads=config.val_num_thread) 57 | else: 58 | val_loader = None 59 | 60 | Trainer = get_trainer(config.trainer) 61 | trainer = Trainer( 62 | config=config, 63 | data_loader=train_loader, 64 | val_data_loader=val_loader, 65 | ) 66 | 67 | trainer.train() 68 | 69 | 70 | if __name__ == "__main__": 71 | logger = logging.getLogger() 72 | config = get_config() 73 | 74 | dconfig = vars(config) 75 | if config.resume_dir and not config.finetune_restart: 76 | resume_config = json.load(open(config.resume_dir + '/config.json', 'r')) 77 | for k in dconfig: 78 | if k not in ['resume_dir'] and k in resume_config: 79 | dconfig[k] = resume_config[k] 80 | dconfig['resume'] = resume_config['out_dir'] + '/checkpoint.pth' 81 | 82 | logging.info('===> Configurations') 83 | for k in dconfig: 84 | logging.info(' {}: {}'.format(k, dconfig[k])) 85 | 86 | # Convert to dict 87 | config = edict(dconfig) 88 | main(config) 89 | -------------------------------------------------------------------------------- /FCGF_APR/util/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/FCGF_APR/util/__init__.py -------------------------------------------------------------------------------- /FCGF_APR/util/file.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from os import listdir 4 | from os.path import isfile, isdir, join, splitext 5 | 6 | 7 | def read_txt(path): 8 | """Read txt file into lines. 9 | """ 10 | with open(path) as f: 11 | lines = f.readlines() 12 | lines = [x.strip() for x in lines] 13 | return lines 14 | 15 | 16 | def ensure_dir(path): 17 | if not os.path.exists(path): 18 | os.makedirs(path, mode=0o755) 19 | 20 | 21 | def sorted_alphanum(file_list_ordered): 22 | 23 | def convert(text): 24 | return int(text) if text.isdigit() else text 25 | 26 | def alphanum_key(key): 27 | return [convert(c) for c in re.split('([0-9]+)', key)] 28 | 29 | return sorted(file_list_ordered, key=alphanum_key) 30 | 31 | 32 | def get_file_list(path, extension=None): 33 | if extension is None: 34 | file_list = [join(path, f) for f in listdir(path) if isfile(join(path, f))] 35 | else: 36 | file_list = [ 37 | join(path, f) 38 | for f in listdir(path) 39 | if isfile(join(path, f)) and splitext(f)[1] == extension 40 | ] 41 | file_list = sorted_alphanum(file_list) 42 | return file_list 43 | 44 | 45 | def get_file_list_specific(path, color_depth, extension=None): 46 | if extension is None: 47 | file_list = [join(path, f) for f in listdir(path) if isfile(join(path, f))] 48 | else: 49 | file_list = [ 50 | join(path, f) 51 | for f in listdir(path) 52 | if isfile(join(path, f)) and color_depth in f and splitext(f)[1] == extension 53 | ] 54 | file_list = sorted_alphanum(file_list) 55 | return file_list 56 | 57 | 58 | def get_folder_list(path): 59 | folder_list = [join(path, f) for f in listdir(path) if isdir(join(path, f))] 60 | folder_list = sorted_alphanum(folder_list) 61 | return folder_list 62 | -------------------------------------------------------------------------------- /FCGF_APR/util/misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import MinkowskiEngine as ME 4 | 5 | 6 | def _hash(arr, M): 7 | if isinstance(arr, np.ndarray): 8 | N, D = arr.shape 9 | else: 10 | N, D = len(arr[0]), len(arr) 11 | 12 | hash_vec = np.zeros(N, dtype=np.int64) 13 | for d in range(D): 14 | if isinstance(arr, np.ndarray): 15 | hash_vec += arr[:, d] * M**d 16 | else: 17 | hash_vec += arr[d] * M**d 18 | return hash_vec 19 | 20 | 21 | def extract_features(model, 22 | xyz, 23 | rgb=None, 24 | normal=None, 25 | voxel_size=0.05, 26 | device=None, 27 | skip_check=False, 28 | is_eval=True): 29 | ''' 30 | xyz is a N x 3 matrix 31 | rgb is a N x 3 matrix and all color must range from [0, 1] or None 32 | normal is a N x 3 matrix and all normal range from [-1, 1] or None 33 | 34 | if both rgb and normal are None, we use Nx1 one vector as an input 35 | 36 | if device is None, it tries to use gpu by default 37 | 38 | if skip_check is True, skip rigorous checks to speed up 39 | 40 | model = model.to(device) 41 | xyz, feats = extract_features(model, xyz) 42 | ''' 43 | if is_eval: 44 | model.eval() 45 | 46 | if not skip_check: 47 | assert xyz.shape[1] == 3 48 | 49 | N = xyz.shape[0] 50 | if rgb is not None: 51 | assert N == len(rgb) 52 | assert rgb.shape[1] == 3 53 | if np.any(rgb > 1): 54 | raise ValueError('Invalid color. Color must range from [0, 1]') 55 | 56 | if normal is not None: 57 | assert N == len(normal) 58 | assert normal.shape[1] == 3 59 | if np.any(normal > 1): 60 | raise ValueError('Invalid normal. Normal must range from [-1, 1]') 61 | 62 | if device is None: 63 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 64 | 65 | feats = [] 66 | if rgb is not None: 67 | # [0, 1] 68 | feats.append(rgb - 0.5) 69 | 70 | if normal is not None: 71 | # [-1, 1] 72 | feats.append(normal / 2) 73 | 74 | if rgb is None and normal is None: 75 | feats.append(np.ones((len(xyz), 1))) 76 | 77 | feats = np.hstack(feats) 78 | 79 | # Voxelize xyz and feats 80 | coords = np.floor(xyz / voxel_size) 81 | coords, inds = ME.utils.sparse_quantize(coords, return_index=True) 82 | # Convert to batched coords compatible with ME 83 | coords = ME.utils.batched_coordinates([coords]) 84 | return_coords = xyz[inds] 85 | 86 | feats = feats[inds] 87 | 88 | feats = torch.tensor(feats, dtype=torch.float32) 89 | coords = torch.tensor(coords, dtype=torch.int32) 90 | 91 | stensor = ME.SparseTensor(feats, coordinates=coords, device=device) 92 | 93 | return return_coords, model(stensor).F 94 | -------------------------------------------------------------------------------- /FCGF_APR/util/pointcloud.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import numpy as np 3 | import math 4 | 5 | import open3d as o3d 6 | from lib.eval import find_nn_cpu 7 | 8 | 9 | def make_open3d_point_cloud(xyz, color=None): 10 | pcd = o3d.geometry.PointCloud() 11 | pcd.points = o3d.utility.Vector3dVector(xyz) 12 | if color is not None: 13 | pcd.colors = o3d.utility.Vector3dVector(color) 14 | return pcd 15 | 16 | 17 | def make_open3d_feature(data, dim, npts): 18 | feature = o3d.pipelines.registration.Feature() 19 | feature.resize(dim, npts) 20 | feature.data = data.cpu().numpy().astype('d').transpose() 21 | return feature 22 | 23 | 24 | def make_open3d_feature_from_numpy(data): 25 | assert isinstance(data, np.ndarray) 26 | assert data.ndim == 2 27 | 28 | feature = o3d.pipelines.registration.Feature() 29 | feature.resize(data.shape[1], data.shape[0]) 30 | feature.data = data.astype('d').transpose() 31 | return feature 32 | 33 | 34 | def prepare_pointcloud(filename, voxel_size): 35 | pcd = o3d.io.read_point_cloud(filename) 36 | T = get_random_transformation(pcd) 37 | pcd.transform(T) 38 | pcd_down = pcd.voxel_down_sample(voxel_size) 39 | return pcd_down, T 40 | 41 | 42 | def compute_overlap_ratio(pcd0, pcd1, trans, voxel_size): 43 | pcd0_down = pcd0.voxel_down_sample(voxel_size) 44 | pcd1_down = pcd1.voxel_down_sample(voxel_size) 45 | matching01 = get_matching_indices(pcd0_down, pcd1_down, trans, voxel_size, 1) 46 | matching10 = get_matching_indices(pcd1_down, pcd0_down, np.linalg.inv(trans), 47 | voxel_size, 1) 48 | overlap0 = len(matching01) / len(pcd0_down.points) 49 | overlap1 = len(matching10) / len(pcd1_down.points) 50 | return max(overlap0, overlap1) 51 | 52 | 53 | def get_matching_indices(source, target, trans, search_voxel_size, K=None): 54 | source_copy = copy.deepcopy(source) 55 | target_copy = copy.deepcopy(target) 56 | source_copy.transform(trans) 57 | pcd_tree = o3d.geometry.KDTreeFlann(target_copy) 58 | 59 | match_inds = [] 60 | for i, point in enumerate(source_copy.points): 61 | [_, idx, _] = pcd_tree.search_radius_vector_3d(point, search_voxel_size) 62 | if K is not None: 63 | idx = idx[:K] 64 | for j in idx: 65 | match_inds.append((i, j)) 66 | return match_inds 67 | 68 | 69 | def evaluate_feature(pcd0, pcd1, feat0, feat1, trans_gth, search_voxel_size): 70 | match_inds = get_matching_indices(pcd0, pcd1, trans_gth, search_voxel_size) 71 | pcd_tree = o3d.geometry.KDTreeFlann(feat1) 72 | dist = [] 73 | for ind in match_inds: 74 | k, idx, _ = pcd_tree.search_knn_vector_xd(feat0.data[:, ind[0]], 1) 75 | dist.append( 76 | np.clip( 77 | np.power(pcd1.points[ind[1]] - pcd1.points[idx[0]], 2), 78 | a_min=0.0, 79 | a_max=1.0)) 80 | return np.mean(dist) 81 | 82 | 83 | def valid_feat_ratio(pcd0, pcd1, feat0, feat1, trans_gth, thresh=0.1): 84 | pcd0_copy = copy.deepcopy(pcd0) 85 | pcd0_copy.transform(trans_gth) 86 | inds = find_nn_cpu(feat0, feat1, return_distance=False) 87 | dist = np.sqrt(((np.array(pcd0_copy.points) - np.array(pcd1.points)[inds])**2).sum(1)) 88 | return np.mean(dist < thresh) 89 | 90 | 91 | def evaluate_feature_3dmatch(pcd0, pcd1, feat0, feat1, trans_gth, inlier_thresh=0.1): 92 | r"""Return the hit ratio (ratio of inlier correspondences and all correspondences). 93 | 94 | inliear_thresh is the inlier_threshold in meter. 95 | """ 96 | if len(pcd0.points) < len(pcd1.points): 97 | hit = valid_feat_ratio(pcd0, pcd1, feat0, feat1, trans_gth, inlier_thresh) 98 | else: 99 | hit = valid_feat_ratio(pcd1, pcd0, feat1, feat0, np.linalg.inv(trans_gth), inlier_thresh) 100 | return hit 101 | 102 | 103 | def get_matching_matrix(source, target, trans, voxel_size, debug_mode): 104 | source_copy = copy.deepcopy(source) 105 | target_copy = copy.deepcopy(target) 106 | source_copy.transform(trans) 107 | pcd_tree = o3d.geometry.KDTreeFlann(target_copy) 108 | matching_matrix = np.zeros((len(source_copy.points), len(target_copy.points))) 109 | 110 | for i, point in enumerate(source_copy.points): 111 | [k, idx, _] = pcd_tree.search_radius_vector_3d(point, voxel_size * 1.5) 112 | if k >= 1: 113 | matching_matrix[i, idx[0]] = 1 # TODO: only the cloest? 114 | 115 | return matching_matrix 116 | 117 | 118 | def get_random_transformation(pcd_input): 119 | 120 | def rot_x(x): 121 | out = np.zeros((3, 3)) 122 | c = math.cos(x) 123 | s = math.sin(x) 124 | out[0, 0] = 1 125 | out[1, 1] = c 126 | out[1, 2] = -s 127 | out[2, 1] = s 128 | out[2, 2] = c 129 | return out 130 | 131 | def rot_y(x): 132 | out = np.zeros((3, 3)) 133 | c = math.cos(x) 134 | s = math.sin(x) 135 | out[0, 0] = c 136 | out[0, 2] = s 137 | out[1, 1] = 1 138 | out[2, 0] = -s 139 | out[2, 2] = c 140 | return out 141 | 142 | def rot_z(x): 143 | out = np.zeros((3, 3)) 144 | c = math.cos(x) 145 | s = math.sin(x) 146 | out[0, 0] = c 147 | out[0, 1] = -s 148 | out[1, 0] = s 149 | out[1, 1] = c 150 | out[2, 2] = 1 151 | return out 152 | 153 | pcd_output = copy.deepcopy(pcd_input) 154 | mean = np.mean(np.asarray(pcd_output.points), axis=0).transpose() 155 | xyz = np.random.uniform(0, 2 * math.pi, 3) 156 | R = np.dot(np.dot(rot_x(xyz[0]), rot_y(xyz[1])), rot_z(xyz[2])) 157 | T = np.zeros((4, 4)) 158 | T[:3, :3] = R 159 | T[:3, 3] = np.dot(-R, mean) 160 | T[3, 3] = 1 161 | return T 162 | -------------------------------------------------------------------------------- /FCGF_APR/util/trajectory.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | 5 | class CameraPose: 6 | 7 | def __init__(self, meta, mat): 8 | self.metadata = meta 9 | self.pose = mat 10 | 11 | def __str__(self): 12 | return 'metadata : ' + ' '.join(map(str, self.metadata)) + '\n' + \ 13 | "pose : " + "\n" + np.array_str(self.pose) 14 | 15 | 16 | def read_trajectory(filename, dim=4): 17 | traj = [] 18 | assert os.path.exists(filename) 19 | with open(filename, 'r') as f: 20 | metastr = f.readline() 21 | while metastr: 22 | metadata = list(map(int, metastr.split())) 23 | mat = np.zeros(shape=(dim, dim)) 24 | for i in range(dim): 25 | matstr = f.readline() 26 | mat[i, :] = np.fromstring(matstr, dtype=float, sep=' \t') 27 | traj.append(CameraPose(metadata, mat)) 28 | metastr = f.readline() 29 | return traj 30 | 31 | 32 | def write_trajectory(traj, filename, dim=4): 33 | with open(filename, 'w') as f: 34 | for x in traj: 35 | p = x.pose.tolist() 36 | f.write(' '.join(map(str, x.metadata)) + '\n') 37 | f.write('\n'.join(' '.join(map('{0:.12f}'.format, p[i])) for i in range(dim))) 38 | f.write('\n') 39 | -------------------------------------------------------------------------------- /FCGF_APR/util/transform_estimation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import MinkowskiEngine as ME 3 | 4 | 5 | def rot_x(x): 6 | out = torch.zeros((3, 3)) 7 | c = torch.cos(x) 8 | s = torch.sin(x) 9 | out[0, 0] = 1 10 | out[1, 1] = c 11 | out[1, 2] = -s 12 | out[2, 1] = s 13 | out[2, 2] = c 14 | return out 15 | 16 | 17 | def rot_y(x): 18 | out = torch.zeros((3, 3)) 19 | c = torch.cos(x) 20 | s = torch.sin(x) 21 | out[0, 0] = c 22 | out[0, 2] = s 23 | out[1, 1] = 1 24 | out[2, 0] = -s 25 | out[2, 2] = c 26 | return out 27 | 28 | 29 | def rot_z(x): 30 | out = torch.zeros((3, 3)) 31 | c = torch.cos(x) 32 | s = torch.sin(x) 33 | out[0, 0] = c 34 | out[0, 1] = -s 35 | out[1, 0] = s 36 | out[1, 1] = c 37 | out[2, 2] = 1 38 | return out 39 | 40 | 41 | def get_trans(x): 42 | trans = torch.eye(4) 43 | trans[:3, :3] = rot_z(x[2]).mm(rot_y(x[1])).mm(rot_x(x[0])) 44 | trans[:3, 3] = x[3:, 0] 45 | return trans 46 | 47 | 48 | def update_pcd(pts, trans): 49 | R = trans[:3, :3] 50 | T = trans[:3, 3] 51 | # pts = R.mm(pts.t()).t() + T.unsqueeze(1).t().expand_as(pts) 52 | pts = torch.t(R @ torch.t(pts)) + T 53 | return pts 54 | 55 | 56 | def build_linear_system(pts0, pts1, weight): 57 | npts0 = pts0.shape[0] 58 | A0 = torch.zeros((npts0, 6)) 59 | A1 = torch.zeros((npts0, 6)) 60 | A2 = torch.zeros((npts0, 6)) 61 | A0[:, 1] = pts0[:, 2] 62 | A0[:, 2] = -pts0[:, 1] 63 | A0[:, 3] = 1 64 | A1[:, 0] = -pts0[:, 2] 65 | A1[:, 2] = pts0[:, 0] 66 | A1[:, 4] = 1 67 | A2[:, 0] = pts0[:, 1] 68 | A2[:, 1] = -pts0[:, 0] 69 | A2[:, 5] = 1 70 | ww1 = weight.repeat(3, 6) 71 | ww2 = weight.repeat(3, 1) 72 | A = ww1 * torch.cat((A0, A1, A2), 0) 73 | b = ww2 * torch.cat( 74 | (pts1[:, 0] - pts0[:, 0], pts1[:, 1] - pts0[:, 1], pts1[:, 2] - pts0[:, 2]), 75 | 0, 76 | ).unsqueeze(1) 77 | return A, b 78 | 79 | 80 | def solve_linear_system(A, b): 81 | temp = torch.inverse(A.t().mm(A)) 82 | return temp.mm(A.t()).mm(b) 83 | 84 | 85 | def compute_weights(pts0, pts1, par): 86 | return par / (torch.norm(pts0 - pts1, dim=1).unsqueeze(1) + par) 87 | 88 | 89 | def est_quad_linear_robust(pts0, pts1, weight=None): 90 | # TODO: 2. residual scheduling 91 | pts0_curr = pts0 92 | trans = torch.eye(4) 93 | 94 | par = 1.0 # todo: need to decide 95 | if weight is None: 96 | weight = torch.ones(pts0.size()[0], 1) 97 | 98 | for i in range(20): 99 | if i > 0 and i % 5 == 0: 100 | par /= 2.0 101 | 102 | # compute weights 103 | A, b = build_linear_system(pts0_curr, pts1, weight) 104 | x = solve_linear_system(A, b) 105 | 106 | # TODO: early termination 107 | # residual = np.linalg.norm(A@x - b) 108 | # print(residual) 109 | 110 | # x = torch.empty(6, 1).uniform_(0, 1) 111 | trans_curr = get_trans(x) 112 | pts0_curr = update_pcd(pts0_curr, trans_curr) 113 | weight = compute_weights(pts0_curr, pts1, par) 114 | trans = trans_curr.mm(trans) 115 | 116 | return trans 117 | 118 | 119 | def pose_estimation(model, 120 | device, 121 | xyz0, 122 | xyz1, 123 | coord0, 124 | coord1, 125 | feats0, 126 | feats1, 127 | return_corr=False): 128 | sinput0 = ME.SparseTensor(feats0.to(device), coordinates=coord0.to(device)) 129 | F0 = model(sinput0).F 130 | 131 | sinput1 = ME.SparseTensor(feats1.to(device), coordinates=coord1.to(device)) 132 | F1 = model(sinput1).F 133 | 134 | corr = F0.mm(F1.t()) 135 | weight, inds = corr.max(dim=1) 136 | weight = weight.unsqueeze(1).cpu() 137 | xyz1_corr = xyz1[inds, :] 138 | 139 | trans = est_quad_linear_robust(xyz0, xyz1_corr, weight) # let's do this later 140 | 141 | if return_corr: 142 | return trans, weight, corr 143 | else: 144 | return trans, weight 145 | -------------------------------------------------------------------------------- /FCGF_APR/util/visualization.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import open3d as o3d 3 | import numpy as np 4 | 5 | from sklearn.manifold import TSNE 6 | from matplotlib import pyplot as plt 7 | 8 | 9 | def get_color_map(x): 10 | colours = plt.cm.Spectral(x) 11 | return colours[:, :3] 12 | 13 | 14 | def mesh_sphere(pcd, voxel_size, sphere_size=0.6): 15 | # Create a mesh sphere 16 | spheres = o3d.geometry.TriangleMesh() 17 | s = o3d.geometry.TriangleMesh.create_sphere(radius=voxel_size * sphere_size) 18 | s.compute_vertex_normals() 19 | 20 | for i, p in enumerate(pcd.points): 21 | si = copy.deepcopy(s) 22 | trans = np.identity(4) 23 | trans[:3, 3] = p 24 | si.transform(trans) 25 | si.paint_uniform_color(pcd.colors[i]) 26 | spheres += si 27 | return spheres 28 | 29 | 30 | def get_colored_point_cloud_feature(pcd, feature, voxel_size): 31 | tsne_results = embed_tsne(feature) 32 | 33 | color = get_color_map(tsne_results) 34 | pcd.colors = o3d.utility.Vector3dVector(color) 35 | spheres = mesh_sphere(pcd, voxel_size) 36 | 37 | return spheres 38 | 39 | 40 | def embed_tsne(data): 41 | """ 42 | N x D np.array data 43 | """ 44 | tsne = TSNE(n_components=1, verbose=1, perplexity=40, n_iter=300, random_state=0) 45 | tsne_results = tsne.fit_transform(data) 46 | tsne_results = np.squeeze(tsne_results) 47 | tsne_min = np.min(tsne_results) 48 | tsne_max = np.max(tsne_results) 49 | return (tsne_results - tsne_min) / (tsne_max - tsne_min) 50 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 liuQuan98 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 | -------------------------------------------------------------------------------- /Predator_APR/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in all 11 | copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | SOFTWARE. 20 | -------------------------------------------------------------------------------- /Predator_APR/README.md: -------------------------------------------------------------------------------- 1 | # APR: Online Distant Point Cloud Registration Through Aggregated Point Cloud Reconstruction, Implemented with PREDATOR: Registration of 3D Point Clouds with Low Overlap ([Predator](https://github.com/overlappredator/OverlapPredator)) 2 | 3 | For many driving safety applications, it is of great importance to accurately register LiDAR point clouds generated on distant moving vehicles. However, such point clouds have extremely different point density and sensor perspective on the same object, making registration on such point clouds very hard. In this paper, we propose a novel feature extraction framework, called APR, for online distant point cloud registration. Specifically, APR leverages an autoencoder design, where the autoencoder reconstructs a denser aggregated point cloud with several frames instead of the original single input point cloud. Our design forces the encoder to extract features with rich local geometry information based on one single input point cloud. Such features are then used for online distant point cloud registration. We conduct extensive experiments against state-of-the-art (SOTA) feature extractors on KITTI and nuScenes datasets. Results show that APR outperforms all other extractors by a large margin, increasing average registration recall of SOTA extractors by 7.1% on LoKITTI and 4.6% on LoNuScenes. 4 | 5 | This repository is an implementation of APR using Predator as the feature extractor. 6 | 7 | ## Requirements 8 | 9 | - Ubuntu 14.04 or higher 10 | - CUDA 11.1 or higher 11 | - Python v3.7 or higher 12 | - Pytorch v1.6 or higher 13 | - [MinkowskiEngine](https://github.com/stanfordvl/MinkowskiEngine) v0.5 or higher 14 | 15 | ## Training & Testing 16 | 17 | First, use Predator_APR as the working directory: 18 | 19 | ``` 20 | cd ./Predator_APR 21 | conda activate apr 22 | ``` 23 | 24 | If you haven't downloaded any of the datasets (KITTI and nuScenes) according to our specification, please refer to the README.md in the [parent directory](../README.md). 25 | 26 | ### Setting the distance between two LiDARs (registration difficulty) 27 | 28 | As the major focus of this paper, we divide the registration datasets into different slices according to the registration difficulty, i.e., the distance $d$ between two LiDARs. Greater $d$ leads to a smaller overlap and more divergent point density, resulting in a higher registration difficulty. We denote range of $d$ with the parameter `pair_min_dist` and `pair_max_dist`, which can be found in `./configs/{$phase}/{$dataset}.yaml`. For example, setting 29 | 30 | ``` 31 | pair_min_dist: 5 32 | pair_max_dist: 20 33 | ``` 34 | 35 | will set $d\in [5m,20m]$. In other words, for every pair of point clouds, the ground-truth euclidean distance betwen two corresponding LiDAR positions (i.e., the origins of the two specified point clouds) obeys a uniform distribution between 5m and 20m. 36 | 37 | ### Training suggestions 38 | 39 | For cases where you want `pair_max_dist` to be larger than 20, we recommend following the two-stage training paradigm as pointed out in Section 5 of our paper: 40 | 41 | 1. Pretrain a model with the following distance parameters: `pair_min_dist: 5` and `pair_max_dist: 20`. Find out the converged model checkpoint file path according to your training config. It shoud be some path like this: `./snapshot/{$exp_dir}/checkpoints/model_best_recall.pth` 42 | 2. Finetune a new model on `pair_min_dist 5` and `pair_max_dist {$YOUR_SPECIFIC_DISTANCE}`, while setting the pretrained checkpoint file path in `pretrain: "{$PRETRAINED_CKPT_PATH}"` in the config file `./configs/{$phase}/{$dataset}.yaml`. Do not forget to set `pretrain_restart: True`. 43 | 44 | Emperically, the pretraining strategy helps a lot in model convergence especially when the distance is large; Otherwise the model just diverges. 45 | 46 | ### Launch the training 47 | 48 | Specify the GPU usage in `main.py`, then train Predator-APR with either of the following command inside conda environment `apr`: 49 | 50 | ``` 51 | python main.py ./configs/train/kitti.yaml 52 | python main.py ./configs/train/nuscenes.yaml 53 | ``` 54 | 55 | Note: The symmetric APR setup is not supported with Predator backbone due to GPU memory issue. 56 | 57 | ### Testing 58 | 59 | To test a specific model, please constitute the model directory into the 'pretrain' section in the `./config/test/{$dataset}.yaml`, then run either of the following: 60 | 61 | ``` 62 | python main.py configs/test/kitti.yaml 63 | python main.py configs/test/nuscenes.yaml 64 | ``` 65 | 66 | ## Pre-trained models 67 | 68 | We provide our [model](https://drive.google.com/file/d/1mLqiahQMgYMRyB4XKhp-HJdy5yavL2fj/view?usp=sharing) trained on Predator+APR with different point cloud distance. To reproduce the results using our model, please extract the model checkpoints to the 'snapshot' directory. 69 | -------------------------------------------------------------------------------- /Predator_APR/common/colors.py: -------------------------------------------------------------------------------- 1 | """Useful color codes""" 2 | ORANGE = [239, 124, 0] 3 | BLUE = [0, 61, 124] -------------------------------------------------------------------------------- /Predator_APR/common/math/random.py: -------------------------------------------------------------------------------- 1 | """Functions for random sampling""" 2 | import numpy as np 3 | 4 | 5 | def uniform_2_sphere(num: int = None): 6 | """Uniform sampling on a 2-sphere 7 | 8 | Source: https://gist.github.com/andrewbolster/10274979 9 | 10 | Args: 11 | num: Number of vectors to sample (or None if single) 12 | 13 | Returns: 14 | Random Vector (np.ndarray) of size (num, 3) with norm 1. 15 | If num is None returned value will have size (3,) 16 | 17 | """ 18 | if num is not None: 19 | phi = np.random.uniform(0.0, 2 * np.pi, num) 20 | cos_theta = np.random.uniform(-1.0, 1.0, num) 21 | else: 22 | phi = np.random.uniform(0.0, 2 * np.pi) 23 | cos_theta = np.random.uniform(-1.0, 1.0) 24 | 25 | theta = np.arccos(cos_theta) 26 | x = np.sin(theta) * np.cos(phi) 27 | y = np.sin(theta) * np.sin(phi) 28 | z = np.cos(theta) 29 | 30 | return np.stack((x, y, z), axis=-1) 31 | 32 | 33 | if __name__ == '__main__': 34 | # Visualize sampling 35 | from vtk_visualizer.plot3d import plotxyz 36 | rand_2s = uniform_2_sphere(10000) 37 | plotxyz(rand_2s, block=True) 38 | -------------------------------------------------------------------------------- /Predator_APR/common/math/se3.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.spatial.transform import Rotation 3 | 4 | 5 | def identity(): 6 | return np.eye(3, 4) 7 | 8 | 9 | def transform(g: np.ndarray, pts: np.ndarray): 10 | """ Applies the SE3 transform 11 | 12 | Args: 13 | g: SE3 transformation matrix of size ([B,] 3/4, 4) 14 | pts: Points to be transformed ([B,] N, 3) 15 | 16 | Returns: 17 | transformed points of size (N, 3) 18 | """ 19 | rot = g[..., :3, :3] # (3, 3) 20 | trans = g[..., :3, 3] # (3) 21 | 22 | transformed = pts[..., :3] @ np.swapaxes(rot, -1, -2) + trans[..., None, :] 23 | return transformed 24 | 25 | 26 | def inverse(g: np.ndarray): 27 | """Returns the inverse of the SE3 transform 28 | 29 | Args: 30 | g: ([B,] 3/4, 4) transform 31 | 32 | Returns: 33 | ([B,] 3/4, 4) matrix containing the inverse 34 | 35 | """ 36 | rot = g[..., :3, :3] # (3, 3) 37 | trans = g[..., :3, 3] # (3) 38 | 39 | inv_rot = np.swapaxes(rot, -1, -2) 40 | inverse_transform = np.concatenate([inv_rot, inv_rot @ -trans[..., None]], axis=-1) 41 | if g.shape[-2] == 4: 42 | inverse_transform = np.concatenate([inverse_transform, [[0.0, 0.0, 0.0, 1.0]]], axis=-2) 43 | 44 | return inverse_transform 45 | 46 | 47 | def concatenate(a: np.ndarray, b: np.ndarray): 48 | """ Concatenate two SE3 transforms 49 | 50 | Args: 51 | a: First transform ([B,] 3/4, 4) 52 | b: Second transform ([B,] 3/4, 4) 53 | 54 | Returns: 55 | a*b ([B, ] 3/4, 4) 56 | 57 | """ 58 | 59 | r_a, t_a = a[..., :3, :3], a[..., :3, 3] 60 | r_b, t_b = b[..., :3, :3], b[..., :3, 3] 61 | 62 | r_ab = r_a @ r_b 63 | t_ab = r_a @ t_b[..., None] + t_a[..., None] 64 | 65 | concatenated = np.concatenate([r_ab, t_ab], axis=-1) 66 | 67 | if a.shape[-2] == 4: 68 | concatenated = np.concatenate([concatenated, [[0.0, 0.0, 0.0, 1.0]]], axis=-2) 69 | 70 | return concatenated 71 | 72 | 73 | def from_xyzquat(xyzquat): 74 | """Constructs SE3 matrix from x, y, z, qx, qy, qz, qw 75 | 76 | Args: 77 | xyzquat: np.array (7,) containing translation and quaterion 78 | 79 | Returns: 80 | SE3 matrix (4, 4) 81 | """ 82 | rot = Rotation.from_quat(xyzquat[3:]) 83 | trans = rot.apply(-xyzquat[:3]) 84 | transform = np.concatenate([rot.as_dcm(), trans[:, None]], axis=1) 85 | transform = np.concatenate([transform, [[0.0, 0.0, 0.0, 1.0]]], axis=0) 86 | 87 | return transform -------------------------------------------------------------------------------- /Predator_APR/common/math/so3.py: -------------------------------------------------------------------------------- 1 | """ 2 | Rotation related functions for numpy arrays 3 | """ 4 | 5 | import numpy as np 6 | from scipy.spatial.transform import Rotation 7 | 8 | 9 | def dcm2euler(mats: np.ndarray, seq: str = 'zyx', degrees: bool = True): 10 | """Converts rotation matrix to euler angles 11 | 12 | Args: 13 | mats: (B, 3, 3) containing the B rotation matricecs 14 | seq: Sequence of euler rotations (default: 'zyx') 15 | degrees (bool): If true (default), will return in degrees instead of radians 16 | 17 | Returns: 18 | 19 | """ 20 | 21 | eulers = [] 22 | for i in range(mats.shape[0]): 23 | r = Rotation.from_dcm(mats[i]) 24 | eulers.append(r.as_euler(seq, degrees=degrees)) 25 | return np.stack(eulers) 26 | 27 | 28 | def transform(g: np.ndarray, pts: np.ndarray): 29 | """ Applies the SO3 transform 30 | 31 | Args: 32 | g: SO3 transformation matrix of size (3, 3) 33 | pts: Points to be transformed (N, 3) 34 | 35 | Returns: 36 | transformed points of size (N, 3) 37 | 38 | """ 39 | rot = g[:3, :3] # (3, 3) 40 | transformed = pts @ rot.transpose() 41 | return transformed 42 | -------------------------------------------------------------------------------- /Predator_APR/common/math_torch/se3.py: -------------------------------------------------------------------------------- 1 | """ 3-d rigid body transformation group 2 | """ 3 | import torch 4 | 5 | 6 | def identity(batch_size): 7 | return torch.eye(3, 4)[None, ...].repeat(batch_size, 1, 1) 8 | 9 | 10 | def inverse(g): 11 | """ Returns the inverse of the SE3 transform 12 | 13 | Args: 14 | g: (B, 3/4, 4) transform 15 | 16 | Returns: 17 | (B, 3, 4) matrix containing the inverse 18 | 19 | """ 20 | # Compute inverse 21 | rot = g[..., 0:3, 0:3] 22 | trans = g[..., 0:3, 3] 23 | inverse_transform = torch.cat([rot.transpose(-1, -2), rot.transpose(-1, -2) @ -trans[..., None]], dim=-1) 24 | 25 | return inverse_transform 26 | 27 | 28 | def concatenate(a, b): 29 | """Concatenate two SE3 transforms, 30 | i.e. return a@b (but note that our SE3 is represented as a 3x4 matrix) 31 | 32 | Args: 33 | a: (B, 3/4, 4) 34 | b: (B, 3/4, 4) 35 | 36 | Returns: 37 | (B, 3/4, 4) 38 | """ 39 | 40 | rot1 = a[..., :3, :3] 41 | trans1 = a[..., :3, 3] 42 | rot2 = b[..., :3, :3] 43 | trans2 = b[..., :3, 3] 44 | 45 | rot_cat = rot1 @ rot2 46 | trans_cat = rot1 @ trans2[..., None] + trans1[..., None] 47 | concatenated = torch.cat([rot_cat, trans_cat], dim=-1) 48 | 49 | return concatenated 50 | 51 | 52 | def transform(g, a, normals=None): 53 | """ Applies the SE3 transform 54 | 55 | Args: 56 | g: SE3 transformation matrix of size ([1,] 3/4, 4) or (B, 3/4, 4) 57 | a: Points to be transformed (N, 3) or (B, N, 3) 58 | normals: (Optional). If provided, normals will be transformed 59 | 60 | Returns: 61 | transformed points of size (N, 3) or (B, N, 3) 62 | 63 | """ 64 | R = g[..., :3, :3] # (B, 3, 3) 65 | p = g[..., :3, 3] # (B, 3) 66 | 67 | if len(g.size()) == len(a.size()): 68 | b = torch.matmul(a, R.transpose(-1, -2)) + p[..., None, :] 69 | else: 70 | raise NotImplementedError 71 | b = R.matmul(a.unsqueeze(-1)).squeeze(-1) + p # No batch. Not checked 72 | 73 | if normals is not None: 74 | rotated_normals = normals @ R.transpose(-1, -2) 75 | return b, rotated_normals 76 | 77 | else: 78 | return b 79 | -------------------------------------------------------------------------------- /Predator_APR/common/misc.py: -------------------------------------------------------------------------------- 1 | """ 2 | Misc utilities 3 | """ 4 | 5 | import argparse 6 | from datetime import datetime 7 | import logging 8 | import os 9 | import shutil 10 | import subprocess 11 | import sys 12 | 13 | import coloredlogs 14 | import git 15 | 16 | 17 | _logger = logging.getLogger() 18 | 19 | 20 | def print_info(opt, log_dir=None): 21 | """ Logs source code configuration 22 | """ 23 | _logger.info('Command: {}'.format(' '.join(sys.argv))) 24 | 25 | # Print commit ID 26 | try: 27 | repo = git.Repo(search_parent_directories=True) 28 | git_sha = repo.head.object.hexsha 29 | git_date = datetime.fromtimestamp(repo.head.object.committed_date).strftime('%Y-%m-%d') 30 | git_message = repo.head.object.message 31 | _logger.info('Source is from Commit {} ({}): {}'.format(git_sha[:8], git_date, git_message.strip())) 32 | 33 | # Also create diff file in the log directory 34 | if log_dir is not None: 35 | with open(os.path.join(log_dir, 'compareHead.diff'), 'w') as fid: 36 | subprocess.run(['git', 'diff'], stdout=fid) 37 | 38 | except git.exc.InvalidGitRepositoryError: 39 | pass 40 | 41 | # Arguments 42 | arg_str = ['{}: {}'.format(key, value) for key, value in vars(opt).items()] 43 | arg_str = ', '.join(arg_str) 44 | #_logger.info('Arguments: {}'.format(arg_str)) 45 | 46 | 47 | def prepare_logger(opt: argparse.Namespace, log_path: str = None): 48 | """Creates logging directory, and installs colorlogs 49 | 50 | Args: 51 | opt: Program arguments, should include --dev and --logdir flag. 52 | See get_parent_parser() 53 | log_path: Logging path (optional). This serves to overwrite the settings in 54 | argparse namespace 55 | 56 | Returns: 57 | logger (logging.Logger) 58 | log_path (str): Logging directory 59 | """ 60 | 61 | if log_path is None: 62 | if opt.dev: 63 | log_path = '../logdev' 64 | shutil.rmtree(log_path, ignore_errors=True) 65 | else: 66 | datetime_str = datetime.now().strftime('%y%m%d_%H%M%S') 67 | if opt.name is not None: 68 | log_path = os.path.join(opt.logdir, datetime_str + '_' + opt.name) 69 | else: 70 | log_path = os.path.join(opt.logdir, datetime_str) 71 | 72 | os.makedirs(log_path, exist_ok=True) 73 | logger = logging.getLogger() 74 | coloredlogs.install(level='INFO', logger=logger) 75 | file_handler = logging.FileHandler('{}/log.txt'.format(log_path)) 76 | log_formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s - %(message)s') 77 | file_handler.setFormatter(log_formatter) 78 | logger.addHandler(file_handler) 79 | print_info(opt, log_path) 80 | logger.info('Output and logs will be saved to {}'.format(log_path)) 81 | 82 | return logger, log_path 83 | -------------------------------------------------------------------------------- /Predator_APR/common/torch.py: -------------------------------------------------------------------------------- 1 | """PyTorch related utility functions 2 | """ 3 | 4 | import logging 5 | import os 6 | import pdb 7 | import shutil 8 | import sys 9 | import time 10 | import traceback 11 | 12 | import numpy as np 13 | import torch 14 | from torch.optim.optimizer import Optimizer 15 | 16 | 17 | def dict_all_to_device(tensor_dict, device): 18 | """Sends everything into a certain device """ 19 | for k in tensor_dict: 20 | if isinstance(tensor_dict[k], torch.Tensor): 21 | tensor_dict[k] = tensor_dict[k].to(device) 22 | 23 | 24 | def to_numpy(tensor): 25 | """Wrapper around .detach().cpu().numpy() """ 26 | if isinstance(tensor, torch.Tensor): 27 | return tensor.detach().cpu().numpy() 28 | elif isinstance(tensor, np.ndarray): 29 | return tensor 30 | else: 31 | raise NotImplementedError 32 | 33 | 34 | class CheckPointManager(object): 35 | """Manager for saving/managing pytorch checkpoints. 36 | 37 | Provides functionality similar to tf.Saver such as 38 | max_to_keep and keep_checkpoint_every_n_hours 39 | """ 40 | def __init__(self, save_path: str = None, max_to_keep=5, keep_checkpoint_every_n_hours=10000.0): 41 | 42 | if max_to_keep <= 0: 43 | raise ValueError('max_to_keep must be at least 1') 44 | 45 | self._max_to_keep = max_to_keep 46 | self._keep_checkpoint_every_n_hours = keep_checkpoint_every_n_hours 47 | 48 | self._ckpt_dir = os.path.dirname(save_path) 49 | self._save_path = save_path + '-{}.pth' if save_path is not None else None 50 | self._logger = logging.getLogger(self.__class__.__name__) 51 | self._checkpoints_fname = os.path.join(self._ckpt_dir, 'checkpoints.txt') 52 | 53 | self._checkpoints_permanent = [] # Will not be deleted 54 | self._checkpoints_buffer = [] # Those which might still be deleted 55 | self._next_save_time = time.time() 56 | self._best_score = -float('inf') 57 | self._best_step = None 58 | 59 | os.makedirs(self._ckpt_dir, exist_ok=True) 60 | self._update_checkpoints_file() 61 | 62 | def _save_checkpoint(self, step, model, optimizer, score): 63 | save_name = self._save_path.format(step) 64 | state = {'state_dict': model.state_dict(), 65 | 'optimizer': optimizer.state_dict(), 66 | 'step': step} 67 | torch.save(state, save_name) 68 | self._logger.info('Saved checkpoint: {}'.format(save_name)) 69 | 70 | self._checkpoints_buffer.append((save_name, time.time())) 71 | 72 | if score > self._best_score: 73 | best_save_name = self._save_path.format('best') 74 | shutil.copyfile(save_name, best_save_name) 75 | self._best_score = score 76 | self._best_step = step 77 | self._logger.info('Checkpoint is current best, score={:.3g}'.format(self._best_score)) 78 | 79 | def _remove_old_checkpoints(self): 80 | while len(self._checkpoints_buffer) > self._max_to_keep: 81 | to_remove = self._checkpoints_buffer.pop(0) 82 | 83 | if to_remove[1] > self._next_save_time: 84 | self._checkpoints_permanent.append(to_remove) 85 | self._next_save_time = to_remove[1] + self._keep_checkpoint_every_n_hours * 3600 86 | else: 87 | os.remove(to_remove[0]) 88 | 89 | def _update_checkpoints_file(self): 90 | checkpoints = [os.path.basename(c[0]) for c in self._checkpoints_permanent + self._checkpoints_buffer] 91 | with open(self._checkpoints_fname, 'w') as fid: 92 | fid.write('\n'.join(checkpoints)) 93 | fid.write('\nBest step: {}'.format(self._best_step)) 94 | 95 | def save(self, model: torch.nn.Module, optimizer: Optimizer, step: int, score: float = 0.0): 96 | """Save model checkpoint to file 97 | 98 | Args: 99 | model: Torch model 100 | optimizer: Torch optimizer 101 | step (int): Step, model will be saved as model-[step].pth 102 | score (float, optional): To determine which model is the best 103 | """ 104 | if self._save_path is None: 105 | raise AssertionError('Checkpoint manager must be initialized with save path for save().') 106 | 107 | self._save_checkpoint(step, model, optimizer, score) 108 | self._remove_old_checkpoints() 109 | self._update_checkpoints_file() 110 | 111 | def load(self, save_path, model: torch.nn.Module = None, optimizer: Optimizer = None): 112 | """Loads saved model from file 113 | 114 | Args: 115 | save_path: Path to saved model (.pth). If a directory is provided instead, model-best.pth is used 116 | model: Torch model to restore weights to 117 | optimizer: Optimizer 118 | """ 119 | if os.path.isdir(save_path): 120 | save_path = os.path.join(save_path, 'model-best.pth') 121 | 122 | state = torch.load(save_path) 123 | 124 | step = 0 125 | if 'step' in state: 126 | step = state['step'] 127 | 128 | if 'state_dict' in state and model is not None: 129 | model.load_state_dict(state['state_dict']) 130 | 131 | if 'optimizer' in state and optimizer is not None: 132 | optimizer.load_state_dict(state['optimizer']) 133 | 134 | self._logger.info('Loaded models from {}'.format(save_path)) 135 | return step 136 | 137 | 138 | class TorchDebugger(torch.autograd.detect_anomaly): 139 | """Enters debugger when anomaly detected""" 140 | def __enter__(self) -> None: 141 | super().__enter__() 142 | 143 | def __exit__(self, type, value, trace): 144 | super().__exit__() 145 | if isinstance(value, RuntimeError): 146 | traceback.print_tb(trace) 147 | print(value) 148 | if sys.gettrace() is None: 149 | pdb.set_trace() 150 | -------------------------------------------------------------------------------- /Predator_APR/configs/kitti/file_LoKITTI_40.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/configs/kitti/file_LoKITTI_40.npy -------------------------------------------------------------------------------- /Predator_APR/configs/kitti/file_LoKITTI_50.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/configs/kitti/file_LoKITTI_50.npy -------------------------------------------------------------------------------- /Predator_APR/configs/kitti/test_kitti.txt: -------------------------------------------------------------------------------- 1 | 8 2 | 9 3 | 10 -------------------------------------------------------------------------------- /Predator_APR/configs/kitti/train_kitti.txt: -------------------------------------------------------------------------------- 1 | 0 2 | 1 3 | 2 4 | 3 5 | 4 6 | 5 -------------------------------------------------------------------------------- /Predator_APR/configs/kitti/val_kitti.txt: -------------------------------------------------------------------------------- 1 | 6 2 | 7 -------------------------------------------------------------------------------- /Predator_APR/configs/models.py: -------------------------------------------------------------------------------- 1 | architectures = dict() 2 | architectures['indoor'] = [ 3 | 'simple', 4 | 'resnetb', 5 | 'resnetb_strided', 6 | 'resnetb', 7 | 'resnetb', 8 | 'resnetb_strided', 9 | 'resnetb', 10 | 'resnetb', 11 | 'resnetb_strided', 12 | 'resnetb', 13 | 'resnetb', 14 | 'nearest_upsample', 15 | 'unary', 16 | 'nearest_upsample', 17 | 'unary', 18 | 'nearest_upsample', 19 | 'last_unary' 20 | ] 21 | 22 | architectures['kitti'] = [ 23 | 'simple', 24 | 'resnetb', 25 | 'resnetb_strided', 26 | 'resnetb', 27 | 'resnetb', 28 | 'resnetb_strided', 29 | 'resnetb', 30 | 'resnetb', 31 | 'resnetb_strided', 32 | 'resnetb', 33 | 'resnetb', 34 | 'nearest_upsample', 35 | 'unary', 36 | 'nearest_upsample', 37 | 'unary', 38 | 'nearest_upsample', 39 | 'last_unary' 40 | ] 41 | 42 | architectures['nuscenes'] = [ 43 | 'simple', 44 | 'resnetb', 45 | 'resnetb_strided', 46 | 'resnetb', 47 | 'resnetb', 48 | 'resnetb_strided', 49 | 'resnetb', 50 | 'resnetb', 51 | 'resnetb_strided', 52 | 'resnetb', 53 | 'resnetb', 54 | 'nearest_upsample', 55 | 'unary', 56 | 'nearest_upsample', 57 | 'unary', 58 | 'nearest_upsample', 59 | 'last_unary' 60 | ] 61 | 62 | architectures['modelnet'] = [ 63 | 'simple', 64 | 'resnetb', 65 | 'resnetb', 66 | 'resnetb_strided', 67 | 'resnetb', 68 | 'resnetb', 69 | 'resnetb_strided', 70 | 'resnetb', 71 | 'resnetb', 72 | 'nearest_upsample', 73 | 'unary', 74 | 'unary', 75 | 'nearest_upsample', 76 | 'unary', 77 | 'last_unary' 78 | ] -------------------------------------------------------------------------------- /Predator_APR/configs/nuscenes/file_LoNUSCENES_40.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/configs/nuscenes/file_LoNUSCENES_40.npy -------------------------------------------------------------------------------- /Predator_APR/configs/nuscenes/file_LoNUSCENES_50.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/configs/nuscenes/file_LoNUSCENES_50.npy -------------------------------------------------------------------------------- /Predator_APR/configs/test/kitti.yaml: -------------------------------------------------------------------------------- 1 | misc: 2 | exp_dir: kitti-rcar-5-20-generalization-test 3 | mode: test 4 | gpu_mode: True 5 | verbose: True 6 | verbose_freq: 500 7 | snapshot_freq: 1 8 | pretrain: snapshot/nuscenes-predator_upper-rcar-5-20_continued/checkpoints/model_best_recall.pth 9 | rot_threshold: 5 10 | trans_threshold: 2 11 | 12 | model: 13 | num_layers: 4 14 | in_points_dim: 3 15 | first_feats_dim: 256 16 | final_feats_dim: 32 17 | first_subsampling_dl: 0.3 18 | in_feats_dim: 1 19 | conv_radius: 4.25 20 | deform_radius: 5.0 21 | num_kernel_points: 15 22 | KP_extent: 2.0 23 | KP_influence: linear 24 | aggregation_mode: sum 25 | fixed_kernel_points: center 26 | use_batch_norm: True 27 | batch_norm_momentum: 0.02 28 | deformable: False 29 | modulated: False 30 | add_cross_score: True 31 | condition_feature: True 32 | model: KPFCNN 33 | generative_model: GenerativeMLP_98 34 | 35 | overlap_attention_module: 36 | gnn_feats_dim: 256 37 | dgcnn_k: 10 38 | num_head: 4 39 | nets: ['self','cross','self'] 40 | 41 | loss: 42 | pos_margin: 0.1 43 | neg_margin: 1.4 44 | log_scale: 48 45 | pos_radius: 0.21 46 | safe_radius: 0.75 47 | overlap_radius: 0.45 48 | matchability_radius: 0.3 49 | w_circle_loss: 1.0 50 | w_overlap_loss: 1.0 51 | w_saliency_loss: 0.0 52 | max_points: 512 53 | loss_ratio: 0.001 54 | inner_loss_ratio: 0.01 55 | regularization_strength: 0.01 56 | inner_regularization_strength: 0.01 57 | 58 | optimiser: 59 | optimizer: SGD 60 | max_epoch: 150 61 | lr: 0.05 62 | weight_decay: 0.000001 63 | momentum: 0.98 64 | scheduler: ExpLR 65 | scheduler_gamma: 0.95 66 | scheduler_freq: 1 67 | iter_size: 1 68 | 69 | dataset: 70 | dataset: kitti 71 | benchmark: odometryKITTI 72 | root: /mnt/disk/KITTIOdometry_Full 73 | batch_size: 1 74 | num_workers: 6 75 | augment_noise: 0.01 76 | augment_shift_range: 2.0 77 | augment_scale_max: 1.2 78 | augment_scale_min: 0.8 79 | pair_min_dist: 5 80 | pair_max_dist: 20 81 | complement_pair_dist: 10 82 | num_complement_one_side: 3 83 | point_generation_ratio: 4 84 | use_old_pose: False 85 | downsample_single: 1 86 | test_augmentation: False 87 | LoKITTI: False 88 | symmetric: False 89 | -------------------------------------------------------------------------------- /Predator_APR/configs/test/nuscenes.yaml: -------------------------------------------------------------------------------- 1 | misc: 2 | exp_dir: nuscenes-rcar-5-20-generalization_test 3 | mode: test 4 | gpu_mode: True 5 | verbose: True 6 | verbose_freq: 500 7 | snapshot_freq: 1 8 | pretrain: snapshot/kitti-rcarIntegrated-rcar-5-20_1e-3/checkpoints/model_best_recall.pth 9 | # pretrain_restart: True 10 | rot_threshold: 5 11 | trans_threshold: 2 12 | 13 | model: 14 | num_layers: 4 15 | in_points_dim: 3 16 | first_feats_dim: 256 17 | final_feats_dim: 32 18 | first_subsampling_dl: 0.3 19 | in_feats_dim: 1 20 | conv_radius: 4.25 21 | deform_radius: 5.0 22 | num_kernel_points: 15 23 | KP_extent: 2.0 24 | KP_influence: linear 25 | aggregation_mode: sum 26 | fixed_kernel_points: center 27 | use_batch_norm: True 28 | batch_norm_momentum: 0.02 29 | deformable: False 30 | modulated: False 31 | add_cross_score: True 32 | condition_feature: True 33 | model: KPFCNN 34 | generative_model: GenerativeMLP_98 35 | 36 | overlap_attention_module: 37 | gnn_feats_dim: 256 38 | dgcnn_k: 10 39 | num_head: 4 40 | nets: ['self','cross','self'] 41 | 42 | loss: 43 | pos_margin: 0.1 44 | neg_margin: 1.4 45 | log_scale: 48 46 | pos_radius: 0.21 47 | safe_radius: 0.75 48 | overlap_radius: 0.45 49 | matchability_radius: 0.3 50 | w_circle_loss: 1.0 51 | w_overlap_loss: 1.0 52 | w_saliency_loss: 0.0 53 | max_points: 512 54 | loss_ratio: 0.0005 55 | inner_loss_ratio: 0.01 56 | regularization_strength: 0.01 57 | inner_regularization_strength: 0.01 58 | 59 | optimiser: 60 | optimizer: SGD 61 | max_epoch: 150 62 | lr: 0.05 63 | weight_decay: 0.000001 64 | momentum: 0.98 65 | scheduler: ExpLR 66 | scheduler_gamma: 0.99 67 | scheduler_freq: 1 68 | iter_size: 1 69 | 70 | dataset: 71 | dataset: nuscenes 72 | benchmark: nucenes 73 | root: /mnt/disk/NUSCENES/nusc_kitti 74 | batch_size: 1 75 | num_workers: 6 76 | augment_noise: 0.01 77 | augment_shift_range: 2.0 78 | augment_scale_max: 1.2 79 | augment_scale_min: 0.8 80 | pair_min_dist: 5 81 | pair_max_dist: 20 82 | complement_pair_dist: 10 83 | num_complement_one_side: 3 84 | point_generation_ratio: 4 85 | downsample_single: 1.0 86 | use_old_pose: True 87 | LoNUSCENES: False 88 | symmetric: False 89 | -------------------------------------------------------------------------------- /Predator_APR/configs/train/kitti.yaml: -------------------------------------------------------------------------------- 1 | misc: 2 | exp_dir: kitti-apr 3 | mode: train 4 | gpu_mode: True 5 | verbose: True 6 | verbose_freq: 150 7 | snapshot_freq: 1 8 | pretrain: "" 9 | pretrain_restart: False 10 | 11 | model: 12 | num_layers: 4 13 | in_points_dim: 3 14 | first_feats_dim: 256 15 | final_feats_dim: 32 16 | first_subsampling_dl: 0.3 17 | in_feats_dim: 1 18 | conv_radius: 4.25 19 | deform_radius: 5.0 20 | num_kernel_points: 15 21 | KP_extent: 2.0 22 | KP_influence: linear 23 | aggregation_mode: sum 24 | fixed_kernel_points: center 25 | use_batch_norm: True 26 | batch_norm_momentum: 0.02 27 | deformable: False 28 | modulated: False 29 | add_cross_score: True 30 | condition_feature: True 31 | model: KPFCNN 32 | generative_model: GenerativeMLP_98 33 | 34 | overlap_attention_module: 35 | gnn_feats_dim: 256 36 | dgcnn_k: 10 37 | num_head: 4 38 | nets: ['self','cross','self'] 39 | 40 | loss: 41 | pos_margin: 0.1 42 | neg_margin: 1.4 43 | log_scale: 48 44 | pos_radius: 0.21 45 | safe_radius: 0.75 46 | overlap_radius: 0.45 47 | matchability_radius: 0.3 48 | w_circle_loss: 1.0 49 | w_overlap_loss: 1.0 50 | w_saliency_loss: 0.0 51 | max_points: 512 52 | loss_ratio: 0.001 53 | regularization_strength: 0.01 54 | 55 | optimiser: 56 | optimizer: SGD 57 | max_epoch: 150 58 | lr: 0.01 59 | weight_decay: 0.000001 60 | momentum: 0.98 61 | scheduler: ExpLR 62 | scheduler_gamma: 0.99 63 | scheduler_freq: 1 64 | iter_size: 1 65 | 66 | dataset: 67 | dataset: kitti 68 | benchmark: odometryKITTI 69 | root: /mnt/disk/KITTIOdometry_Full 70 | batch_size: 1 71 | num_workers: 6 72 | augment_noise: 0.01 73 | augment_shift_range: 2.0 74 | augment_scale_max: 1.2 75 | augment_scale_min: 0.8 76 | pair_min_dist: 5 77 | pair_max_dist: 40 78 | complement_pair_dist: 6 79 | num_complement_one_side: 5 80 | point_generation_ratio: 4 81 | use_old_pose: False 82 | test_augmentation: False 83 | mutate_neighbour_percentage: 0.0 84 | LoKITTI: False 85 | symmetric: False 86 | -------------------------------------------------------------------------------- /Predator_APR/configs/train/nuscenes.yaml: -------------------------------------------------------------------------------- 1 | misc: 2 | exp_dir: nuscenes-get-APC 3 | mode: train 4 | gpu_mode: True 5 | verbose: True 6 | verbose_freq: 150 7 | snapshot_freq: 1 8 | pretrain: snapshot/nuscenes-predator_upper-rcar-5-20_continued/checkpoints/model_best_recall.pth 9 | pretrain_restart: False 10 | 11 | model: 12 | num_layers: 4 13 | in_points_dim: 3 14 | first_feats_dim: 256 15 | final_feats_dim: 32 16 | first_subsampling_dl: 0.3 17 | in_feats_dim: 1 18 | conv_radius: 4.25 19 | deform_radius: 5.0 20 | num_kernel_points: 15 21 | KP_extent: 2.0 22 | KP_influence: linear 23 | aggregation_mode: sum 24 | fixed_kernel_points: center 25 | use_batch_norm: True 26 | batch_norm_momentum: 0.02 27 | deformable: False 28 | modulated: False 29 | add_cross_score: True 30 | condition_feature: True 31 | model: KPFCNN 32 | generative_model: GenerativeMLP_98 33 | 34 | overlap_attention_module: 35 | gnn_feats_dim: 256 36 | dgcnn_k: 10 37 | num_head: 4 38 | nets: ['self','cross','self'] 39 | 40 | loss: 41 | pos_margin: 0.1 42 | neg_margin: 1.4 43 | log_scale: 48 44 | pos_radius: 0.21 45 | safe_radius: 0.75 46 | overlap_radius: 0.45 47 | matchability_radius: 0.3 48 | w_circle_loss: 1.0 49 | w_overlap_loss: 1.0 50 | w_saliency_loss: 0.0 51 | max_points: 512 52 | loss_ratio: 0.0005 53 | inner_loss_ratio: 0.01 54 | regularization_strength: 0.01 55 | inner_regularization_strength: 0.01 56 | 57 | optimiser: 58 | optimizer: SGD 59 | max_epoch: 150 60 | lr: 0.05 61 | weight_decay: 0.000001 62 | momentum: 0.98 63 | scheduler: ExpLR 64 | scheduler_gamma: 0.99 65 | scheduler_freq: 1 66 | iter_size: 1 67 | 68 | dataset: 69 | dataset: nuscenes 70 | benchmark: nucenes 71 | root: /mnt/disk/NUSCENES/nusc_kitti 72 | batch_size: 1 73 | num_workers: 6 74 | augment_noise: 0.01 75 | augment_shift_range: 2.0 76 | augment_scale_max: 1.2 77 | augment_scale_min: 0.8 78 | pair_min_dist: 5 79 | pair_max_dist: 20 80 | complement_pair_dist: 10 81 | num_complement_one_side: 3 82 | point_generation_ratio: 4 83 | mutate_neighbour_percentage: 0 84 | use_old_pose: True 85 | LoNUSCENES: False 86 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | # Compile cpp neighbors 9 | cd cpp_neighbors 10 | python3 setup.py build_ext --inplace 11 | cd .. -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_neighbors/build.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | py setup.py build_ext --inplace 3 | 4 | 5 | pause -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "neighbors.h" 3 | 4 | 5 | void brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose) 6 | { 7 | 8 | // Initialize variables 9 | // ****************** 10 | 11 | // square radius 12 | float r2 = radius * radius; 13 | 14 | // indices 15 | int i0 = 0; 16 | 17 | // Counting vector 18 | int max_count = 0; 19 | vector> tmp(queries.size()); 20 | 21 | // Search neigbors indices 22 | // *********************** 23 | 24 | for (auto& p0 : queries) 25 | { 26 | int i = 0; 27 | for (auto& p : supports) 28 | { 29 | if ((p0 - p).sq_norm() < r2) 30 | { 31 | tmp[i0].push_back(i); 32 | if (tmp[i0].size() > max_count) 33 | max_count = tmp[i0].size(); 34 | } 35 | i++; 36 | } 37 | i0++; 38 | } 39 | 40 | // Reserve the memory 41 | neighbors_indices.resize(queries.size() * max_count); 42 | i0 = 0; 43 | for (auto& inds : tmp) 44 | { 45 | for (int j = 0; j < max_count; j++) 46 | { 47 | if (j < inds.size()) 48 | neighbors_indices[i0 * max_count + j] = inds[j]; 49 | else 50 | neighbors_indices[i0 * max_count + j] = -1; 51 | } 52 | i0++; 53 | } 54 | 55 | return; 56 | } 57 | 58 | void ordered_neighbors(vector& queries, 59 | vector& supports, 60 | vector& neighbors_indices, 61 | float radius) 62 | { 63 | 64 | // Initialize variables 65 | // ****************** 66 | 67 | // square radius 68 | float r2 = radius * radius; 69 | 70 | // indices 71 | int i0 = 0; 72 | 73 | // Counting vector 74 | int max_count = 0; 75 | float d2; 76 | vector> tmp(queries.size()); 77 | vector> dists(queries.size()); 78 | 79 | // Search neigbors indices 80 | // *********************** 81 | 82 | for (auto& p0 : queries) 83 | { 84 | int i = 0; 85 | for (auto& p : supports) 86 | { 87 | d2 = (p0 - p).sq_norm(); 88 | if (d2 < r2) 89 | { 90 | // Find order of the new point 91 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 92 | int index = std::distance(dists[i0].begin(), it); 93 | 94 | // Insert element 95 | dists[i0].insert(it, d2); 96 | tmp[i0].insert(tmp[i0].begin() + index, i); 97 | 98 | // Update max count 99 | if (tmp[i0].size() > max_count) 100 | max_count = tmp[i0].size(); 101 | } 102 | i++; 103 | } 104 | i0++; 105 | } 106 | 107 | // Reserve the memory 108 | neighbors_indices.resize(queries.size() * max_count); 109 | i0 = 0; 110 | for (auto& inds : tmp) 111 | { 112 | for (int j = 0; j < max_count; j++) 113 | { 114 | if (j < inds.size()) 115 | neighbors_indices[i0 * max_count + j] = inds[j]; 116 | else 117 | neighbors_indices[i0 * max_count + j] = -1; 118 | } 119 | i0++; 120 | } 121 | 122 | return; 123 | } 124 | 125 | void batch_ordered_neighbors(vector& queries, 126 | vector& supports, 127 | vector& q_batches, 128 | vector& s_batches, 129 | vector& neighbors_indices, 130 | float radius) 131 | { 132 | 133 | // Initialize variables 134 | // ****************** 135 | 136 | // square radius 137 | float r2 = radius * radius; 138 | 139 | // indices 140 | int i0 = 0; 141 | 142 | // Counting vector 143 | int max_count = 0; 144 | float d2; 145 | vector> tmp(queries.size()); 146 | vector> dists(queries.size()); 147 | 148 | // batch index 149 | int b = 0; 150 | int sum_qb = 0; 151 | int sum_sb = 0; 152 | 153 | 154 | // Search neigbors indices 155 | // *********************** 156 | 157 | for (auto& p0 : queries) 158 | { 159 | // Check if we changed batch 160 | if (i0 == sum_qb + q_batches[b]) 161 | { 162 | sum_qb += q_batches[b]; 163 | sum_sb += s_batches[b]; 164 | b++; 165 | } 166 | 167 | // Loop only over the supports of current batch 168 | vector::iterator p_it; 169 | int i = 0; 170 | for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ ) 171 | { 172 | d2 = (p0 - *p_it).sq_norm(); 173 | if (d2 < r2) 174 | { 175 | // Find order of the new point 176 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 177 | int index = std::distance(dists[i0].begin(), it); 178 | 179 | // Insert element 180 | dists[i0].insert(it, d2); 181 | tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i); 182 | 183 | // Update max count 184 | if (tmp[i0].size() > max_count) 185 | max_count = tmp[i0].size(); 186 | } 187 | i++; 188 | } 189 | i0++; 190 | } 191 | 192 | // Reserve the memory 193 | neighbors_indices.resize(queries.size() * max_count); 194 | i0 = 0; 195 | for (auto& inds : tmp) 196 | { 197 | for (int j = 0; j < max_count; j++) 198 | { 199 | if (j < inds.size()) 200 | neighbors_indices[i0 * max_count + j] = inds[j]; 201 | else 202 | neighbors_indices[i0 * max_count + j] = supports.size(); 203 | } 204 | i0++; 205 | } 206 | 207 | return; 208 | } 209 | 210 | 211 | void batch_nanoflann_neighbors(vector& queries, 212 | vector& supports, 213 | vector& q_batches, 214 | vector& s_batches, 215 | vector& neighbors_indices, 216 | float radius) 217 | { 218 | 219 | // Initialize variables 220 | // ****************** 221 | 222 | // indices 223 | int i0 = 0; 224 | 225 | // Square radius 226 | float r2 = radius * radius; 227 | 228 | // Counting vector 229 | int max_count = 0; 230 | float d2; 231 | vector>> all_inds_dists(queries.size()); 232 | 233 | // batch index 234 | int b = 0; 235 | int sum_qb = 0; 236 | int sum_sb = 0; 237 | 238 | // Nanoflann related variables 239 | // *************************** 240 | 241 | // CLoud variable 242 | PointCloud current_cloud; 243 | 244 | // Tree parameters 245 | nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); 246 | 247 | // KDTree type definition 248 | typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , 249 | PointCloud, 250 | 3 > my_kd_tree_t; 251 | 252 | // Pointer to trees 253 | my_kd_tree_t* index; 254 | 255 | // Build KDTree for the first batch element 256 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 257 | index = new my_kd_tree_t(3, current_cloud, tree_params); 258 | index->buildIndex(); 259 | 260 | 261 | // Search neigbors indices 262 | // *********************** 263 | 264 | // Search params 265 | nanoflann::SearchParams search_params; 266 | search_params.sorted = true; 267 | 268 | for (auto& p0 : queries) 269 | { 270 | 271 | // Check if we changed batch 272 | if (i0 == sum_qb + q_batches[b]) 273 | { 274 | sum_qb += q_batches[b]; 275 | sum_sb += s_batches[b]; 276 | b++; 277 | 278 | // Change the points 279 | current_cloud.pts.clear(); 280 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 281 | 282 | // Build KDTree of the current element of the batch 283 | delete index; 284 | index = new my_kd_tree_t(3, current_cloud, tree_params); 285 | index->buildIndex(); 286 | } 287 | 288 | // Initial guess of neighbors size 289 | all_inds_dists[i0].reserve(max_count); 290 | 291 | // Find neighbors 292 | float query_pt[3] = { p0.x, p0.y, p0.z}; 293 | size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params); 294 | 295 | // Update max count 296 | if (nMatches > max_count) 297 | max_count = nMatches; 298 | 299 | // Increment query idx 300 | i0++; 301 | } 302 | 303 | // Reserve the memory 304 | neighbors_indices.resize(queries.size() * max_count); 305 | i0 = 0; 306 | sum_sb = 0; 307 | sum_qb = 0; 308 | b = 0; 309 | for (auto& inds_dists : all_inds_dists) 310 | { 311 | // Check if we changed batch 312 | if (i0 == sum_qb + q_batches[b]) 313 | { 314 | sum_qb += q_batches[b]; 315 | sum_sb += s_batches[b]; 316 | b++; 317 | } 318 | 319 | for (int j = 0; j < max_count; j++) 320 | { 321 | if (j < inds_dists.size()) 322 | neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb; 323 | else 324 | neighbors_indices[i0 * max_count + j] = supports.size(); 325 | } 326 | i0++; 327 | } 328 | 329 | delete index; 330 | 331 | return; 332 | } 333 | 334 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_neighbors/neighbors/neighbors.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | #include "../../cpp_utils/nanoflann/nanoflann.hpp" 5 | 6 | #include 7 | #include 8 | 9 | using namespace std; 10 | 11 | 12 | void ordered_neighbors(vector& queries, 13 | vector& supports, 14 | vector& neighbors_indices, 15 | float radius); 16 | 17 | void batch_ordered_neighbors(vector& queries, 18 | vector& supports, 19 | vector& q_batches, 20 | vector& s_batches, 21 | vector& neighbors_indices, 22 | float radius); 23 | 24 | void batch_nanoflann_neighbors(vector& queries, 25 | vector& supports, 26 | vector& q_batches, 27 | vector& s_batches, 28 | vector& neighbors_indices, 29 | float radius); 30 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 11 | "neighbors/neighbors.cpp", 12 | "wrapper.cpp"] 13 | 14 | module = Extension(name="radius_neighbors", 15 | sources=SOURCES, 16 | extra_compile_args=['-std=c++11', 17 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 18 | 19 | 20 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_neighbors/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "neighbors/neighbors.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides two methods to compute radius neighbors from pointclouds or batch of pointclouds"; 12 | 13 | static char batch_query_docstring[] = "Method to get radius neighbors in a batch of stacked pointclouds"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *batch_neighbors(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "batch_query", (PyCFunction)batch_neighbors, METH_VARARGS | METH_KEYWORDS, batch_query_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "radius_neighbors", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_radius_neighbors(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Definition of the batch_subsample method 56 | // ********************************** 57 | 58 | static PyObject* batch_neighbors(PyObject* self, PyObject* args, PyObject* keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject* queries_obj = NULL; 66 | PyObject* supports_obj = NULL; 67 | PyObject* q_batches_obj = NULL; 68 | PyObject* s_batches_obj = NULL; 69 | 70 | // Keywords containers 71 | static char* kwlist[] = { "queries", "supports", "q_batches", "s_batches", "radius", NULL }; 72 | float radius = 0.1; 73 | 74 | // Parse the input 75 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "OOOO|$f", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &radius)) 76 | { 77 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 78 | return NULL; 79 | } 80 | 81 | 82 | // Interpret the input objects as numpy arrays. 83 | PyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY); 84 | PyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY); 85 | PyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY); 86 | PyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY); 87 | 88 | // Verify data was load correctly. 89 | if (queries_array == NULL) 90 | { 91 | Py_XDECREF(queries_array); 92 | Py_XDECREF(supports_array); 93 | Py_XDECREF(q_batches_array); 94 | Py_XDECREF(s_batches_array); 95 | PyErr_SetString(PyExc_RuntimeError, "Error converting query points to numpy arrays of type float32"); 96 | return NULL; 97 | } 98 | if (supports_array == NULL) 99 | { 100 | Py_XDECREF(queries_array); 101 | Py_XDECREF(supports_array); 102 | Py_XDECREF(q_batches_array); 103 | Py_XDECREF(s_batches_array); 104 | PyErr_SetString(PyExc_RuntimeError, "Error converting support points to numpy arrays of type float32"); 105 | return NULL; 106 | } 107 | if (q_batches_array == NULL) 108 | { 109 | Py_XDECREF(queries_array); 110 | Py_XDECREF(supports_array); 111 | Py_XDECREF(q_batches_array); 112 | Py_XDECREF(s_batches_array); 113 | PyErr_SetString(PyExc_RuntimeError, "Error converting query batches to numpy arrays of type int32"); 114 | return NULL; 115 | } 116 | if (s_batches_array == NULL) 117 | { 118 | Py_XDECREF(queries_array); 119 | Py_XDECREF(supports_array); 120 | Py_XDECREF(q_batches_array); 121 | Py_XDECREF(s_batches_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting support batches to numpy arrays of type int32"); 123 | return NULL; 124 | } 125 | 126 | // Check that the input array respect the dims 127 | if ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3) 128 | { 129 | Py_XDECREF(queries_array); 130 | Py_XDECREF(supports_array); 131 | Py_XDECREF(q_batches_array); 132 | Py_XDECREF(s_batches_array); 133 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : query.shape is not (N, 3)"); 134 | return NULL; 135 | } 136 | if ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3) 137 | { 138 | Py_XDECREF(queries_array); 139 | Py_XDECREF(supports_array); 140 | Py_XDECREF(q_batches_array); 141 | Py_XDECREF(s_batches_array); 142 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : support.shape is not (N, 3)"); 143 | return NULL; 144 | } 145 | if ((int)PyArray_NDIM(q_batches_array) > 1) 146 | { 147 | Py_XDECREF(queries_array); 148 | Py_XDECREF(supports_array); 149 | Py_XDECREF(q_batches_array); 150 | Py_XDECREF(s_batches_array); 151 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : queries_batches.shape is not (B,) "); 152 | return NULL; 153 | } 154 | if ((int)PyArray_NDIM(s_batches_array) > 1) 155 | { 156 | Py_XDECREF(queries_array); 157 | Py_XDECREF(supports_array); 158 | Py_XDECREF(q_batches_array); 159 | Py_XDECREF(s_batches_array); 160 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : supports_batches.shape is not (B,) "); 161 | return NULL; 162 | } 163 | if ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0)) 164 | { 165 | Py_XDECREF(queries_array); 166 | Py_XDECREF(supports_array); 167 | Py_XDECREF(q_batches_array); 168 | Py_XDECREF(s_batches_array); 169 | PyErr_SetString(PyExc_RuntimeError, "Wrong number of batch elements: different for queries and supports "); 170 | return NULL; 171 | } 172 | 173 | // Number of points 174 | int Nq = (int)PyArray_DIM(queries_array, 0); 175 | int Ns= (int)PyArray_DIM(supports_array, 0); 176 | 177 | // Number of batches 178 | int Nb = (int)PyArray_DIM(q_batches_array, 0); 179 | 180 | // Call the C++ function 181 | // ********************* 182 | 183 | // Convert PyArray to Cloud C++ class 184 | vector queries; 185 | vector supports; 186 | vector q_batches; 187 | vector s_batches; 188 | queries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq); 189 | supports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns); 190 | q_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb); 191 | s_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb); 192 | 193 | // Create result containers 194 | vector neighbors_indices; 195 | 196 | // Compute results 197 | //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 198 | batch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 199 | 200 | // Check result 201 | if (neighbors_indices.size() < 1) 202 | { 203 | PyErr_SetString(PyExc_RuntimeError, "Error"); 204 | return NULL; 205 | } 206 | 207 | // Manage outputs 208 | // ************** 209 | 210 | // Maximal number of neighbors 211 | int max_neighbors = neighbors_indices.size() / Nq; 212 | 213 | // Dimension of output containers 214 | npy_intp* neighbors_dims = new npy_intp[2]; 215 | neighbors_dims[0] = Nq; 216 | neighbors_dims[1] = max_neighbors; 217 | 218 | // Create output array 219 | PyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_INT); 220 | PyObject* ret = NULL; 221 | 222 | // Fill output array with values 223 | size_t size_in_bytes = Nq * max_neighbors * sizeof(int); 224 | memcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes); 225 | 226 | // Merge results 227 | ret = Py_BuildValue("N", res_obj); 228 | 229 | // Clean up 230 | // ******** 231 | 232 | Py_XDECREF(queries_array); 233 | Py_XDECREF(supports_array); 234 | Py_XDECREF(q_batches_array); 235 | Py_XDECREF(s_batches_array); 236 | 237 | return ret; 238 | } 239 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_subsampling/build.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | py setup.py build_ext --inplace 3 | 4 | 5 | pause -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initialize variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initialize variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 101 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 102 | } 103 | } 104 | 105 | return; 106 | } 107 | 108 | 109 | void batch_grid_subsampling(vector& original_points, 110 | vector& subsampled_points, 111 | vector& original_features, 112 | vector& subsampled_features, 113 | vector& original_classes, 114 | vector& subsampled_classes, 115 | vector& original_batches, 116 | vector& subsampled_batches, 117 | float sampleDl, 118 | int max_p) 119 | { 120 | // Initialize variables 121 | // ****************** 122 | 123 | int b = 0; 124 | int sum_b = 0; 125 | 126 | // Number of points in the cloud 127 | size_t N = original_points.size(); 128 | 129 | // Dimension of the features 130 | size_t fdim = original_features.size() / N; 131 | size_t ldim = original_classes.size() / N; 132 | 133 | // Handle max_p = 0 134 | if (max_p < 1) 135 | max_p = N; 136 | 137 | // Loop over batches 138 | // ***************** 139 | 140 | for (b = 0; b < original_batches.size(); b++) 141 | { 142 | 143 | // Extract batch points features and labels 144 | vector b_o_points = vector(original_points.begin () + sum_b, 145 | original_points.begin () + sum_b + original_batches[b]); 146 | 147 | vector b_o_features; 148 | if (original_features.size() > 0) 149 | { 150 | b_o_features = vector(original_features.begin () + sum_b * fdim, 151 | original_features.begin () + (sum_b + original_batches[b]) * fdim); 152 | } 153 | 154 | vector b_o_classes; 155 | if (original_classes.size() > 0) 156 | { 157 | b_o_classes = vector(original_classes.begin () + sum_b * ldim, 158 | original_classes.begin () + sum_b + original_batches[b] * ldim); 159 | } 160 | 161 | 162 | // Create result containers 163 | vector b_s_points; 164 | vector b_s_features; 165 | vector b_s_classes; 166 | 167 | // Compute subsampling on current batch 168 | grid_subsampling(b_o_points, 169 | b_s_points, 170 | b_o_features, 171 | b_s_features, 172 | b_o_classes, 173 | b_s_classes, 174 | sampleDl, 175 | 0); 176 | 177 | // Stack batches points features and labels 178 | // **************************************** 179 | 180 | // If too many points remove some 181 | if (b_s_points.size() <= max_p) 182 | { 183 | subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.end()); 184 | 185 | if (original_features.size() > 0) 186 | subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.end()); 187 | 188 | if (original_classes.size() > 0) 189 | subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.end()); 190 | 191 | subsampled_batches.push_back(b_s_points.size()); 192 | } 193 | else 194 | { 195 | subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.begin() + max_p); 196 | 197 | if (original_features.size() > 0) 198 | subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.begin() + max_p * fdim); 199 | 200 | if (original_classes.size() > 0) 201 | subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.begin() + max_p * ldim); 202 | 203 | subsampled_batches.push_back(max_p); 204 | } 205 | 206 | // Stack new batch lengths 207 | sum_b += original_batches[b]; 208 | } 209 | 210 | return; 211 | } 212 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | void grid_subsampling(vector& original_points, 83 | vector& subsampled_points, 84 | vector& original_features, 85 | vector& subsampled_features, 86 | vector& original_classes, 87 | vector& subsampled_classes, 88 | float sampleDl, 89 | int verbose); 90 | 91 | void batch_grid_subsampling(vector& original_points, 92 | vector& subsampled_points, 93 | vector& original_features, 94 | vector& subsampled_features, 95 | vector& original_classes, 96 | vector& subsampled_classes, 97 | vector& original_batches, 98 | vector& subsampled_batches, 99 | float sampleDl, 100 | int max_p); 101 | 102 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 11 | "grid_subsampling/grid_subsampling.cpp", 12 | "wrapper.cpp"] 13 | 14 | module = Extension(name="grid_subsampling", 15 | sources=SOURCES, 16 | extra_compile_args=['-std=c++11', 17 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 18 | 19 | 20 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initialize limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initialize limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /Predator_APR/cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | struct PointCloud 151 | { 152 | 153 | std::vector pts; 154 | 155 | // Must return the number of data points 156 | inline size_t kdtree_get_point_count() const { return pts.size(); } 157 | 158 | // Returns the dim'th component of the idx'th point in the class: 159 | // Since this is inlined and the "dim" argument is typically an immediate value, the 160 | // "if/else's" are actually solved at compile time. 161 | inline float kdtree_get_pt(const size_t idx, const size_t dim) const 162 | { 163 | if (dim == 0) return pts[idx].x; 164 | else if (dim == 1) return pts[idx].y; 165 | else return pts[idx].z; 166 | } 167 | 168 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 169 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 170 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 171 | template 172 | bool kdtree_get_bbox(BBOX& /* bb */) const { return false; } 173 | 174 | }; 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | -------------------------------------------------------------------------------- /Predator_APR/datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/datasets/__init__.py -------------------------------------------------------------------------------- /Predator_APR/datasets/indoor.py: -------------------------------------------------------------------------------- 1 | """ 2 | Author: Shengyu Huang 3 | Last modified: 30.11.2020 4 | """ 5 | 6 | import os,sys,glob,torch 7 | import numpy as np 8 | from scipy.spatial.transform import Rotation 9 | from torch.utils.data import Dataset 10 | import open3d as o3d 11 | from lib.benchmark_utils import to_o3d_pcd, to_tsfm, get_correspondences 12 | 13 | 14 | class IndoorDataset(Dataset): 15 | """ 16 | Load subsampled coordinates, relative rotation and translation 17 | Output(torch.Tensor): 18 | src_pcd: [N,3] 19 | tgt_pcd: [M,3] 20 | rot: [3,3] 21 | trans: [3,1] 22 | """ 23 | def __init__(self,infos,config,data_augmentation=True): 24 | super(IndoorDataset,self).__init__() 25 | self.infos = infos 26 | self.base_dir = config.root 27 | self.overlap_radius = config.overlap_radius 28 | self.data_augmentation=data_augmentation 29 | self.config = config 30 | 31 | self.rot_factor=1. 32 | self.augment_noise = config.augment_noise 33 | self.max_points = 30000 34 | 35 | def __len__(self): 36 | return len(self.infos['rot']) 37 | 38 | def __getitem__(self,item): 39 | # get transformation 40 | rot=self.infos['rot'][item] 41 | trans=self.infos['trans'][item] 42 | 43 | # get pointcloud 44 | src_path=os.path.join(self.base_dir,self.infos['src'][item]) 45 | tgt_path=os.path.join(self.base_dir,self.infos['tgt'][item]) 46 | src_pcd = torch.load(src_path) 47 | tgt_pcd = torch.load(tgt_path) 48 | 49 | # if we get too many points, we do some downsampling 50 | if(src_pcd.shape[0] > self.max_points): 51 | idx = np.random.permutation(src_pcd.shape[0])[:self.max_points] 52 | src_pcd = src_pcd[idx] 53 | if(tgt_pcd.shape[0] > self.max_points): 54 | idx = np.random.permutation(tgt_pcd.shape[0])[:self.max_points] 55 | tgt_pcd = tgt_pcd[idx] 56 | 57 | # add gaussian noise 58 | if self.data_augmentation: 59 | # rotate the point cloud 60 | euler_ab=np.random.rand(3)*np.pi*2/self.rot_factor # anglez, angley, anglex 61 | rot_ab= Rotation.from_euler('zyx', euler_ab).as_matrix() 62 | if(np.random.rand(1)[0]>0.5): 63 | src_pcd=np.matmul(rot_ab,src_pcd.T).T 64 | rot=np.matmul(rot,rot_ab.T) 65 | else: 66 | tgt_pcd=np.matmul(rot_ab,tgt_pcd.T).T 67 | rot=np.matmul(rot_ab,rot) 68 | trans=np.matmul(rot_ab,trans) 69 | 70 | src_pcd += (np.random.rand(src_pcd.shape[0],3) - 0.5) * self.augment_noise 71 | tgt_pcd += (np.random.rand(tgt_pcd.shape[0],3) - 0.5) * self.augment_noise 72 | 73 | if(trans.ndim==1): 74 | trans=trans[:,None] 75 | 76 | # get correspondence at fine level 77 | tsfm = to_tsfm(rot, trans) 78 | correspondences = get_correspondences(to_o3d_pcd(src_pcd), to_o3d_pcd(tgt_pcd), tsfm,self.overlap_radius) 79 | 80 | src_feats=np.ones_like(src_pcd[:,:1]).astype(np.float32) 81 | tgt_feats=np.ones_like(tgt_pcd[:,:1]).astype(np.float32) 82 | rot = rot.astype(np.float32) 83 | trans = trans.astype(np.float32) 84 | 85 | return src_pcd,tgt_pcd,src_feats,tgt_feats,rot,trans, correspondences, src_pcd, tgt_pcd, torch.ones(1) -------------------------------------------------------------------------------- /Predator_APR/kernels/dispositions/k_015_center_3D.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/kernels/dispositions/k_015_center_3D.ply -------------------------------------------------------------------------------- /Predator_APR/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/lib/__init__.py -------------------------------------------------------------------------------- /Predator_APR/lib/timer.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | class AverageMeter(object): 5 | """Computes and stores the average and current value""" 6 | 7 | def __init__(self): 8 | self.reset() 9 | 10 | def reset(self): 11 | self.val = 0 12 | self.avg = 0 13 | self.sum = 0.0 14 | self.sq_sum = 0.0 15 | self.count = 0 16 | 17 | def update(self, val, n=1): 18 | self.val = val 19 | self.sum += val * n 20 | self.count += n 21 | self.avg = self.sum / self.count 22 | self.sq_sum += val ** 2 * n 23 | self.var = self.sq_sum / self.count - self.avg ** 2 24 | 25 | 26 | class Timer(object): 27 | """A simple timer.""" 28 | 29 | def __init__(self): 30 | self.total_time = 0. 31 | self.calls = 0 32 | self.start_time = 0. 33 | self.diff = 0. 34 | self.avg = 0. 35 | 36 | def reset(self): 37 | self.total_time = 0 38 | self.calls = 0 39 | self.start_time = 0 40 | self.diff = 0 41 | self.avg = 0 42 | 43 | def tic(self): 44 | # using time.time instead of time.clock because time time.clock 45 | # does not normalize for multithreading 46 | self.start_time = time.time() 47 | 48 | def toc(self, average=True): 49 | self.diff = time.time() - self.start_time 50 | self.total_time += self.diff 51 | self.calls += 1 52 | self.avg = self.total_time / self.calls 53 | if average: 54 | return self.avg 55 | else: 56 | return self.diff 57 | -------------------------------------------------------------------------------- /Predator_APR/lib/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | General utility functions 3 | 4 | Author: Shengyu Huang 5 | Last modified: 30.11.2020 6 | """ 7 | 8 | import os,re,sys,json,yaml,random, argparse, torch, pickle 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import torch.optim as optim 12 | import numpy as np 13 | from scipy.spatial.transform import Rotation 14 | 15 | from sklearn.neighbors import NearestNeighbors 16 | from scipy.spatial.distance import minkowski 17 | _EPS = 1e-7 # To prevent division by zero 18 | 19 | 20 | class Logger: 21 | def __init__(self, path): 22 | self.path = path 23 | self.fw = open(self.path+'/log','a') 24 | 25 | def write(self, text): 26 | self.fw.write(text) 27 | self.fw.flush() 28 | 29 | def close(self): 30 | self.fw.close() 31 | 32 | def save_obj(obj, path ): 33 | """ 34 | save a dictionary to a pickle file 35 | """ 36 | with open(path, 'wb') as f: 37 | pickle.dump(obj, f) 38 | 39 | def load_obj(path): 40 | """ 41 | read a dictionary from a pickle file 42 | """ 43 | with open(path, 'rb') as f: 44 | return pickle.load(f) 45 | 46 | def load_config(path): 47 | """ 48 | Loads config file: 49 | 50 | Args: 51 | path (str): path to the config file 52 | 53 | Returns: 54 | config (dict): dictionary of the configuration parameters, merge sub_dicts 55 | 56 | """ 57 | with open(path,'r') as f: 58 | cfg = yaml.safe_load(f) 59 | 60 | config = dict() 61 | for key, value in cfg.items(): 62 | for k,v in value.items(): 63 | config[k] = v 64 | 65 | return config 66 | 67 | 68 | def setup_seed(seed): 69 | """ 70 | fix random seed for deterministic training 71 | """ 72 | torch.manual_seed(seed) 73 | torch.cuda.manual_seed_all(seed) 74 | np.random.seed(seed) 75 | random.seed(seed) 76 | torch.backends.cudnn.deterministic = True 77 | 78 | def square_distance(src, dst, normalised = False): 79 | """ 80 | Calculate Euclid distance between each two points. 81 | Args: 82 | src: source points, [B, N, C] 83 | dst: target points, [B, M, C] 84 | Returns: 85 | dist: per-point square distance, [B, N, M] 86 | """ 87 | B, N, _ = src.shape 88 | _, M, _ = dst.shape 89 | dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) 90 | if(normalised): 91 | dist += 2 92 | else: 93 | dist += torch.sum(src ** 2, dim=-1)[:, :, None] 94 | dist += torch.sum(dst ** 2, dim=-1)[:, None, :] 95 | 96 | dist = torch.clamp(dist, min=1e-12, max=None) 97 | return dist 98 | 99 | 100 | def validate_gradient(model): 101 | """ 102 | Confirm all the gradients are non-nan and non-inf 103 | """ 104 | for name, param in model.named_parameters(): 105 | if param.grad is not None: 106 | if torch.any(torch.isnan(param.grad)): 107 | return False 108 | if torch.any(torch.isinf(param.grad)): 109 | return False 110 | return True 111 | 112 | 113 | def natural_key(string_): 114 | """ 115 | Sort strings by numbers in the name 116 | """ 117 | return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_)] -------------------------------------------------------------------------------- /Predator_APR/main.py: -------------------------------------------------------------------------------- 1 | import os, torch, time, shutil, json,glob, argparse, shutil 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 3 | 4 | import numpy as np 5 | from easydict import EasyDict as edict 6 | 7 | from datasets.dataloader import get_dataloader, get_datasets 8 | from models.architectures import KPFCNN, KPFCNNDecoder 9 | from models.mlp import get_GenerativeMLP 10 | from lib.utils import setup_seed, load_config 11 | from lib.tester import get_trainer 12 | from lib.loss import MetricLoss 13 | from configs.models import architectures 14 | 15 | from torch import optim 16 | from torch import nn 17 | setup_seed(0) 18 | 19 | 20 | if __name__ == '__main__': 21 | # indicate the gpu this thread is using 22 | print(f"Using GPU number " + os.environ["CUDA_VISIBLE_DEVICES"]) 23 | # load configs 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('config', type=str, help= 'Path to the config file.') 26 | args = parser.parse_args() 27 | config = load_config(args.config) 28 | config['snapshot_dir'] = 'snapshot/%s' % config['exp_dir'] 29 | config['tboard_dir'] = 'snapshot/%s/tensorboard' % config['exp_dir'] 30 | config['save_dir'] = 'snapshot/%s/checkpoints' % config['exp_dir'] 31 | config = edict(config) 32 | 33 | os.makedirs(config.snapshot_dir, exist_ok=True) 34 | os.makedirs(config.save_dir, exist_ok=True) 35 | os.makedirs(config.tboard_dir, exist_ok=True) 36 | json.dump( 37 | config, 38 | open(os.path.join(config.snapshot_dir, 'config.json'), 'w'), 39 | indent=4, 40 | ) 41 | if config.gpu_mode: 42 | config.device = torch.device('cuda') 43 | else: 44 | config.device = torch.device('cpu') 45 | 46 | # backup the files 47 | os.system(f'cp -r models {config.snapshot_dir}') 48 | os.system(f'cp -r datasets {config.snapshot_dir}') 49 | os.system(f'cp -r lib {config.snapshot_dir}') 50 | shutil.copy2('main.py',config.snapshot_dir) 51 | 52 | # model initialization 53 | config.switch_to_decoder = False 54 | config.architecture = architectures[config.dataset] 55 | config.model = KPFCNN(config) 56 | config.model_name = "KPFCNN" 57 | 58 | # build symmetrical or asymmetrical decoders 59 | config.switch_to_decoder = True 60 | if config.symmetric: 61 | config.generative_model = KPFCNNDecoder(config) 62 | else: 63 | config.generative_model = get_GenerativeMLP(config) 64 | 65 | # create optimizer 66 | if config.optimizer == 'SGD': 67 | config.optimizer = optim.SGD( 68 | params = [ 69 | {'params': config.model.parameters()}, 70 | {'params': config.generative_model.parameters()} 71 | ], 72 | lr=config.lr, 73 | momentum=config.momentum, 74 | weight_decay=config.weight_decay, 75 | ) 76 | elif config.optimizer == 'ADAM': 77 | config.optimizer = optim.Adam( 78 | params = [ 79 | {'params': config.model.parameters()}, 80 | {'params': config.generator_model.parameters()} 81 | ], 82 | lr=config.lr, 83 | betas=(0.9, 0.999), 84 | weight_decay=config.weight_decay, 85 | ) 86 | 87 | # create learning rate scheduler 88 | config.scheduler = optim.lr_scheduler.ExponentialLR( 89 | config.optimizer, 90 | gamma=config.scheduler_gamma, 91 | ) 92 | 93 | # create dataset and dataloader 94 | train_set, val_set, benchmark_set = get_datasets(config) 95 | config.train_loader, neighborhood_limits = get_dataloader(dataset=train_set, 96 | batch_size=config.batch_size, 97 | shuffle=True, 98 | num_workers=config.num_workers, 99 | ) 100 | config.val_loader, _ = get_dataloader(dataset=val_set, 101 | batch_size=config.batch_size, 102 | shuffle=False, 103 | num_workers=1, 104 | neighborhood_limits=neighborhood_limits 105 | ) 106 | 107 | config.test_loader, _ = get_dataloader(dataset=benchmark_set, 108 | batch_size=config.batch_size, 109 | shuffle=False, 110 | num_workers=1, 111 | neighborhood_limits=neighborhood_limits) 112 | 113 | # create evaluation metrics 114 | config.desc_loss = MetricLoss(config) 115 | trainer = get_trainer(config) 116 | if(config.mode=='train'): 117 | trainer.train() 118 | elif(config.mode =='val'): 119 | trainer.eval() 120 | else: 121 | trainer.test() -------------------------------------------------------------------------------- /Predator_APR/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/Predator_APR/models/__init__.py -------------------------------------------------------------------------------- /Predator_APR/models/gcn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | import torch.nn as nn 4 | from copy import deepcopy 5 | import torch.utils.checkpoint as checkpoint 6 | from lib.utils import square_distance 7 | 8 | 9 | def get_graph_feature(coords, feats, k=10): 10 | """ 11 | Apply KNN search based on coordinates, then concatenate the features to the centroid features 12 | Input: 13 | X: [B, 3, N] 14 | feats: [B, C, N] 15 | Return: 16 | feats_cat: [B, 2C, N, k] 17 | """ 18 | # apply KNN search to build neighborhood 19 | B, C, N = feats.size() 20 | dist = square_distance(coords.transpose(1,2), coords.transpose(1,2)) 21 | 22 | idx = dist.topk(k=k+1, dim=-1, largest=False, sorted=True)[1] #[B, N, K+1], here we ignore the smallest element as it's the query itself 23 | idx = idx[:,:,1:] #[B, N, K] 24 | 25 | idx = idx.unsqueeze(1).repeat(1,C,1,1) #[B, C, N, K] 26 | all_feats = feats.unsqueeze(2).repeat(1, 1, N, 1) # [B, C, N, N] 27 | 28 | neighbor_feats = torch.gather(all_feats, dim=-1,index=idx) #[B, C, N, K] 29 | 30 | # concatenate the features with centroid 31 | feats = feats.unsqueeze(-1).repeat(1,1,1,k) 32 | 33 | feats_cat = torch.cat((feats, neighbor_feats-feats),dim=1) 34 | 35 | return feats_cat 36 | 37 | 38 | 39 | class SelfAttention(nn.Module): 40 | def __init__(self,feature_dim,k=10): 41 | super(SelfAttention, self).__init__() 42 | self.conv1 = nn.Conv2d(feature_dim*2, feature_dim, kernel_size=1, bias=False) 43 | self.in1 = nn.InstanceNorm2d(feature_dim) 44 | 45 | self.conv2 = nn.Conv2d(feature_dim*2, feature_dim * 2, kernel_size=1, bias=False) 46 | self.in2 = nn.InstanceNorm2d(feature_dim * 2) 47 | 48 | self.conv3 = nn.Conv2d(feature_dim * 4, feature_dim, kernel_size=1, bias=False) 49 | self.in3 = nn.InstanceNorm2d(feature_dim) 50 | 51 | self.k = k 52 | 53 | def forward(self, coords, features): 54 | """ 55 | Here we take coordinats and features, feature aggregation are guided by coordinates 56 | Input: 57 | coords: [B, 3, N] 58 | feats: [B, C, N] 59 | Output: 60 | feats: [B, C, N] 61 | """ 62 | B, C, N = features.size() 63 | 64 | x0 = features.unsqueeze(-1) #[B, C, N, 1] 65 | 66 | x1 = get_graph_feature(coords, x0.squeeze(-1), self.k) 67 | x1 = F.leaky_relu(self.in1(self.conv1(x1)), negative_slope=0.2) 68 | x1 = x1.max(dim=-1,keepdim=True)[0] 69 | 70 | x2 = get_graph_feature(coords, x1.squeeze(-1), self.k) 71 | x2 = F.leaky_relu(self.in2(self.conv2(x2)), negative_slope=0.2) 72 | x2 = x2.max(dim=-1, keepdim=True)[0] 73 | 74 | x3 = torch.cat((x0,x1,x2),dim=1) 75 | x3 = F.leaky_relu(self.in3(self.conv3(x3)), negative_slope=0.2).view(B, -1, N) 76 | 77 | return x3 78 | 79 | 80 | def MLP(channels: list, do_bn=True): 81 | """ Multi-layer perceptron """ 82 | n = len(channels) 83 | layers = [] 84 | for i in range(1, n): 85 | layers.append( 86 | nn.Conv1d(channels[i - 1], channels[i], kernel_size=1, bias=True)) 87 | if i < (n-1): 88 | if do_bn: 89 | layers.append(nn.InstanceNorm1d(channels[i])) 90 | layers.append(nn.ReLU()) 91 | return nn.Sequential(*layers) 92 | 93 | 94 | def attention(query, key, value): 95 | dim = query.shape[1] 96 | scores = torch.einsum('bdhn,bdhm->bhnm', query, key) / dim**.5 97 | prob = torch.nn.functional.softmax(scores, dim=-1) 98 | return torch.einsum('bhnm,bdhm->bdhn', prob, value), prob 99 | 100 | 101 | class MultiHeadedAttention(nn.Module): 102 | """ Multi-head attention to increase model expressivitiy """ 103 | def __init__(self, num_heads: int, d_model: int): 104 | super().__init__() 105 | assert d_model % num_heads == 0 106 | self.dim = d_model // num_heads 107 | self.num_heads = num_heads 108 | self.merge = nn.Conv1d(d_model, d_model, kernel_size=1) 109 | self.proj = nn.ModuleList([deepcopy(self.merge) for _ in range(3)]) 110 | 111 | def forward(self, query, key, value): 112 | batch_dim = query.size(0) 113 | query, key, value = [l(x).view(batch_dim, self.dim, self.num_heads, -1) 114 | for l, x in zip(self.proj, (query, key, value))] 115 | x, _ = attention(query, key, value) 116 | return self.merge(x.contiguous().view(batch_dim, self.dim*self.num_heads, -1)) 117 | 118 | 119 | class AttentionalPropagation(nn.Module): 120 | def __init__(self, feature_dim: int, num_heads: int): 121 | super().__init__() 122 | self.attn = MultiHeadedAttention(num_heads, feature_dim) 123 | self.mlp = MLP([feature_dim*2, feature_dim*2, feature_dim]) 124 | nn.init.constant_(self.mlp[-1].bias, 0.0) 125 | 126 | def forward(self, x, source): 127 | message = self.attn(x, source, source) 128 | return self.mlp(torch.cat([x, message], dim=1)) 129 | 130 | 131 | class MultiHeadedAttentionCat(nn.Module): 132 | """ Multi-head attention to increase model expressivitiy """ 133 | def __init__(self, num_heads: int, d_model: int): 134 | super().__init__() 135 | assert d_model % num_heads == 0 136 | self.dim = d_model // num_heads 137 | self.num_heads = num_heads 138 | self.distribute = nn.Conv1d(d_model, d_model, kernel_size=1) 139 | self.proj = nn.ModuleList([self.distribute, deepcopy(self.distribute), deepcopy(self.distribute)]) 140 | self.merge = nn.Conv1d(d_model+7*4, d_model+7*4, kernel_size=1) 141 | 142 | def forward(self, query, key, value, coords0, coords1): 143 | batch_dim = query.size(0) 144 | num_q_points = coords0.size(2) 145 | query, key, value = [l(x).view(batch_dim, self.dim, self.num_heads, -1) 146 | for l, x in zip(self.proj, (query, key, value))] 147 | # print(f"value shape: {value.shape}; coords shape: {coords1.shape}", flush=True) 148 | value = torch.cat([value, coords1.repeat(1, 4, 1).view(1, 4, 3, -1).transpose(1,2)], dim=1) 149 | x, _ = attention(query, key, value) 150 | augment1 = x[:, self.dim:self.dim+3, :, :] - coords0.repeat(1, 4, 1).view(1, 4, 3, -1).transpose(1,2) 151 | augment2 = torch.norm(augment1, dim=1, keepdim=True) 152 | y = torch.zeros((batch_dim, self.dim+7, self.num_heads, num_q_points), dtype=augment2.dtype).to(query.device) 153 | for i in range(batch_dim): 154 | y[i] = torch.cat([x[i], augment1[i], augment2[i]], dim=0) 155 | del x 156 | return self.merge(y.contiguous().view(batch_dim, (self.dim+7)*self.num_heads, -1)) 157 | 158 | 159 | class AttentionalPropagationCat(nn.Module): 160 | def __init__(self, feature_dim: int, num_heads: int): 161 | super().__init__() 162 | self.attn = MultiHeadedAttentionCat(num_heads, feature_dim) 163 | self.mlp = MLP([feature_dim*2+7*4, feature_dim*2, feature_dim]) 164 | nn.init.constant_(self.mlp[-1].bias, 0.0) 165 | 166 | def forward(self, feats0, feats1, coords0, coords1): 167 | message = self.attn(feats0, feats1, feats1, coords0, coords1) 168 | return self.mlp(torch.cat([feats0, message], dim=1)) 169 | 170 | 171 | class GCN(nn.Module): 172 | """ 173 | Alternate between self-attention and cross-attention 174 | Input: 175 | coords: [B, 3, N] 176 | feats: [B, C, N] 177 | Output: 178 | feats: [B, C, N] 179 | """ 180 | def __init__(self, num_head: int, feature_dim: int, k: int, layer_names: list): 181 | super().__init__() 182 | self.layers=[] 183 | for atten_type in layer_names: 184 | if atten_type == 'cross': 185 | self.layers.append(AttentionalPropagation(feature_dim,num_head)) 186 | elif atten_type == 'cross_cat': 187 | self.layers.append(AttentionalPropagationCat(feature_dim,num_head)) 188 | elif atten_type == 'self': 189 | self.layers.append(SelfAttention(feature_dim, k)) 190 | self.layers = nn.ModuleList(self.layers) 191 | self.names = layer_names 192 | 193 | def forward(self, coords0, coords1, desc0, desc1): 194 | for layer, name in zip(self.layers, self.names): 195 | if name == 'cross': 196 | # desc0 = desc0 + checkpoint.checkpoint(layer, desc0, desc1) 197 | # desc1 = desc1 + checkpoint.checkpoint(layer, desc1, desc0) 198 | desc0 = desc0 + layer(desc0, desc1) 199 | desc1 = desc1 + layer(desc1, desc0) 200 | elif name == 'cross_cat': 201 | desc0 = desc0 + layer(desc0, desc1, coords0, coords1) 202 | desc1 = desc1 + layer(desc1, desc0, coords1, coords0) 203 | elif name == 'self': 204 | desc0 = layer(coords0, desc0) 205 | desc1 = layer(coords1, desc1) 206 | return desc0, desc1 -------------------------------------------------------------------------------- /Predator_APR/models/mlp.py: -------------------------------------------------------------------------------- 1 | # import torch.nn as nn 2 | # # import MinkowskiEngine as ME 3 | # # import MinkowskiEngine.MinkowskiFunctional as MEF 4 | # # from model.common import get_norm 5 | 6 | # class GenerativeMLP(nn.Module): 7 | # CHANNELS = [None, 512, 128, None] 8 | 9 | # def __init__(self, 10 | # in_channel=125, 11 | # out_points=6, 12 | # bn_momentum=0.1): 13 | # super().__init__() 14 | # CHANNELS = self.CHANNELS 15 | # self.mlp = nn.Sequential( 16 | # nn.Linear(in_channel, CHANNELS[1]), 17 | # nn.ReLU(), 18 | # nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 19 | # nn.Linear(CHANNELS[1], CHANNELS[2]), 20 | # nn.ReLU(), 21 | # nn.BatchNorm1d(CHANNELS[2], momentum=bn_momentum), 22 | # nn.Linear(CHANNELS[2], out_points*3), 23 | # nn.ReLU() 24 | # ) 25 | 26 | # def forward(self, x): 27 | # y = self.mlp(x) 28 | # # print(y) 29 | # return y 30 | 31 | 32 | # class GenerativeMLP_99(GenerativeMLP): 33 | # CHANNELS = [None, 512, 512, None] 34 | 35 | 36 | # class GenerativeMLP_98(GenerativeMLP): 37 | # CHANNELS = [None, 512, 256, None] 38 | 39 | 40 | # class GenerativeMLP_54(GenerativeMLP): 41 | # CHANNELS = [None, 32, 16, None] 42 | 43 | 44 | # class GenerativeMLP_4(nn.Module): 45 | # CHANNELS = [None, 16, None] 46 | 47 | # def __init__(self, 48 | # in_channel=125, 49 | # out_points=6, 50 | # bn_momentum=0.1): 51 | # super().__init__() 52 | # CHANNELS = self.CHANNELS 53 | # self.mlp = nn.Sequential( 54 | # nn.Linear(in_channel, CHANNELS[1]), 55 | # nn.ReLU(), 56 | # nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 57 | # nn.Linear(CHANNELS[1], out_points*3), 58 | # nn.ReLU() 59 | # ) 60 | 61 | # def forward(self, x): 62 | # y = self.mlp(x) 63 | # # print(y) 64 | # return y 65 | 66 | 67 | # class GenerativeMLP_11_10_9(nn.Module): 68 | # CHANNELS = [None, 2048, 1024, 512, None] 69 | 70 | # def __init__(self, 71 | # in_channel=125, 72 | # out_points=6, 73 | # bn_momentum=0.1): 74 | # super().__init__() 75 | # CHANNELS = self.CHANNELS 76 | # self.mlp = nn.Sequential( 77 | # nn.Linear(in_channel, CHANNELS[1]), 78 | # nn.ReLU(), 79 | # nn.BatchNorm1d(CHANNELS[1], momentum=bn_momentum), 80 | # nn.Linear(CHANNELS[1], CHANNELS[2]), 81 | # nn.ReLU(), 82 | # nn.BatchNorm1d(CHANNELS[2], momentum=bn_momentum), 83 | # nn.Linear(CHANNELS[2], CHANNELS[3]), 84 | # nn.ReLU(), 85 | # nn.BatchNorm1d(CHANNELS[3], momentum=bn_momentum), 86 | # nn.Linear(CHANNELS[3], out_points*3), 87 | # nn.ReLU() 88 | # ) 89 | 90 | # def forward(self, x): 91 | # y = self.mlp(x) 92 | # # print(y) 93 | # return y 94 | 95 | 96 | # def get_GenerativeMLP(config): 97 | # models = [GenerativeMLP_4, GenerativeMLP_98, GenerativeMLP_99, GenerativeMLP_54, GenerativeMLP_11_10_9] 98 | # mdict = {model.__name__: model for model in models} 99 | # return mdict[config.generative_model](in_channel=config.final_feats_dim, 100 | # out_points=config.point_generation_ratio, 101 | # bn_momentum=config.batch_norm_momentum) 102 | 103 | import torch.nn as nn 104 | # import MinkowskiEngine as ME 105 | # import MinkowskiEngine.MinkowskiFunctional as MEF 106 | # from model.common import get_norm 107 | 108 | class GenerativeMLP(nn.Module): 109 | CHANNELS = [None, 512, 128, None] 110 | 111 | def __init__(self, 112 | in_channel=125, 113 | out_points=6, 114 | radius = 1, 115 | bn_momentum=0.1): 116 | super().__init__() 117 | # print(in_channel) 118 | self.CHANNELS[0] = in_channel 119 | self.CHANNELS[-1] = out_points*3 120 | self.list_modules = [] 121 | self.radius = radius 122 | for layer_idx in range(len(self.CHANNELS)-1): 123 | # print(layer_idx) 124 | if layer_idx < len(self.CHANNELS)-1: # middle layer 125 | self.list_modules.append( 126 | nn.Sequential( 127 | nn.Linear(self.CHANNELS[layer_idx], self.CHANNELS[layer_idx+1]), 128 | nn.ReLU(), 129 | nn.BatchNorm1d(self.CHANNELS[layer_idx+1], momentum=bn_momentum), 130 | ) 131 | ) 132 | else: # last layer 133 | self.list_modules.append( 134 | nn.Sequential( 135 | nn.Linear(self.CHANNELS[layer_idx], self.CHANNELS[layer_idx+1]), 136 | nn.ReLU(), 137 | ) 138 | ) 139 | self.list_modules = nn.ModuleList(self.list_modules) 140 | 141 | 142 | def forward(self, x): 143 | # print(x.size(), self.CHANNELS) 144 | for module in self.list_modules: 145 | x = module(x) 146 | if self.radius is None: 147 | return x 148 | else: 149 | return x, self.radius 150 | 151 | 152 | class GenerativeMLP_99(GenerativeMLP): 153 | CHANNELS = [None, 512, 512, None] 154 | 155 | 156 | class GenerativeMLP_98(GenerativeMLP): 157 | CHANNELS = [None, 512, 256, None] 158 | 159 | 160 | class GenerativeMLP_54(GenerativeMLP): 161 | CHANNELS = [None, 32, 16, None] 162 | 163 | 164 | class GenerativeMLP_4(GenerativeMLP): 165 | CHANNELS = [None, 16, None] 166 | 167 | 168 | class GenerativeMLP_11_10_9(GenerativeMLP): 169 | CHANNELS = [None, 2048, 1024, 512, None] 170 | 171 | 172 | def get_GenerativeMLP(config, radius=None, in_channels=None): 173 | models = [GenerativeMLP_4, GenerativeMLP_98, GenerativeMLP_99, GenerativeMLP_54, GenerativeMLP_11_10_9] 174 | mdict = {model.__name__: model for model in models} 175 | if in_channels is None: 176 | in_channels = config.final_feats_dim 177 | return mdict[config.generative_model](in_channel=in_channels, 178 | out_points=config.point_generation_ratio, 179 | radius=radius, 180 | bn_momentum=config.batch_norm_momentum) -------------------------------------------------------------------------------- /Predator_APR/requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib==3.3.3 2 | numpy==1.19.4 3 | nibabel==3.2.1 4 | tqdm==4.38.0 5 | open3d==0.10.0.0 6 | easydict==1.9 7 | scipy==1.5.4 8 | coloredlogs==15.0 9 | PyYAML==5.4.1 10 | scikit_learn==0.24.1 11 | tensorboardX==2.1 12 | vtk_visualizer==0.9.6 13 | nibabel==3.2.1 14 | h5py==3.2.1 15 | coloredlogs==15.0 16 | gitpython==3.1.41 17 | chamferdist==1.0.0 -------------------------------------------------------------------------------- /Predator_APR/scripts/cal_overlap.py: -------------------------------------------------------------------------------- 1 | """ 2 | We use this script to calculate the overlap ratios for all the train/test fragment pairs 3 | """ 4 | import os,sys,glob 5 | import open3d as o3d 6 | from lib.utils import natural_key 7 | import numpy as np 8 | from tqdm import tqdm 9 | import multiprocessing as mp 10 | 11 | def determine_epsilon(): 12 | """ 13 | We follow Learning Compact Geomtric Features to compute this hyperparameter, which unfortunately we didn't use later. 14 | """ 15 | base_dir='../dataset/3DMatch/test/*/03_Transformed/*.ply' 16 | files=sorted(glob.glob(base_dir),key=natural_key) 17 | etas=[] 18 | for eachfile in files: 19 | pcd=o3d.io.read_point_cloud(eachfile) 20 | pcd=pcd.voxel_down_sample(0.025) 21 | pcd_tree = o3d.geometry.KDTreeFlann(pcd) 22 | distances=[] 23 | for i, point in enumerate(pcd.points): 24 | [count,vec1, vec2] = pcd_tree.search_knn_vector_3d(point,2) 25 | distances.append(np.sqrt(vec2[1])) 26 | etai=np.median(distances) 27 | etas.append(etai) 28 | return np.median(etas) 29 | 30 | 31 | def get_overlap_ratio(source,target,threshold=0.03): 32 | """ 33 | We compute overlap ratio from source point cloud to target point cloud 34 | """ 35 | pcd_tree = o3d.geometry.KDTreeFlann(target) 36 | 37 | match_count=0 38 | for i, point in enumerate(source.points): 39 | [count, _, _] = pcd_tree.search_radius_vector_3d(point, threshold) 40 | if(count!=0): 41 | match_count+=1 42 | 43 | overlap_ratio = match_count / len(source.points) 44 | return overlap_ratio 45 | 46 | def cal_overlap_per_scene(c_folder): 47 | base_dir=os.path.join(c_folder,'03_Transformed') 48 | fragments=sorted(glob.glob(base_dir+'/*.ply'),key=natural_key) 49 | n_fragments=len(fragments) 50 | 51 | with open(f'{c_folder}/overlaps_ours.txt','w') as f: 52 | for i in tqdm(range(n_fragments-1)): 53 | for j in range(i+1,n_fragments): 54 | path1,path2=fragments[i],fragments[j] 55 | 56 | # load, downsample and transform 57 | pcd1=o3d.io.read_point_cloud(path1) 58 | pcd2=o3d.io.read_point_cloud(path2) 59 | pcd1=pcd1.voxel_down_sample(0.01) 60 | pcd2=pcd2.voxel_down_sample(0.01) 61 | 62 | # calculate overlap 63 | c_overlap = get_overlap_ratio(pcd1,pcd2) 64 | f.write(f'{i},{j},{c_overlap:.4f}\n') 65 | f.close() 66 | 67 | if __name__=='__main__': 68 | base_dir='your data folder' 69 | scenes = sorted(glob.glob(base_dir)) 70 | 71 | p = mp.Pool(processes=mp.cpu_count()) 72 | p.map(cal_overlap_mat,scenes) 73 | p.close() 74 | p.join() -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # APR: Online Distant Point Cloud Registration Through Aggregated Point Cloud Reconstruction (IJCAI'23) 2 | 3 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/apr-online-distant-point-cloud-registration/point-cloud-registration-on-kitti-distant-pcr)](https://paperswithcode.com/sota/point-cloud-registration-on-kitti-distant-pcr?p=apr-online-distant-point-cloud-registration) 4 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/apr-online-distant-point-cloud-registration/point-cloud-registration-on-nuscenes-distant)](https://paperswithcode.com/sota/point-cloud-registration-on-nuscenes-distant?p=apr-online-distant-point-cloud-registration) 5 | 6 | For many driving safety applications, it is of great importance to accurately register LiDAR point clouds generated on distant moving vehicles. However, such point clouds have extremely different point density and sensor perspective on the same object, making registration on such point clouds very hard. In this paper, we propose a novel feature extraction framework, called APR, for online distant point cloud registration. Specifically, APR leverages an autoencoder design, where the autoencoder reconstructs a denser aggregated point cloud with several frames instead of the original single input point cloud. Our design forces the encoder to extract features with rich local geometry information based on one single input point cloud. Such features are then used for online distant point cloud registration. We conduct extensive experiments against state-of-the-art (SOTA) feature extractors on KITTI and nuScenes datasets. Results show that APR outperforms all other extractors by a large margin, increasing average registration recall of SOTA extractors by 7.1% on LoKITTI and 4.6% on LoNuScenes. 7 | 8 | Camera-ready:[https://www.ijcai.org/proceedings/2023/134](https://www.ijcai.org/proceedings/2023/134) 9 | Arxiv: [https://arxiv.org/abs/2305.02893](https://arxiv.org/abs/2305.02893) 10 | 11 | ## News 12 | 13 | 20230720 - Our new advancement, GCL (ICCV 2023), achieved up to 40% RR increase on KITTI long range scenarios! Be sure to also [check it out](https://www.github.com/liuQuan98/GCL)! 14 | 15 | 20230508 - Source code and pretrained weights are released. 16 | 17 | 20230420 - Our paper has been accepted by IJCAI'23! 18 | 19 | ## Aggregated Point Cloud Reconstruction (APR) 20 | 21 | ![1683269100777](resources/1683269100777.png) 22 | Traditional feature extractors dedicated to registration take in point clouds and output point-wise features. We propose to extend the workflow to treat the original feature extractor as the encoder (Key-frame Feature Extractor, KFE) of an autoencoder structure, while the decoder (Non-key-frame Point-cloud Reconstruction, NPR) decodes an accumulated denser environmental point cloud, i.e., the Aggregated Point Cloud. It is obtained by accumulating the past and future frames around a key frame (in Aggregated Point-cloud Generation, APG). The encoder is encouraged to enhance its feature representativeness in order to allow the decoder to guess the dense from the sparse, and infer the complete from the incomplete. Please refer to the paper for a detailed description. 23 | 24 |
25 | 26 |
27 | 28 | This is a typical Aggregated Point Cloud (in blue) versus the corresponding key frame (in orange). 29 | 30 | ## Important notes before you proceed 31 | 32 | This repository is a combination of two separate implementation of APR, one upon FCGF and another upon Predator. **During the execution of both methods, please treat `./FCGF_APR` or `./Predator_APR` as the working directory.** 33 | 34 | **For method-specific execution instructions, please refer to the README.md files in the child directories, including [FCGF_APR](./FCGF_APR/README.md) and [Predator_APR](./Predator_APR/README.md).** 35 | 36 | ## Requirements 37 | 38 | - Ubuntu 14.04 or higher 39 | - CUDA 11.1 or higher 40 | - Python v3.7 or higher 41 | - Pytorch v1.6 or higher 42 | - [MinkowskiEngine](https://github.com/stanfordvl/MinkowskiEngine) v0.5 or higher 43 | 44 | ## Dataset Preparation 45 | 46 | ### KITTI 47 | 48 | For KITTI dataset preparation, please first follow the [KITTI official instructions](https://www.cvlibs.net/datasets/kitti/eval_odometry.php) to download the 'velodyne laser data', 'calibration files', and (optionally) 'ground truth poses'. 49 | 50 | Since the GT poses provided in KITTI drift a lot, we recommend using the pose labels provided by [SemanticKITTI](http://www.semantic-kitti.org/dataset.html#download) instead, as they are more accurate. Please follow the official instruction to download the split called 'SemanticKITTI label data'. 51 | 52 | Extract all compressed files in the same folder and we are done. We denote KITTI_ROOT as the directory that have the following structure: `{$KITTI_ROOT}/dataset/poses` and `{$KITTI_ROOT}/dataset/sequences/XX`. 53 | 54 | The option to use KITTI original pose is still preserved which can be enabled by setting `use_old_pose` to True in the scripts for FCGF_APR (or config files for Predator_APR), although we highly discourage doing so due to performance degredation. Please note that all of the methods reported in our paper are retrained on the label of SemanticKITTI instead of OdometryKITTI. 55 | 56 | ### nuScenes 57 | 58 | The vanilla nuScenes dataset structure is not friendly to the registration task, so we propose to convert the lidar part into KITTI format for ease of development and extension. Thanks to the code provided by nuscenes-devkit, the conversion requires only minimal modification. 59 | 60 | To download nuScenes, please follow the [nuscenes official page](https://www.nuscenes.org/nuscenes#download) to obtain the 'lidar blobs' (inside 'file blobs') and 'Metadata' of the 'trainval' and 'test' split in the 'Full dataset (v1.0)' section. Only LiDAR scans and pose annotations are used in APR. 61 | 62 | Next, execute the following commands to deploy [nuscenes-devkit](https://github.com/nutonomy/nuscenes-devkit) and our conversion script: 63 | 64 | ``` 65 | git clone https://github.com/nutonomy/nuscenes-devkit.git 66 | conda create -n nuscenes-devkit python=3.8 67 | conda activate nuscenes-devkit 68 | pip install nuscenes-devkit 69 | cp ./resources/export_kitti_minimal.py ./nuscenes-devkit/python-sdk/nuscenes/scripts/export_kitti_minimal.py 70 | ``` 71 | 72 | Cater the `nusc_dir` and `nusc_kitti_dir` parameter in `./nuscenes-devkit/python-sdk/nuscenes/scripts/export_kitti_minimal.py` line 51 & 52 to your preferred path. Parameter `nusc_dir` specifies the path to the nuScenes dataset, and `nusc_kitti_dir` will be the path to store the converted nuScenes LiDAR data. Start conversion by executing the following instructions: 73 | 74 | ``` 75 | cd ./nuscenes-devkit/python-sdk 76 | python nuscenes/scripts/export_kitti_minimal.py 77 | ``` 78 | 79 | The process may be slow (can take hours). 80 | 81 | ## Installation 82 | 83 | We recommend conda for installation. First, we need to create a basic environment to setup MinkowskiEngine: 84 | 85 | ``` 86 | conda create -n apr python=3.7 pip=21.1 87 | conda activate apr 88 | conda install pytorch=1.9.0 torchvision cudatoolkit=11.1 -c pytorch -c nvidia 89 | pip install numpy 90 | ``` 91 | 92 | Then install [Minkowski Engine](https://github.com/NVIDIA/MinkowskiEngine) along with other dependencies: 93 | 94 | ``` 95 | pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps --install-option="--blas_include_dirs=${CONDA_PREFIX}/include" --install-option="--blas=openblas" 96 | pip install -r ./FCGF_APR/requirements.txt 97 | pip install -r ./Predator_APR/requirements.txt 98 | ``` 99 | 100 | ## Train and test 101 | 102 | For method-specific instructions and pretrained models, please refer to the README.md files in the child directories, [FCGF_APR](./FCGF_APR/README.md) and [Predator_APR](./Predator_APR/README.md). 103 | 104 | ## Acknowlegdements 105 | 106 | We thank [FCGF](https://github.com/chrischoy/FCGF) and [Predator](https://github.com/prs-eth/OverlapPredator) for the wonderful baseline methods, [Chamferdist](https://github.com/krrish94/chamferdist) for providing the torch-compatible implementation of chamfer loss, and [nuscenes-devkit](https://github.com/nutonomy/nuscenes-devkit) for the convenient dataset conversion code. 107 | -------------------------------------------------------------------------------- /resources/1683269100777.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/resources/1683269100777.png -------------------------------------------------------------------------------- /resources/1683269374754.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liuQuan98/APR/8effc071b1a6942f7902e2f74613f58dcc8b15ff/resources/1683269374754.png -------------------------------------------------------------------------------- /resources/export_kitti_minimal.py: -------------------------------------------------------------------------------- 1 | # nuScenes dev-kit. 2 | # Code written by Holger Caesar, 2019. 3 | 4 | """ 5 | This script converts nuScenes data to KITTI format and KITTI results to nuScenes. 6 | It is used for compatibility with software that uses KITTI-style annotations. 7 | 8 | We do not encourage this, as: 9 | - KITTI has only front-facing cameras, whereas nuScenes has a 360 degree horizontal fov. 10 | - KITTI has no radar data. 11 | - The nuScenes database format is more modular. 12 | - KITTI fields like occluded and truncated cannot be exactly reproduced from nuScenes data. 13 | - KITTI has different categories. 14 | 15 | Limitations: 16 | - We don't specify the KITTI imu_to_velo_kitti projection in this code base. 17 | - We map nuScenes categories to nuScenes detection categories, rather than KITTI categories. 18 | - Attributes are not part of KITTI and therefore set to '' in the nuScenes result format. 19 | - Velocities are not part of KITTI and therefore set to 0 in the nuScenes result format. 20 | - This script uses the `train` and `val` splits of nuScenes, whereas standard KITTI has `training` and `testing` splits. 21 | 22 | This script includes one main function: 23 | - nuscenes_construct_kitti_PCR_data(): Converts nuScenes LiDAR data and pose annotation to KITTI format. 24 | 25 | To launch these scripts run: 26 | - python export_kitti_minimal.py 27 | 28 | To work with the original KITTI dataset, use these parameters: 29 | --nusc_kitti_dir /data/sets/kitti --split training 30 | 31 | See https://www.nuscenes.org/object-detection for more information on the nuScenes result format. 32 | """ 33 | import json 34 | import os 35 | from typing import List, Dict, Any 36 | 37 | import matplotlib.pyplot as plt 38 | import numpy as np 39 | from PIL import Image 40 | from pyquaternion import Quaternion 41 | 42 | from nuscenes.nuscenes import NuScenes 43 | from nuscenes.utils.data_classes import LidarPointCloud 44 | from nuscenes.utils.geometry_utils import transform_matrix 45 | from nuscenes.utils.splits import create_splits_logs 46 | 47 | 48 | class KittiConverter: 49 | def __init__(self, 50 | nusc_dir: str = '/mnt/disk/NUSCENES', 51 | nusc_kitti_dir: str = '/mnt/disk/NUSCENES/nusc_kitti', 52 | lidar_name: str = 'LIDAR_TOP', 53 | nusc_version: str = 'v1.0-trainval', 54 | split: str = 'val'): 55 | """ 56 | :param nusc_dir: Root of nuScenes dataset. 57 | :param nusc_kitti_dir: Where to write the KITTI-style annotations. 58 | :param lidar_name: Name of the lidar sensor. 59 | :param image_count: Number of images to convert. 60 | :param nusc_version: nuScenes version to use. 61 | :param split: Dataset split to use. 62 | """ 63 | self.nusc_kitti_dir = os.path.expanduser(nusc_kitti_dir) 64 | self.lidar_name = lidar_name 65 | self.nusc_version = nusc_version 66 | self.split = split 67 | 68 | # Create nusc_kitti_dir. 69 | if not os.path.isdir(self.nusc_kitti_dir): 70 | os.makedirs(self.nusc_kitti_dir) 71 | 72 | # Select subset of the data to look at. 73 | self.nusc = NuScenes(version=nusc_version, dataroot=nusc_dir, verbose=True) 74 | 75 | def nuscenes_construct_kitti_PCR_data(self) -> None: 76 | """ 77 | Converts nuScenes Lidar sequences and poses into KITTI form 78 | """ 79 | kitti_to_nu_lidar = Quaternion(axis=(0, 0, 1), angle=np.pi / 2) 80 | kitti_to_nu_lidar_inv = kitti_to_nu_lidar.inverse 81 | 82 | # Get assignment of scenes to splits. 83 | split_logs = create_splits_logs(self.split, self.nusc) 84 | # print(split_logs[0]) 85 | 86 | # Create output folder. 87 | base_folder = os.path.join(self.nusc_kitti_dir, self.split, 'sequences') 88 | # indice_folder = os.path.join(self.nusc_kitti_dir, self.split, 'indice') 89 | for folder in [base_folder]: 90 | if not os.path.isdir(folder): 91 | os.makedirs(folder) 92 | 93 | for log in split_logs: 94 | # Use only the samples from the current split. 95 | sample_tokens = self._split_to_samples(log) 96 | 97 | token_idx = 0 # Start tokens from 0. 98 | trans = [] 99 | 100 | log_folder = os.path.join(base_folder, log, 'velodyne') 101 | if not os.path.isdir(log_folder): 102 | os.makedirs(log_folder) 103 | 104 | for sample_token in sample_tokens: 105 | print(f"Processing {log}, {token_idx}") 106 | # Get sample data. 107 | sample = self.nusc.get('sample', sample_token) 108 | lidar_token = sample['data'][self.lidar_name] 109 | 110 | # Retrieve sensor records. 111 | sd_record_lid = self.nusc.get('sample_data', lidar_token) 112 | cs_record_lid = self.nusc.get('calibrated_sensor', sd_record_lid['calibrated_sensor_token']) 113 | 114 | # Get ego pose. Note that ego pose is the position of imu, not that of Lidar, thus it needs correcting. 115 | pos = self.nusc.get('ego_pose', sd_record_lid['ego_pose_token']) 116 | ego_to_world = transform_matrix(pos['translation'], Quaternion(pos['rotation']), 117 | inverse=False) 118 | lid_to_ego = transform_matrix(cs_record_lid['translation'], Quaternion(cs_record_lid['rotation']), 119 | inverse=False) 120 | lid_to_world = np.dot(ego_to_world, lid_to_ego) 121 | lid_to_world_kitti = np.dot(lid_to_world, kitti_to_nu_lidar.transformation_matrix) 122 | trans.append(lid_to_world_kitti) 123 | 124 | # Retrieve the token from the lidar. 125 | # Note that this may be confusing as the filename of the camera will include the timestamp of the lidar, 126 | # not the camera. 127 | filename_lid_full = sd_record_lid['filename'] 128 | token = '%06d' % token_idx # We use KITTI names instead of nuScenes names 129 | token_idx += 1 130 | 131 | # Convert lidar. 132 | # Note that we are only using a single sweep, instead of the commonly used n sweeps. 133 | src_lid_path = os.path.join(self.nusc.dataroot, filename_lid_full) 134 | dst_lid_path = os.path.join(log_folder, token + '.bin') 135 | assert not dst_lid_path.endswith('.pcd.bin') 136 | pcl = LidarPointCloud.from_file(src_lid_path) 137 | pcl.rotate(kitti_to_nu_lidar_inv.rotation_matrix) # In KITTI lidar frame. 138 | with open(dst_lid_path, "w") as lid_file: 139 | pcl.points.T.tofile(lid_file) 140 | 141 | # Save poses of a single log sequence into one file 142 | trans = np.array(trans) 143 | pose_path = os.path.join(base_folder, log, 'poses') 144 | np.save(pose_path, trans) 145 | 146 | 147 | if __name__ == '__main__': 148 | for convert_split in ['train', 'val', 'test']: 149 | converter = KittiConverter(split=convert_split) 150 | converter.nuscenes_construct_kitti_PCR_data() 151 | --------------------------------------------------------------------------------