code release

.gitignore

@@ -0,0 +1,7 @@
+snapshot/
+*.so
+*.zip
+data_ThreeDMatch/
+weights/
+__pycache__/
+build/

README.md

@@ -1,27 +1,90 @@
-### PREDATOR: Registration of 3D Point Clouds with Low Overlap
+# PREDATOR: Registration of 3D Point Clouds with Low Overlap
 This repository provides code and data required to train and evaluate PREDATOR, a  model  for  **p**airwise point-cloud **re**gistration with **d**eep **at**tention to the **o**verlap **r**egion. It represents the official implementation of the paper:
 
-#### [PREDATOR: Registration of 3D Point Clouds with Low Overlap](https://addlink)
+### [PREDATOR: Registration of 3D Point Clouds with Low Overlap](https://addlink)
+
 \*[Shengyu Huang](https://shengyuh.github.io)
 , \*[Zan Gojcic](https://zgojcic.github.io/)
-, [Mikhail Usvyatsov](https://prs.igp.ethz.ch/content/specialinterest/baug/institute-igp/photogrammetry-and-remote-sensing/en/group/people/person-detail.html?persid=242711)
+, [Mikhail Usvyatsov](https://aelphy.github.io)
 , [Andreas Wieser](https://gseg.igp.ethz.ch/people/group-head/prof-dr--andreas-wieser.html)
-, [Konrad Schiendler](https://prs.igp.ethz.ch/group/people/person-detail.schindler.html)\
+, [Konrad Schindler](https://prs.igp.ethz.ch/group/people/person-detail.schindler.html)\
 |[ETH Zurich](https://igp.ethz.ch/) |\
 \* Equal contribution
 
 
 
 ![Predator_teaser](figures/teaser_predator.jpg?raw=true)
 
-**Code coming soon, stay tuned!**
 
 ### Citation
 
 If you find this code useful for your work or use it in your project, please consider citing:
 
-
 ```shell
 add bibtex
 ```
 
+### Contact
+If you have any questions, please let us know: Shengyu Huang {[email protected]}
+
+## News
+- 2020-11-30: Code release
+
+## Instructions
+This code has been tested on 
+- Python 3.6.9, PyTorch 1.4.0, CUDA 11.0, gcc 7.5, TITAN Xp
+- Python 3.7.4, PyTorch 1.5.1+cu101, CUDA 10.1, gcc 6.3.0, GeForce RTX 2080 Ti
+- Python 3.8.5, PyTorch 1.8.0.dev20201124+cu110, CUDA 11.1, gcc 9.3.0, GeForce RTX 3090
+
+**Note**: We observe data loader random crash due to memory issues on machines with less than 64GB CPU RAM.
+
+### Requirements
+Under your working folder, our virtual environment and requirements can be installed by running:
+```shell
+git clone https://github.com/ShengyuH/OverlapPredator.git
+virtualenv --no-site-packages predator -p python3; source predator/bin/activate
+cd OverlapPredator; pip install -r requirements.txt
+cd cpp_wrappers; sh compile_wrappers.sh; cd ..
+```
+
+### Datasets
+We provide preprocessed 3DMatch pairwise datasets, you can download them [here](https://drive.google.com/file/d/11oD5YsLn4OBNpLp4d-VEZtHegWpHaa_K/view?usp=sharing)(500MB). Please unzip it and move to ```OverlapPredator```.
+
+### Pretrained weights
+We provide two pretrained models on 3DMatch dataset, Predator and bigPredator. bigPredator is a wider network which is trained on a single GeForce RTX 3090, you can download them [here](https://drive.google.com/file/d/1xLqv1CBiFukRUn7fHYiLTXGuc0q1xo3x/view?usp=sharing)(275MB). Please unzip it and move to ```OverlapPredator```.
+
+| Model       | first_feats_dim   | gnn_feats_dim | # parameters| performance |
+|:-----------:|:-------------------:|:-------:|:-------:|:-------:|
+| Predator | 128               | 256 | 7.43M| recall: 34%|
+| bigPredator | 256                | 512 | 29.67M| recall: 40%|
+
+
+### Train
+After having virtual environment and datasets prepared, you can train from scratch by running:
+```shell
+python main.py --mode train --exp_dir predator_3dmatch --first_feats_dim 128 --gnn_feats_dim 256
+```
+or 
+```shell
+python main.py --mode train --exp_dir bigpredator_3dmatch --first_feats_dim 256 --gnn_feats_dim 512
+```
+
+### Evaluate
+To evaluate PREDATOR, we first extract features and scores and store them as .pth files, then run RANSAC. To extract features and scores on 3DLoMatch benchmark, run: 
+```shell
+python main.py --mode test --exp_dir val_predator --pretrain model_zoo/Predator.pth --first_feats_dim 128 --gnn_feats_dim 256 --test_info 3DLoMatch
+```
+the features will be saved to ```snapshot/{exp_dir}/3DLoMatch```. Then we can run RANSAC by:
+```shell
+python evaluate_predator.py --source_path snapshot/val_predator/3DLoMatch --n_points 1000 --exp_dir est_3dlomatch_1000
+```
+this might take a few minutes, depends on ```n_points``` used by RANSAC. 
+
+### 
+### Demo
+TODO
+
+
+
+
+

config.py

@@ -0,0 +1,97 @@
+import argparse
+import time
+import os
+
+arg_lists = []
+parser = argparse.ArgumentParser()
+
+
+def add_argument_group(name):
+  arg = parser.add_argument_group(name)
+  arg_lists.append(arg)
+  return arg
+
+
+def str2bool(v):
+  return v.lower() in ('true', '1')
+
+
+# snapshot configurations
+snapshot_arg = add_argument_group('Snapshot')
+snapshot_arg.add_argument('--exp_dir', type=str, default='test')
+snapshot_arg.add_argument('--snapshot_interval',type=int, default=1)
+
+# KPConv Network configuration
+net_arg = add_argument_group('Network')
+net_arg.add_argument('--num_layers', type=int, default=4)
+net_arg.add_argument('--in_points_dim', type=int, default=3)
+net_arg.add_argument('--first_feats_dim', type=int, default=256) 
+net_arg.add_argument('--final_feats_dim', type=int, default=32)
+net_arg.add_argument('--first_subsampling_dl', type=float, default=0.025)
+net_arg.add_argument('--in_features_dim', type=int, default=1)
+net_arg.add_argument('--conv_radius', type=float, default=2.5)
+net_arg.add_argument('--deform_radius', type=float, default=5.0)
+net_arg.add_argument('--num_kernel_points', type=int, default=15)
+net_arg.add_argument('--KP_extent', type=float, default=2.0)
+net_arg.add_argument('--KP_influence', type=str, default='linear')
+net_arg.add_argument('--aggregation_mode', type=str, default='sum', choices=['closest', 'sum'])
+net_arg.add_argument('--fixed_kernel_points', type=str, default='center', choices=['center', 'verticals', 'none'])
+net_arg.add_argument('--use_batch_norm', type=str2bool, default=True)
+net_arg.add_argument('--batch_norm_momentum', type=float, default=0.02)
+net_arg.add_argument('--deformable', type=str2bool, default=False)
+net_arg.add_argument('--modulated', type=str2bool, default=False)
+
+# GNN Network configuration
+net_arg.add_argument('--gnn_feats_dim', type = int, default=512,help='feature dimention for DGCNN')
+net_arg.add_argument('--dgcnn_k', type = int, default=10,help='knn graph in DGCNN')
+net_arg.add_argument('--num_head', type = int, default=4,help='cross attention head')
+net_arg.add_argument('--nets', type = str, default="['self','cross','self']",help='GNN configuration')
+
+# Loss configurations
+loss_arg = add_argument_group('Loss')
+loss_arg.add_argument('--pos_margin', type=float, default=0.1)
+loss_arg.add_argument('--neg_margin', type=float, default=1.4)
+loss_arg.add_argument('--log_scale', type=float, default=24)
+loss_arg.add_argument('--pos_radius', type=float,default = 0.025 * 1.5)
+loss_arg.add_argument('--matchability_radius', type=float, default=0.025 * 2)
+loss_arg.add_argument('--safe_radius', type =float, default=0.025 * 4)
+loss_arg.add_argument('--w_circle_loss', type=float, default=1.0)
+loss_arg.add_argument('--w_overlap_loss', type = float, default=1.0)
+loss_arg.add_argument('--w_saliency_loss', type=float, default=0.0)
+loss_arg.add_argument('--max_points', type=int, default=256,help='maximal points for circle loss')
+
+# Optimizer configurations
+opt_arg = add_argument_group('Optimizer')
+opt_arg.add_argument('--optimizer', type=str, default='SGD', choices=['SGD', 'ADAM'])
+opt_arg.add_argument('--max_epoch', type=int, default=40)
+opt_arg.add_argument('--lr', type=float, default=5e-3)
+opt_arg.add_argument('--weight_decay', type=float, default=1e-6)
+opt_arg.add_argument('--momentum', type=float, default=0.98)
+opt_arg.add_argument('--scheduler', type=str, default='ExpLR')
+opt_arg.add_argument('--scheduler_gamma', type=float, default=0.95)
+opt_arg.add_argument('--scheduler_interval', type=int, default=1)
+opt_arg.add_argument('--iter_size', type = int, default=1)
+
+
+# Dataset and dataloader configurations
+data_arg = add_argument_group('Data')
+data_arg.add_argument('--root', type=str, default='data_ThreeDMatch')
+data_arg.add_argument('--augment_noise', type=float, default=0.005)
+data_arg.add_argument('--batch_size', type=int, default=1)
+data_arg.add_argument('--num_workers', type=int, default=6)
+data_arg.add_argument('--train_info', type = str, default='configs/train_info.pkl')
+data_arg.add_argument('--val_info', type = str, default='configs/val_info.pkl')
+data_arg.add_argument('--test_info', type =str, default='3DMatch')
+
+
+# Other configurations
+misc_arg = add_argument_group('Misc')
+misc_arg.add_argument('--mode', type=str, default='test')
+misc_arg.add_argument('--gpu_mode', type=str2bool, default=True)
+misc_arg.add_argument('--verbose', type=str2bool, default=True)
+misc_arg.add_argument('--verbose_freq', type = int, default=1000)
+misc_arg.add_argument('--pretrain', type=str, default=None)
+
+def get_config():
+  args = parser.parse_args()
+  return args