Awesome
Neighborhood Normalization for Robust Geometric Feature Learning
<img src="https://user-images.githubusercontent.com/10382126/154413903-2bae50d3-9f17-4570-b301-7e57efc446b8.png" width="640">This codebase implements the method described in the paper:
Neighborhood Normalization for Robust Geometric Feature Learning
Xingtong Liu*, Benjamin D. Killeen*, Ayushi Sinha, Masaru Ishii, Gregory D. Hager, Russell H. Taylor, Mathias Unberath
In 2021 Conference on Computer Vision and Pattern Recognition (CVPR)
Please contact Xingtong Liu (xingtongliu@jhu.edu) or Mathias Unberath (unberath@jhu.edu) if you have any questions.
Dependencies
MinkowskiEngine(0.5.4 tested), PyTorch(1.7.1 tested), open3d(0.9.0 tested), pyvista, pyacvd, tensorboardX, MulticoreTSNE, umap-learn, torch-geometric, pycuda, scikit-image, scikit-learn, opencv-python, scipy, psutil, tqdm, pathlib, numpy
Datasets
3DMatch
- Download original dataset using this script.
- Download information of pairs with scene overlap here. Note the train/val/test split provided here is the one used in this work.
- Generate mesh dataset using
scripts/generate_3dmatch_mesh.py
. One example is/path/to/python /path/to/generate_3dmatch_mesh.py --data_root /path/to/3dmatch/original/dataset --output_root /path/to/3dmatch/mesh/output
- Split the generated mesh dataset in step 3 in the same way as the pair information in step 2. Store these two types of data in the same folder. Because temporary data will be generated during network training, at least 100 GB storage space is required for this dataset.
- For evaluation on the standard benchmark, download the test point cloud dataset here.
- For evaluation on the resolution mismatch benchmark, use the mesh data in the test split generated in step 3. Before
evaluation, copy the
*-evaluation
folders downloaded in step 5 to the test split generated in step 3.
KITTI
- Download velodyne dataset from this link. 85 GB storage space is needed.
- train/val/test split has been specified in the dataset class
KITTIPairDataset
indatasets/kitti_dataset.py
and no further operations are needed.
Nasal Cavity
- Download dataset for network training and
evluation here
. Data inside
low_resolution
folder is used for computation efficiency reason. - The 52 head CT scans from several public datasets and the corresponding statistical shape models are provided here . Note this is not needed for this work and we provide it for other potential research topics.
- For experiments in this work, train/val/test split on this dataset simply use different range of PCA mode weights.
Network Training
train_3dmatch.py
is for training a 3D geometric descriptor on the 3DMatch dataset. One example to run this script is/path/to/python /path/to/train_3dmatch.py --config_path /path/to/config/file
. Example config files are provided in thescripts
folder. To train a network for the standard benchmark (fixed resolution), one example config file istrain_3dmatch_standard.json
. To train a network with Batch-Neighborhood Normalization ( B-NHN) for the resolution mismatch benchmark, the example config file istrain_3dmatch_mismatch_bnhn.json
. To train with Neighborhood Normalization (NHN), one example config file istrain_3dmatch_nhn.json
. For the two-stage training strategy mentioned in the paper, a network should first be trained withtrain_3dmatch_nhn.json
until convergence. Then it should be trained withtrain_3dmatch_mismatch_nhn_to_bnhn.json
to adjust the weights of B-NHNs.train_kitti.py
is for training a 3D geometric descriptor on the KITTI dataset. One example to run this script is/path/to/python /path/to/train_kitti.py --config_path /path/to/config/file
. The example config files for the standard and resolution mismatch benchmarks aretrain_kitti_standard.json
andtrain_kitti_mismatch.json
, respectively.train_nasal.py
is for training a 3D geometric descriptor on the Nasal Cavity dataset. One example to run this script is/path/to/python /path/to/train_nasal.py --config_path /path/to/config/file
. The example config file for the resolution mismatch benchmark istrain_nasal_mismatch.json
.
Evaluation
-
evaluation/evaluation_3dmatch.py
is to evaluate Feature Match Recall (FMR) on the 3DMatch dataset. One example to run this script is/path/to/python /path/to/evaluation_3dmatch.py --config_path /path/to/config/file
. Example config files for the standard and resolution mismatch benchmarks areeval_3dmatch_standard.json
amdeval_3dmatch_mismatch.json
, respectively. -
evaluation/evaluation_kitti_standard.py
is to evaluate FMR and Registration Error on the standard benchmark of KITTI dataset. One example to run this script is/path/to/python /path/to/evaluation_kitti_standard.py --config_path /path/to/config/file
. One example config file iseval_kitti_standard.json
. -
evaluation/evaluation_kitti_mismatch.py
is to evaluate FMR on the resolution mismatch benchmark of KITTI dataset. One example to run this script is/path/to/python /path/to/evaluation_kitti_mismatch.py --config_path /path/to/config/file
. One example config file iseval_kitti_mismatch.json
. -
evaluation/evaluation_nasal.py
is to evaluate FMR on the resolution mismatch benchmark of Nasal Cavity dataset. One example to run this script is/path/to/python /path/to/evaluation_nasal.py --config_path /path/to/config/file
. One example config file iseval_nasal_mismatch.json
.
Pre-trained Models
We provide pre-trained models with Batch-Neighborhood Normalization on the datasets used in this work.
net_norm_type
and net_upsample_type
in the config files corresponds to the two terms separated by +
below. These models produced experiment results reported in the paper for the B-NHN normalization type.
Benchmark / Dataset | 3DMatch | KITTI | Nasal Cavity |
---|---|---|---|
Standard | B-NHN + Transpose | B-NHN + Pool | N/A |
Resolution mismatch | B-NHN + Pool | B-NHN + Pool | B-NHN + Pool |
Visualization
Here we present feature embeddings of several sample pairs from the three datasets with resolution mismatch. Features are reduced to 3-dimension for color display using UMAP.