initial commit

.gitattributes

@@ -33,3 +33,35 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/r3pmnet_overview.png filter=lfs diff=lfs merge=lfs -text
+assets/sioux_cranfield.png filter=lfs diff=lfs merge=lfs -text
+assets/sioux_scans.png filter=lfs diff=lfs merge=lfs -text
+assets/success_cases.png filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_classifier/models/best_model_snap.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_classifier/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_classifier/models/best_ptnet_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_curvenet/models/model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_dcp/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_flownet/models/model.best.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_ipcrnet/models/best_model_v1.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_ipcrnet/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_ipcrnet/models/best_ptnet_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_0.01.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_0.6.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_0.7.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_0.8.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_0.9.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_100.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_200.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_300.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_400.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_masknet2/models/best_model_500.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_pcn/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_pnlk/models/best_model_snap.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_pnlk/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_pnlk/models/best_ptnet_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_prnet/models/best_model.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/pretrained/exp_prnet/models/model.99.t7 filter=lfs diff=lfs merge=lfs -text
+thirdparty/learning3d/utils/lib/build/lib.linux-x86_64-3.5/pointnet2_cuda.cpython-35m-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,12 @@
+__pycache__/
+*:*Zone.Identifier
+.ipynb_checkpoints/
+*.ipynb
+
+checkpoints/
+data/
+results/
+registration_plys/
+logs/
+notebooks/
+kernels/

README.md

@@ -0,0 +1,251 @@
+<!-- # R3PM-Net
+
+
+
+This repository contains the official implementation of the paper:
+
+<p align="center">
+  <strong><a href="https://arxiv.org/abs/2604.05060">R3PM-Net: Real-time, Robust, Real-world Point Matching Network</a></strong><br>
+  <strong>(AI4RWC@CVPRW 2026 - Oral Presentation)</strong>
+</p> -->
+
+<p align="center">
+
+  <h1 align="center">R3PM-Net: Real-time, Robust, Real-world Point Matching Network</h1>
+<p align="center"> <strong>AI4RWC@CVPRW 2026 - Oral Presentation</strong></p>
+  <h3 align="center"><a href="https://arxiv.org/abs/2604.05060">Paper</a> | <a href="https://yasiikb.github.io/R3PM-Net/">Project Page</a> | <a href="https://huggingface.co/datasets/YasiiKB/R3PM-Net">Dataset</a></h3>
+  <div align="center"></div>
+</p>
+<p align="center"> <img src="assets/r3pmnet_overview.png" width="95%"> </p>
+<p align="left"><i>Figure 1. Overview of the R3PM-Net Architecture. R3PM-Net employs a global-aware feature extraction module with shared weights to learn geometric similarities across a full receptive field.</i></p>
+
+## Introduction
+
+R3PM-Net is a lightweight, global-aware, object-level point matching network designed to bridge the gap between approaches trained and evaluated on clean, dense, synthetic and real-world industrial point cloud data by prioritizing both generalizability and real-time efficiency.
+
+<p align="center"> <img src="assets/teaser.png" width="40%"> </p>
+<p align="left"><i>Figure 2. Examples of R3PM-Net performance on the Sioux-Cranfield dataset.</i></p>
+
+### Datasets
+
+We propose two datasets; **Sioux-Cranfield** and **Sioux-Scans**, to address the gap between synthetic datasets and real-world industrial data.
+
+<p align="center">
+  <table>
+    <tr>
+      <td align="center">
+        <img src="assets/sioux_cranfield.png" height="250">
+        <br>
+        <sub><b>Sioux-Cranfield</b></sub>
+      </td>
+      <td align="center">
+        <img src="assets/sioux_scans.png" height="250">
+        <br>
+        <sub><b>Sioux-Scans</b></sub>
+      </td>
+    </tr>
+  </table>
+</p>
+<p align="left"><i>Figure 3. CAD models of the Sioux-Cranfield dataset (Left). The first six belong to the Cranfield Assembly benchmark and the rest are contributions of this paper (Sioux dataset). Sioux-Scans point cloud data (Right). Target (blue) and Source (yellow) point clouds for seven distinct objects.</i></p>
+
+## Environment Setup
+
+```bash
+# 1. Create environment
+conda env create -f environment.yml
+conda activate r3pm_net
+
+# Optionally, install the dependencies and run manually:
+pip install -e .
+```
+
+To run the evaluations, please refer to each method's repo to set up the environment:
+[Predator](https://github.com/prs-eth/OverlapPredator),
+[GeoTransformer](https://github.com/qinzheng93/geotransformer),
+[LoGDesc](https://github.com/karim416/LoGDesc), and 
+[RegTR](https://github.com/yewzijian/regtr).
+
+Everything must be installed into the **same** conda enviromnet.
+
+## Data Preparation
+
+### ModelNet40
+
+Download the dataset from [ModelNet40](http://modelnet.cs.princeton.edu/ModelNet40.zip) and extract it to:
+
+```
+data/ModelNet40
+```
+
+To save time, download the downsampled ModelNet40 test set from [ModelNet40_Downsampled](https://huggingface.co/datasets/YasiiKB/R3PM-Net/blob/main/down_sampled_modelnet40.zip) and put it in:
+
+```
+data/down_sampled_modelnet40
+```
+
+### Sioux-Cranfield
+
+Download the dataset from [Sioux_Cranfiled](https://huggingface.co/datasets/YasiiKB/R3PM-Net/blob/main/sioux_cranfield.zip) and put it in:
+
+```
+data/sioux_cranfield
+```
+
+### Sioux-Scans
+
+Download the dataset from [Sioux_Scans](https://huggingface.co/datasets/YasiiKB/R3PM-Net/blob/main/sioux_scans.zip) and put it in:
+
+```
+data/sioux_scans
+```
+
+### Fine-tune 
+
+Download the pickle files (.pkl) from [here](https://huggingface.co/datasets/YasiiKB/R3PM-Net/blob/main/simulators.zip) and put them in:
+
+```
+data/simulators
+```
+These pickle files are created from a subset of the Sioux-Cranfield containing the "teeth", "cube", "lime" and "lego" CAD models. There are 320 point cloud pairs, with 80-20 train-test split. 
+
+Optionally, to create your own datasets, use the scripts in `dataloader`, refering to the README file in that directory.
+
+## Pre-trained Models
+
+Please download the pretrained model of each method from their repo (links provided above) and follow their instructions as to where to put them.
+
+We use RPMNet's pre-trained model (*clean-trained*) for our Zero-shot version. Download it from [here](https://github.com/vinits5/learning3d/tree/master/pretrained/exp_rpmnet/models) and put it in:
+
+```
+checkpoints/
+```
+
+*Note:* You need to fine-tune the model yourself (see bleow) to get the fine-tuned weights which then you can put in the same directory. 
+
+## Folder Structure
+
+```text
+r3pm_net/
+├── assets/                     
+├── config/
+│   ├── default.yaml            # Training defaults
+│   └── eval.yaml               # Paths for evaluation scripts
+├── checkpoints/                # Pre-trained models' weights     
+├── data/                       
+│   ├── down_sampled_modelnet40/
+│   ├── ModelNet40/
+│   ├── sioux_cranfield/
+│   └── sioux_scans/
+├── dataloader/                 # Dataset dict generation & loaders
+├── logs/                       # Experiment logs
+├── r3pm_net/                   # Core package (model, feature extractor, config)
+├── scripts/                    # SLURM/Bash and evaluation scripts
+│   ├── eval_modelnet40.py
+│   ├── eval_sioux_cranfield.py
+│   ├── eval_sioux_scans.py
+│   ├── modelnet40.sh
+│   ├── sioux_cranfield.sh
+│   └── sioux_scans.sh
+├── src/
+│   └── train.py                # Training 
+├── thirdparty/learning3d/      # learning3d (RPMNet, losses, ops, …)
+├── tools/                      # Registration eval, metrics, visualization
+├── environment.yml 
+├── pyproject.toml
+└── README.md
+```
+
+## Train 
+
+To train the model using `data/simulators` or your own dataset run:
+```bash
+python src/train.py
+```
+
+## Evaluation
+
+Scripts are provided in `scripts/` to reproduce results.
+
+**ModelNet40**
+
+```bash
+bash scripts/modelnet40.sh
+```
+
+**Sioux-Cranfield**
+
+```bash
+bash scripts/sioux_cranfield.sh
+```
+
+**Sioux-Scans**
+This evaluates the proposed hybrid Coarse-to-Fine Registration approach.
+
+```bash
+bash scripts/sioux_scans.sh
+```
+
+### Manual Execution
+
+For example for evaluation on `Sioux-Cranfield`, run:
+
+```bash
+python scripts/eval_sioux_cranfield.py
+```
+
+## Results
+*IMPORTANT NOTE: Unfortunately, we cannot release the feature-extraction model and the fine-tuned weights. Therefore, to re-poduce these results you need to implement the feature extractor (based on the paper) and fine-tune it with the provided data.*
+
+### ModelNet40
+
+
+| Method              | RRE [°] ↓         | RTE [cm] ↓        | CD [cm] ↓         | Fitness ↑         | In. RMSE [cm] ↓ | Time [s] ↓        |
+| ------------------- | ----------------- | ----------------- | ----------------- | ----------------- | ------------------ | ----------------- |
+| RPMNet              | 30.898    | **0.002** | 0.153     | *0.998* | 0.094      | *0.021* |
+| Predator            | 7.262      | 0.028 | *0.045* | **1.000** | *0.026*  | 0.071      |
+| GeoTransformer      | 50.357 | 0.215 | 0.255 | 0.921 | 0.101 | 0.065 |
+| RegTR               | **1.712** | *0.007* | **0.017** | **1.000** | **0.009**  | 0.045     |
+| LoGDesc             | 42.762 | 0.158 | 0.183 | 0.978 | 0.097 | 0.075     |
+| **R3PM-Net (ours)** | *5.198* | 0.010   | 0.052     | **1.000** | 0.029      | **0.007** |
+
+
+> **Notes:** **Best** results are in bold; *Second-best* results are underlined.
+
+### Sioux-Cranfield
+
+
+| Method              | RRE [°] ↓         | RTE [cm] ↓        | CD [cm] ↓         | Fitness ↑         | In. RMSE [cm] ↓ | Time [s] ↓        |
+| ------------------- | ----------------- | ----------------- | ----------------- | ----------------- | ------------------ | ----------------- |
+| RPMNet              | 32.217    | **0.002** | 0.160   | *0.997* | 0.098 | 0.021  |
+| Predator            | 16.448    | 0.044     | 0.072   | **1.000** | 0.042      | 0.071     |
+| GeoTrans.           | 45.582    | 0.183     | 0.297   | 0.906     | 0.111       | 0.065     |
+| RegTR               | **1.311** | *0.004* | **0.023** | **1.000** | **0.012**  | 0.045   |
+| LoGDesc             | 121.224 | 0.773    | 0.692   | 0.718     | 0.224 | 0.075     |
+| **R3PM-Net (ours)** | *5.451* | 0.006 | *0.054* | **1.000** | *0.030*  | **0.006** |
+
+
+### Sioux-Scans
+<p align="center"> <img src="assets/success_cases.png" width="85%"> </p>
+
+<p align="left"><i>Figure 4. Qualitative registration results of R3PM-Net on real-world event-camera data. It successfully aligns the "teeth" and "cube" models. The fine-tuned version also solves the "lime" and "house".</i></p>
+
+## Acknowledgement
+
+We adapted some codes from some awesome repositories including [Learning3D](https://github.com/vinits5/learning3d) and [RPMNet](https://github.com/yewzijian/RPMNet). Thanks for making the codes publicly available.
+
+## Citation
+
+If you find this repository useful, please consider citing:
+
+```bibtex
+@misc{kashefbahrami2026r3pmnetrealtimerobustrealworld,
+      title={R3PM-Net: Real-time, Robust, Real-world Point Matching Network}, 
+      author={Yasaman Kashefbahrami and Erkut Akdag and Panagiotis Meletis and Evgeniya Balmashnova and Dip Goswami and Egor Bondarau},
+      year={2026},
+      eprint={2604.05060},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2604.05060}, 
+}
+```
+

config/default.yaml

@@ -0,0 +1,24 @@
+# Default training and data paths for R3PM-Net 
+
+exp_name: exp_r3pmnet
+eval: false
+save_dir: ""
+
+fine_tune_feature_extractor: tune
+transfer_weights: "" # Optional: leave empty to skip loading
+emb_dims: 1024
+symfn: max
+
+seed: 1234
+workers: 4
+batch_size: 5
+epochs: 2
+start_epoch: 0
+optimizer: Adam
+resume: ""
+pretrained: ""
+device: cuda:0
+
+# Pickled Registration Dataset dicts (keys: template, source, transformation)
+train_dict_path: data/simulators/data_dict_train.pkl
+test_dict_path: data/simulators/data_dict_test.pkl

config/eval.yaml

@@ -0,0 +1,28 @@
+# Paths for scripts/evaluation (loaded by scripts/*.py).
+
+data_root: data
+pretrained_rpmnet_dir: checkpoints
+
+modelnet40:
+  dataset_path: data/ModelNet40
+  cache_dir: data/down_sampled_modelnet40
+
+sioux:
+  base_dir: data
+  
+methods:
+  geotransformer:
+    root: GeoTransformer
+    exp_subdir: GeoTransformer/experiments/geotransformer.modelnet.rpmnet.stage4.gse.k3.max.oacl.stage2.sinkhorn
+    weights_path: GeoTransformer/weights/geotransformer-modelnet.pth.tar
+  predator:
+    root: OverlapPredator
+    config_path: OverlapPredator/configs/test/modelnet.yaml
+    weights_path: null
+  logdesc:
+    root: LoGDesc
+    weights_path: LoGDesc/pre-trained/best_model.pth
+  regtr:
+    root: RegTR
+    ckpt_path: RegTR/trained_models/modelnet/ckpt/model-best.pth
+    config_path: RegTR/trained_models/modelnet/config.yaml

dataloader/README.md

@@ -0,0 +1,71 @@
+# Dataloaders
+
+This directory contains scripts to generate Simulated datasets for training and testing.
+It uses functionalities from `tools` folder to:
+
+`generate_dataset`
+
+- load and downsample data
+- compute normals if needed
+- apply random transformations
+- add augmentations (noise, outliers and occlusion)
+- save point clouds
+
+**Note: Two input point clouds must have same length.** 
+
+`generate_dataset_dict`
+
+- save generated dataset in dictionaries suitable to train models (following Learning3d requirments)
+- checks dimensions (to meet Learning3d requirments)
+
+`combine_dataset_dict`
+
+- shuffle and combine all generated dictionaries into one 
+- split train and test sets
+
+## How to generate datasets?
+
+Modify the `args.txt` file to contain the correct paths and other specifications e.g. downsampling rate, noise level, etc. Other default arguments in `data_dict_generator.py` can also be changed.
+
+1. Generate transformed target point clouds + GT transforms
+Change `--action in dataloader/args.txt` to `generate_dataset` and run:
+```
+python dataloader/data_dict_generator.py @dataloader/args.txt
+```
+
+2. Generate train/test .pkl dicts 
+Change `--action in dataloader/args.txt` to `generate_dataset_dict`, then run the above script again.
+
+3. Combine multiple object dicts 
+Set `--action combine_dataset_dict` and run again to get train and test `dict.pkl` files.
+
+### Manual Run (without args.txt)
+Optinally you can manually run: 
+```
+python dataloader/data_dict_generator.py \
+  --pcdPath /path/to/source_scan.pcd \
+  --cadPath /path/to/object.stl \
+  --name teeth \
+  --action generate_dataset \
+  --every_k_points 100 \
+  --num_transformation 50 \
+  --angles 0 90 180 \
+  --translation_range -1 1 \
+  --index 0 \
+  --noise_level 0 \
+  --outlier_level 0 \
+  --outlier_bounds -0.05 0.05 \
+  --occ_level 0 \
+  --save_path data/simulators
+```
+then:
+```
+python dataloader/data_dict_generator.py \
+  --pcdPath /path/to/source_scan.pcd \
+  --cadPath /path/to/object.stl \
+  --name teeth \
+  --action generate_dataset_dict \
+  --dataset_size 50 \
+  --index 0 \
+  --save_path data/simulators
+```

dataloader/__init__.py

@@ -0,0 +1 @@
+# Dataset loaders and generators

dataloader/args.txt

@@ -0,0 +1,14 @@
+--pcdPath data/sioux_scans/teeth_clean.ply
+--cadPath data/sioux_cranfield/teeth.stl
+--action combine_dataset_dict
+--name teeth
+--every_k_points 100
+--num_transformation 50
+--angles 0 90 180
+--translation_range -1 1
+--dataset_size 50
+--index 0
+--noise_level 0
+--outlier_level 0
+--outlier_bounds -0.05 0.05
+--occ_level 0

dataloader/data_dict_generator.py

@@ -0,0 +1,119 @@
+import argparse
+import copy
+import logging
+import os
+import shlex
+import sys
+from pathlib import Path
+
+import numpy as np
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tools import data
+from dataloader.dataset_generator import combine_dataset_dict, generate_dataset, generate_dataset_dict
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def main():
+    # Set up the argument parser
+    parser = argparse.ArgumentParser(description='Automate dataset generation and processing.')
+
+    # Define arguments (change these as needed)
+    parser.add_argument('--pcdPath', type=str, required=True, help='Path to the PCD file')
+    parser.add_argument('--cadPath', type=str, required=True, help='Path to the CAD file')
+    parser.add_argument('--action', type=str, choices=['generate_dataset', 'generate_dataset_dict', 'combine_dataset_dict'], required=True, help='Action to perform')
+    parser.add_argument('--compute_normals', action='store_true', help='Flag to compute normals')
+    parser.add_argument('--every_k_points', type=int, default=1, help='Sampling rate for points')
+    parser.add_argument('--save', action='store_true', help='Flag to save the generated dataset')
+    parser.add_argument(
+        '--save_path',
+        type=str,
+        default='data/simulators',
+        help='Directory to save generated datasets (relative to repo root if not absolute)',
+    )
+    parser.add_argument('--name', type=str, required=True, help='Name identifier for the dataset (e.g., teeth, cube, etc.)')
+
+    # Additional parameters for dataset generation (change these as needed)
+    parser.add_argument('--num_transformation', type=int, default=50, help='Number of transformations')
+    parser.add_argument('--angles', type=int, nargs='+', default=list(range(0, 360, 10)), help='Rotation angles')
+    parser.add_argument('--translation_range', type=float, nargs=2, default=(-1, 1), help='Translation range')
+    parser.add_argument('--dataset_size', type=int, default=400, help='Size of the dataset to generate')
+    parser.add_argument('--index', type=int, default=0, help='Index for dataset generation')
+    parser.add_argument('--noise_level', type=float, default=0, help='Noise level')
+    parser.add_argument('--outlier_level', type=float, default=0, help='Outlier level')
+    parser.add_argument('--outlier_bounds', type=float, nargs=2, default=(-10, 10), help='Outlier bounds')
+    parser.add_argument('--occ_level', type=float, default=0, help='Occlusion level')
+
+    # Parse the arguments
+
+    # Check if an argument file is being used
+    if sys.argv[1].startswith('@'):
+        args_file = sys.argv[1][1:]  # Strip the '@' from the filename
+        with open(args_file, 'r') as file:
+            # Read and split arguments from the file
+            args = parser.parse_args(shlex.split(file.read()))
+    else:
+        args = parser.parse_args()
+
+    # Print out the arguments to verify
+    print(vars(args))
+
+    # Load the data
+    np.random.seed(42)
+    if args.compute_normals:
+        _, cad, _, cad_normals = data.load_data(args.pcdPath, args.cadPath, every_k_points=args.every_k_points, same_length=True, compute_normals=True)
+        suffix = '_with_normals'
+    else:
+        _, cad = data.load_data(args.pcdPath, args.cadPath, every_k_points=args.every_k_points, same_length=True)
+        cad_normals = None
+        suffix = ''
+    source = copy.deepcopy(cad)
+
+    rp = Path(args.save_path)
+    if not rp.is_absolute():
+        rp = _REPO_ROOT / args.save_path
+    ROOT_DIR = str(rp.resolve())
+    if not ROOT_DIR.endswith(os.sep):
+        ROOT_DIR += os.sep
+
+    # Perform the selected action
+    if args.action == 'generate_dataset':
+        logging.info('Generating dataset...')
+        generate_dataset(source, args.pcdPath, args.cadPath, args.num_transformation, args.angles, args.translation_range, args.index, args.noise_level, args.outlier_level, args.outlier_bounds, args.occ_level, save_dir=ROOT_DIR)
+
+    elif args.action == 'generate_dataset_dict':
+        logging.info('Generating dataset dictionary...')
+        output_train_file = f'{ROOT_DIR}data_dict_train_{args.name}{suffix}.pkl'
+        output_test_file = f'{ROOT_DIR}data_dict_test_{args.name}{suffix}.pkl'
+        generate_dataset_dict(source, args.dataset_size, args.index, output_train_file, output_test_file, cad_normals)
+
+    elif args.action == 'combine_dataset_dict':
+        logging.info('Combining dataset dictionaries...')
+        train_files = [
+            f'{ROOT_DIR}data_dict_train_teeth{suffix}.pkl'
+            # f'{ROOT_DIR}data_dict_train_elephant{suffix}.pkl',
+            # f'{ROOT_DIR}data_dict_train_house{suffix}.pkl',
+            # f'{ROOT_DIR}data_dict_train_shoe{suffix}.pkl'
+        ]
+
+        test_files = [
+            f'{ROOT_DIR}data_dict_test_teeth{suffix}.pkl'
+            # f'{ROOT_DIR}data_dict_test_elephant{suffix}.pkl',
+            # f'{ROOT_DIR}data_dict_test_house{suffix}.pkl',
+            # f'{ROOT_DIR}data_dict_test_shoe{suffix}.pkl'
+        ]
+
+        output_train_file = f'{ROOT_DIR}data_dict_train_{suffix}.pkl'
+        output_test_file = f'{ROOT_DIR}data_dict_test_{suffix}.pkl'
+
+        combine_dataset_dict(train_files, test_files, output_train_file, output_test_file)
+
+    else:
+        logging.warning('No valid action selected.')
+
+if __name__ == '__main__':
+    main()

dataloader/dataset_generator.py

@@ -0,0 +1,258 @@
+import copy
+import os
+import pickle
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import open3d as o3
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tools import augmentation, data, transformations
+
+_SIM_DATA = _REPO_ROOT / "data" / "simulators"
+'''
+This module provides functions to generate a dataset of point clouds with random transformations, with options for noise, outliers, and occlusions.
+It also includes functions to check the shape of the data and to generate a data dictionary for training and testing,
+and a function to combine multiple dataset dictionaries.
+'''
+
+def generate_dataset(pcd, pcdPath, cadPath, num_transformation, angles, translation_range, index, noise_level = 0, outlier_level = 0, outlier_bounds = (-10, 10), occ_level = 0, save_dir=None):
+    '''
+    A function to generate a dataset of point clouds with random transformations.
+
+    Args:
+        pcd (open3d.geometry.PointCloud): The source point cloud
+        pcdPath (str): The path to the source point cloud
+        cadPath (str): The path to the target point cloud
+        num_transformation (int): The number of transformations to generate
+        angles (numpy.ndarray): The range of angles for the random transformations
+        translation_range (tuple): The range of translations for the random transformations
+        index (int): The index to start saving the generated dataset
+        noise_level (float): The level of noise to add to the point clouds
+        outlier_level (float): The level of outliers to add to the point clouds
+        occ_level (float): The level of occlusions to add to the point clouds
+        save (bool): A flag to save the generated dataset
+
+    Returns:
+        None
+    '''
+    np.random.seed(42)
+    target_list = []
+    gt_transformation_list = []
+
+    for i in range(num_transformation):
+        # Generate random gt transformation
+        x_angle= np.random.uniform(angles[0], angles[-1], size=1)
+        y_angle= np.random.uniform(angles[0], angles[-1], size=1)
+        z_angle= np.random.uniform(angles[0], angles[-1], size=1)
+        gt_transformation = transformations.create_transformation(x_angle, y_angle, z_angle, translation_range)
+
+        target = copy.deepcopy(pcd)
+        target.transform(gt_transformation)
+
+        if noise_level != 0:
+            target = augmentation.apply_noise(target, noise_level)
+            print('Noise applied')
+
+        if outlier_level != 0 or occ_level != 0:
+            _, another_cad = data.load_data(pcdPath, cadPath, every_k_points=1)
+            target = copy.deepcopy(another_cad).transform(gt_transformation)
+            if occ_level != 0:
+                target, _ = augmentation.apply_occlusion(target, occ_level)
+                print('Occlusion applied')
+            if outlier_level != 0:
+                target = augmentation.add_outliers(target, outlier_level, outlier_lowerbound=outlier_bounds[0], outlier_upperbound=outlier_bounds[1])
+                print('Outliers applied')
+
+        # randomly take points away from target to get to same length as source
+        if len(target.points) >= len(pcd.points):
+            np.random.seed(42)
+            target_points = np.asarray(target.points) 
+            indices = np.random.choice(len(target_points), 1441, replace=False)  # change len(source.points) to a specific num if you want to have a fixed number of points
+            sampled_points = target_points[indices]
+            target.points = o3.utility.Vector3dVector(sampled_points)
+        else:
+            print('Target has fewer points than source and can\'t be downsampled to the same length.')
+
+        print(f'size of source and target: {len(pcd.points)}, {len(target.points)}')
+        target_list.append(target)
+        gt_transformation_list.append(gt_transformation)
+
+    # Save the generated dataset
+    if save_dir is not None:
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+
+        for i, (target, transformation) in enumerate(zip(target_list, gt_transformation_list)):
+            target_path = os.path.join(save_dir, f"target_{i+index}.pcd")
+            transformation_path = os.path.join(save_dir, f"transformation_{i+index}.npy")
+            o3.io.write_point_cloud(target_path, target)
+            np.save(transformation_path, transformation)
+
+def check_shape(data, expected_shape_3d, expected_shape_6d):
+    return data.shape == expected_shape_3d or data.shape == expected_shape_6d
+
+def generate_dataset_dict(source, dataset_size, index, output_train_file_path, output_test_file_path, source_normals = None):
+    '''
+    This function shuffles the dataset and generates a data_dict for the training and testing data following the pattern acceptable to Learning3D.
+    
+    Args:
+        source (open3d.geometry.PointCloud): The source point cloud
+        dataset_size (int): The size of the dataset
+
+    Returns:
+        None
+    '''
+    np.random.seed(42)
+    transformed_pcds = []
+    gt_transformations = []
+
+    # Load the transformed point clouds and ground truth transformations
+    for i in range(index,index+dataset_size):
+        transformed_pcd = o3.io.read_point_cloud(str(_SIM_DATA / f"target_{i}.pcd"))
+        gt_transformation = np.load(str(_SIM_DATA / f"transformation_{i}.npy"))
+
+        if source_normals is not None: # we also need target normals 
+            M = np.linalg.inv(gt_transformation).T
+            target_normals = np.dot(source_normals, M[:3,:3]) # transformed_normals = normals * (transformation)^-1.T
+            transformed_points = np.concatenate((np.asarray(transformed_pcd.points), target_normals), axis=1)
+        else:
+            transformed_points = np.asarray(transformed_pcd.points).astype(np.float32)
+
+        transformed_pcds.append(transformed_points)
+        gt_transformations.append(gt_transformation)
+
+    # Shuffle the transformed point clouds and ground truth transformations in the same way
+    temp = list(zip(transformed_pcds, gt_transformations))
+    random.shuffle(temp) 
+    transformed_pcds, gt_transformations = zip(*temp)
+
+    # Convert lists to numpy arrays
+    transformed_pcds_np = np.array(transformed_pcds)
+    gt_transformations_np = np.array(gt_transformations)
+
+    if source_normals is not None:
+        source = np.concatenate((np.asarray(source.points), source_normals), axis=1)
+    else:
+        source = np.asarray(source.points).astype(np.float32)
+
+    data_dict = {
+        'template': np.tile(source, (dataset_size, 1, 1)),
+        'source': transformed_pcds_np,
+        'transformation': gt_transformations_np
+    }
+
+    # Split the data_dict into training and testing data_dict
+    train_size = int(0.8 * dataset_size)
+    test_size = dataset_size - train_size
+    num_points = len(source)
+
+    data_dict_train = {}
+    data_dict_test = {}
+    for key in data_dict.keys():
+        data_dict_train[key] = data_dict[key][0:train_size]
+        data_dict_test[key] = data_dict[key][train_size:]
+    
+    assert set(data_dict_train.keys()) == {'template', 'source', 'transformation'}
+    assert set(data_dict_test.keys()) == {'template', 'source', 'transformation'}
+
+    expected_shape_3d_train = (train_size, num_points, 3)
+    expected_shape_6d_train = (train_size, num_points, 6)
+
+    assert check_shape(data_dict_train['template'], expected_shape_3d_train, expected_shape_6d_train), f"Expected shape: {expected_shape_3d_train} or {expected_shape_6d_train}, but got {data_dict_train['template'].shape}"
+    assert check_shape(data_dict_train['source'], expected_shape_3d_train, expected_shape_6d_train), f"Expected shape: {expected_shape_3d_train} or {expected_shape_6d_train}, but got {data_dict_train['source'].shape}"
+    assert data_dict_train['transformation'].shape == (train_size, 4, 4), f"Expected shape: {(train_size, 4, 4)}, but got {data_dict_train['transformation'].shape}"
+
+    expected_shape_3d_test = (test_size, num_points, 3)
+    expected_shape_6d_test = (test_size, num_points, 6)
+
+    assert check_shape(data_dict_test['template'], expected_shape_3d_test, expected_shape_6d_test), f"Expected shape: {expected_shape_3d_test} or {expected_shape_6d_test}, but got {data_dict_test['template'].shape}"
+    assert check_shape(data_dict_test['source'], expected_shape_3d_test, expected_shape_6d_test), f"Expected shape: {expected_shape_3d_test} or {expected_shape_6d_test}, but got {data_dict_test['source'].shape}"
+    assert data_dict_test['transformation'].shape == (test_size, 4, 4), f"Expected shape: {(test_size, 4, 4)}, but got {data_dict_test['transformation'].shape}"
+    
+    with open(output_train_file_path, 'wb') as f:
+        pickle.dump(data_dict_train, f)
+    print(f"train_dict saved to {output_train_file_path}")
+
+    with open(output_test_file_path, 'wb') as f:
+        pickle.dump(data_dict_test, f)
+    print(f"test_dict saved to {output_test_file_path}")
+
+
+def combine_dataset_dict(train_files, test_files, output_train_file_path, output_test_file_path):
+    '''
+    Combine and shuffle dictionaries from multiple files.
+
+    Args:
+        train_files (list of str): List of file paths to training dictionaries.
+        test_files (list of str): List of file paths to testing dictionaries.
+        output_train_file (str): Output file path for the combined training dictionary.
+        output_test_file (str): Output file path for the combined testing dictionary.
+    '''
+    
+    # Load the dictionaries from the .pkl files
+    train_dicts = [pickle.load(open(file, 'rb')) for file in train_files]
+    test_dicts = [pickle.load(open(file, 'rb')) for file in test_files]
+
+    # Combine the dictionaries
+    combined_train_dict = {}
+    combined_test_dict = {}
+
+    for key in train_dicts[0].keys():
+        combined_train_dict[key] = np.concatenate([d[key] for d in train_dicts], axis=0)
+        combined_test_dict[key] = np.concatenate([d[key] for d in test_dicts], axis=0)
+
+    # Shuffle
+    train_combined_list = list(zip(combined_train_dict['template'], combined_train_dict['source'], combined_train_dict['transformation']))
+    test_combined_list = list(zip(combined_test_dict['template'], combined_test_dict['source'], combined_test_dict['transformation']))
+
+    random.shuffle(train_combined_list)
+    random.shuffle(test_combined_list)
+
+    combined_train_dict['template'], combined_train_dict['source'], combined_train_dict['transformation'] = zip(*train_combined_list)
+    combined_test_dict['template'], combined_test_dict['source'], combined_test_dict['transformation'] = zip(*test_combined_list)
+
+    # Convert back to numpy arrays
+    combined_train_dict['template'] = np.array(combined_train_dict['template'])
+    combined_train_dict['source'] = np.array(combined_train_dict['source'])
+    combined_train_dict['transformation'] = np.array(combined_train_dict['transformation'])
+
+    combined_test_dict['template'] = np.array(combined_test_dict['template'])
+    combined_test_dict['source'] = np.array(combined_test_dict['source'])
+    combined_test_dict['transformation'] = np.array(combined_test_dict['transformation'])
+
+    # Checks 
+    train_size = len(combined_train_dict['source'])
+    test_size = len(combined_test_dict['source'])
+    num_points = combined_train_dict['source'].shape[1]
+    
+    assert set(combined_train_dict.keys()) == {'template', 'source', 'transformation'}
+    assert set(combined_test_dict.keys()) == {'template', 'source', 'transformation'}
+
+    expected_shape_3d_train = (train_size, num_points, 3)
+    expected_shape_6d_train = (train_size, num_points, 6)
+
+    assert check_shape(combined_train_dict['template'], expected_shape_3d_train, expected_shape_6d_train), f"Expected shape: {expected_shape_3d_train} or {expected_shape_6d_train}, but got {combined_train_dict['template'].shape}"
+    assert check_shape(combined_train_dict['source'], expected_shape_3d_train, expected_shape_6d_train), f"Expected shape: {expected_shape_3d_train} or {expected_shape_6d_train}, but got {combined_train_dict['source'].shape}"
+    assert combined_train_dict['transformation'].shape == (train_size, 4, 4), f"Expected shape: {(train_size, 4, 4)}, but got {combined_train_dict['transformation'].shape}"
+
+    expected_shape_3d_test = (test_size, num_points, 3)
+    expected_shape_6d_test = (test_size, num_points, 6)
+
+    assert check_shape(combined_test_dict['template'], expected_shape_3d_test, expected_shape_6d_test), f"Expected shape: {expected_shape_3d_test} or {expected_shape_6d_test}, but got {combined_test_dict['template'].shape}"
+    assert check_shape(combined_test_dict['source'], expected_shape_3d_test, expected_shape_6d_test), f"Expected shape: {expected_shape_3d_test} or {expected_shape_6d_test}, but got {combined_test_dict['source'].shape}"
+    assert combined_test_dict['transformation'].shape == (test_size, 4, 4), f"Expected shape: {(test_size, 4, 4)}, but got {combined_test_dict['transformation'].shape}"
+    
+    # Save the dictionaries
+    with open(output_train_file_path, 'wb') as f:
+        pickle.dump(combined_train_dict, f)
+    print(f"combined_train_dict saved to {output_train_file_path}")
+
+    with open(output_test_file_path, 'wb') as f:
+        pickle.dump(combined_test_dict, f)
+    print(f"combined_test_dict saved to {output_train_file_path}")

dataloader/user_data.py

@@ -0,0 +1,127 @@
+import os
+import numpy as np
+import torch
+
+class ClassificationData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.pcs = self.find_attribute('pcs')
+		self.labels = self.find_attribute('labels')
+		self.check_data()
+
+	def find_attribute(self, attribute):
+		try:
+			attribute_data = self.data_dict[attribute]
+		except:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.pcs.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.pcs.shape)
+		assert 0 < len(self.labels.shape) < 3, "Error in dimension of labels! Given data dimension: {}".format(self.labels.shape)
+		
+		if len(self.pcs.shape)==2: self.pcs = self.pcs.reshape(1, -1, 3)
+		if len(self.labels.shape) == 1: self.labels = self.labels.reshape(1, -1)
+
+		assert self.pcs.shape[0] == self.labels.shape[0], "Inconsistency in the number of point clouds and number of ground truth labels!"
+
+
+	def __len__(self):
+		return self.pcs.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.pcs[index]).float(), torch.from_numpy(self.labels[idx]).type(torch.LongTensor)
+
+
+class RegistrationData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.template = self.find_attribute('template')
+		self.source = self.find_attribute('source')
+		self.transformation = self.find_attribute('transformation')
+		self.check_data()
+
+	# def find_attribute(self, attribute):
+	# 	try:
+	# 		attribute_data = self.data[attribute]
+	# 	except:
+	# 		print("Given data directory has no key attribute \"{}\"".format(attribute))
+	# 	return attribute_data
+	
+	def find_attribute(self, attribute):
+		attribute_data = None
+		if attribute in self.data_dict:
+			attribute_data = self.data_dict[attribute]
+		else:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.template.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.template.shape)
+		assert 1 < len(self.source.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.source.shape)
+		assert 1 < len(self.transformation.shape) < 4, "Error in dimension of transformations! Given data dimension: {}".format(self.transformation.shape)
+
+		if len(self.template.shape)==2: self.template = self.template.reshape(1, -1, 3)
+		if len(self.source.shape)==2: self.source = self.source.reshape(1, -1, 3)
+		if len(self.transformation.shape) == 2: self.transformation = self.transformation.reshape(1, 4, 4)
+
+		assert self.template.shape[0] == self.source.shape[0], "Inconsistency in the number of template and source point clouds!"
+		assert self.source.shape[0] == self.transformation.shape[0], "Inconsistency in the number of transformation and source point clouds!"
+
+	def __len__(self):
+		return self.template.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.template[index]).float(), torch.tensor(self.source[index]).float(), torch.tensor(self.transformation[index]).float()
+
+
+class FlowData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.frame1 = self.find_attribute('frame1')
+		self.frame2 = self.find_attribute('frame2')
+		self.flow = self.find_attribute('flow')
+		self.check_data()
+
+	def find_attribute(self, attribute):
+		try:
+			attribute_data = self.data[attribute]
+		except:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.frame1.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame1.shape)
+		assert 1 < len(self.frame2.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame2.shape)
+		assert 1 < len(self.flow.shape) < 4, "Error in dimension of flow! Given data dimension: {}".format(self.flow.shape)
+
+		if len(self.frame1.shape)==2: self.frame1 = self.frame1.reshape(1, -1, 3)
+		if len(self.frame2.shape)==2: self.frame2 = self.frame2.reshape(1, -1, 3)
+		if len(self.flow.shape) == 2: self.flow = self.flow.reshape(1, -1, 3)
+
+		assert self.frame1.shape[0] == self.frame2.shape[0], "Inconsistency in the number of frame1 and frame2 point clouds!"
+		assert self.frame2.shape[0] == self.flow.shape[0], "Inconsistency in the number of flow and frame2 point clouds!"
+
+	def __len__(self):
+		return self.frame1.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.frame1[index]).float(), torch.tensor(self.frame2[index]).float(), torch.tensor(self.flow[index]).float()
+
+
+class UserData:
+	def __init__(self, application, data_dict):
+		self.application = application
+
+		if self.application == 'classification':
+			self.data_class = ClassificationData(data_dict)
+		elif self.application == 'registration':
+			self.data_class = RegistrationData(data_dict)
+		elif self.application == 'flow_estimation':
+			self.data_class = FlowData(data_dict)
+
+	def __len__(self):
+		return len(self.data_class)
+
+	def __getitem__(self, index):
+		return self.data_class[index]

environment.yml

@@ -0,0 +1,16 @@
+name: r3pm_net
+channels:
+  - conda-forge
+  - pytorch
+dependencies:
+  - python=3.9 
+  - pip
+  - open3d
+  - pytorch
+  - hatchling
+  - ipykernel
+  - pip:
+      - tabulate
+      - pyyaml
+      - -e .
+

pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "r3pm_net"
+version = "0.1.0"
+description = "R3PM-Net point cloud registration"
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+  "numpy",
+  "torch",
+  "tensorboardX",
+  "tqdm",
+  "pyyaml",
+  "open3d",
+  "tabulate",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["r3pm_net", "tools", "dataloader", "thirdparty"]
+
+[tool.jupytext]
+formats = "ipynb,py:light"

r3pm_net/__init__.py

@@ -0,0 +1,5 @@
+"""R3PM-Net: point cloud registration with PointNet features."""
+
+from .model import R3PMNet
+
+__all__ = ["R3PMNet"]

r3pm_net/config_loader.py

@@ -0,0 +1,164 @@
+"""Load YAML training config and merge with argparse."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import Any, Mapping
+
+import yaml
+
+from r3pm_net.paths import REPO_ROOT
+
+
+def _resolve_maybe_relative(path_str: str | None) -> str | None:
+    if path_str is None or path_str == "":
+        return path_str
+    p = Path(path_str)
+    if p.is_absolute():
+        return str(p)
+    return str(REPO_ROOT / p)
+
+
+def load_yaml_config(path: str | Path) -> dict[str, Any]:
+    with open(path, "r", encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    if data is None:
+        return {}
+    if not isinstance(data, Mapping):
+        raise ValueError(f"Config must be a mapping, got {type(data)}")
+    return dict(data)
+
+
+def _extract_config_argv(argv: list[str], default_cfg: str) -> tuple[str, list[str]]:
+    """Return (config path for YAML, argv without --config ...)."""
+    path = default_cfg
+    out: list[str] = []
+    i = 0
+    while i < len(argv):
+        if argv[i] == "--config" and i + 1 < len(argv):
+            path = argv[i + 1]
+            i += 2
+            continue
+        if argv[i].startswith("--config="):
+            path = argv[i].split("=", 1)[1]
+            i += 1
+            continue
+        out.append(argv[i])
+        i += 1
+    return path, out
+
+
+def parse_train_args(argv: list[str], build_parser) -> argparse.Namespace:
+    """Load YAML from --config (default: config/default.yaml), merge as argparse defaults, then parse CLI."""
+    default_cfg = str(REPO_ROOT / "config" / "default.yaml")
+    cfg_path, argv_rest = _extract_config_argv(list(argv), default_cfg)
+    cfg = load_yaml_config(cfg_path) if Path(cfg_path).is_file() else {}
+    parser = build_parser(cfg_path)
+    if cfg:
+        known = {
+            a.dest
+            for a in parser._actions
+            if getattr(a, "dest", None) and a.dest not in ("help", argparse.SUPPRESS)
+        }
+        filtered = {k: v for k, v in cfg.items() if k in known}
+        parser.set_defaults(**filtered)
+    return parser.parse_args(argv_rest)
+
+
+def resolve_path_args(ns: Any, path_keys: tuple[str, ...]) -> None:
+    """Mutate namespace: resolve listed keys to absolute paths under REPO_ROOT when relative."""
+    for key in path_keys:
+        val = getattr(ns, key, None)
+        if isinstance(val, str) and val:
+            setattr(ns, key, _resolve_maybe_relative(val))
+
+
+def load_eval_yaml() -> dict[str, Any]:
+    """Load ``config/eval.yaml`` if present; otherwise return an empty dict."""
+    path = REPO_ROOT / "config" / "eval.yaml"
+    if not path.is_file():
+        return {}
+    return load_yaml_config(path)
+
+
+def get_pretrained_rpmnet_dir() -> str:
+    """Directory containing ``clean-trained.pth``, ``best_model_PointNet*.t7``, etc.
+
+    ``R3PM_NET_PRETRAINED_ROOT`` overrides ``pretrained_rpmnet_dir`` in ``config/eval.yaml``.
+    """
+    env = os.environ.get("R3PM_NET_PRETRAINED_ROOT")
+    if env:
+        return str(Path(env).expanduser().resolve())
+    cfg = load_eval_yaml()
+    rel = (cfg.get("pretrained_rpmnet_dir") or "checkpoints").strip()
+    if not rel:
+        rel = "checkpoints"
+    out = _resolve_maybe_relative(rel)
+    return out if out else str(REPO_ROOT / "checkpoints")
+
+
+def get_sioux_data_root() -> str:
+    """Base data directory for Sioux scripts (``data`` / ``sioux_cranfield``, etc.)."""
+    cfg = load_eval_yaml()
+    sioux = cfg.get("sioux") or {}
+    base = sioux.get("base_dir") or cfg.get("data_root") or "data"
+    out = _resolve_maybe_relative(str(base).strip())
+    return out if out else str(REPO_ROOT / "data")
+
+
+def get_modelnet40_paths() -> tuple[str, str]:
+    """Return ``(dataset_path, cache_dir)`` for ModelNet40 evaluation."""
+    cfg = load_eval_yaml()
+    m = cfg.get("modelnet40") or {}
+    ds = m.get("dataset_path", "data/ModelNet40")
+    cache = m.get("cache_dir", "data/down_sampled_modelnet40")
+    dsr = _resolve_maybe_relative(ds)
+    cr = _resolve_maybe_relative(cache)
+    return (
+        dsr if dsr else str(REPO_ROOT / "data" / "ModelNet40"),
+        cr if cr else str(REPO_ROOT / "data" / "down_sampled_modelnet40"),
+    )
+
+
+def get_method_paths() -> dict[str, Any]:
+    """Return resolved path configuration for external registration methods."""
+    cfg = load_eval_yaml()
+    methods = cfg.get("methods") or {}
+    out: dict[str, Any] = {}
+    for method_name, method_cfg in methods.items():
+        if not isinstance(method_cfg, Mapping):
+            continue
+        method_out: dict[str, Any] = {}
+        for k, v in method_cfg.items():
+            if isinstance(v, str) and v.strip():
+                rv = _resolve_maybe_relative(v.strip())
+                method_out[k] = rv if rv else v
+            else:
+                method_out[k] = v
+        out[str(method_name)] = method_out
+    return out
+
+
+def get_sioux_paths() -> dict[str, Any]:
+    """Return Sioux eval paths from config/eval.yaml with absolute paths."""
+    cfg = load_eval_yaml()
+    sioux = cfg.get("sioux") or {}
+    out: dict[str, Any] = {}
+    for k, v in sioux.items():
+        if isinstance(v, str) and v.strip():
+            rv = _resolve_maybe_relative(v.strip())
+            out[k] = rv if rv else v
+        elif isinstance(v, list):
+            vals = []
+            for item in v:
+                if isinstance(item, str) and item.strip():
+                    rv = _resolve_maybe_relative(item.strip())
+                    vals.append(rv if rv else item)
+                else:
+                    vals.append(item)
+            out[k] = vals
+        else:
+            out[k] = v
+    return out

r3pm_net/feature_extractor.py

@@ -0,0 +1,8 @@
+# Feature extractor for R3PM-Net
+'''
+Unfortunately, the feature extractor cannot be provided in the repository due to copyright issues.
+Please implement the feature extractor for R3PM-Net as described in the paper and place it in this file.
+Currently, the feature extractor is set to PPFNet (same as RPMNet).
+'''
+from thirdparty.learning3d.models import PPFNet
+feature_extractor = PPFNet()

r3pm_net/model.py

@@ -0,0 +1,382 @@
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from thirdparty.learning3d.utils import square_distance, angle_difference
+from thirdparty.learning3d.ops.transform_functions import convert2transformation
+
+_EPS = 1e-5  # To prevent division by zero
+
+class ParameterPredictionNet(nn.Module):
+	def __init__(self, weights_dim):
+		"""PointNet based Parameter prediction network
+
+		Args:
+			weights_dim: Number of weights to predict (excluding beta), should be something like
+						 [3], or [64, 3], for 3 types of features
+		"""
+
+		super().__init__()
+
+		self._logger = logging.getLogger(self.__class__.__name__)
+
+		self.weights_dim = weights_dim
+
+		# Pointnet
+		self.prepool = nn.Sequential(
+			nn.Conv1d(4, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 128, 1),
+			nn.GroupNorm(8, 128),
+			nn.ReLU(),
+
+			nn.Conv1d(128, 1024, 1),
+			nn.GroupNorm(16, 1024),
+			nn.ReLU(),
+		)
+		self.pooling = nn.AdaptiveMaxPool1d(1)
+		self.postpool = nn.Sequential(
+			nn.Linear(1024, 512),
+			nn.GroupNorm(16, 512),
+			nn.ReLU(),
+
+			nn.Linear(512, 256),
+			nn.GroupNorm(16, 256),
+			nn.ReLU(),
+
+			nn.Linear(256, 2 + np.prod(weights_dim)),
+		)
+
+		self._logger.info('Predicting weights with dim {}.'.format(self.weights_dim))
+
+	def forward(self, x):
+		""" Returns alpha, beta, and gating_weights (if needed)
+
+		Args:
+			x: List containing two point clouds, x[0] = src (B, J, 3), x[1] = ref (B, K, 3)
+
+		Returns:
+			beta, alpha, weightings
+		"""
+		# X and Y concatenated
+		src_padded = F.pad(x[0], (0, 1), mode='constant', value=0)
+		ref_padded = F.pad(x[1], (0, 1), mode='constant', value=1)
+		concatenated = torch.cat([src_padded, ref_padded], dim=1)
+
+		prepool_feat = self.prepool(concatenated.permute(0, 2, 1)) 
+		pooled = torch.flatten(self.pooling(prepool_feat), start_dim=-2)
+		raw_weights = self.postpool(pooled)
+
+		# softplus to ensure positivity
+		beta = F.softplus(raw_weights[:, 0])
+		alpha = F.softplus(raw_weights[:, 1])
+
+		return beta, alpha
+
+
+
+def to_numpy(tensor):
+	"""Wrapper around .detach().cpu().numpy() """
+	if isinstance(tensor, torch.Tensor):
+		return tensor.detach().cpu().numpy()
+	elif isinstance(tensor, np.ndarray):
+		return tensor
+	else:
+		raise NotImplementedError
+
+
+def se3_transform(g, a, normals=None):
+	""" Applies the SE3 transform
+
+	Args:
+		g: SE3 transformation matrix of size ([1,] 3/4, 4) or (B, 3/4, 4)
+		a: Points to be transformed (N, 3) or (B, N, 3)
+		normals: (Optional). If provided, normals will be transformed
+
+	Returns:
+		transformed points of size (N, 3) or (B, N, 3)
+
+	"""
+	R = g[..., :3, :3]  # (B, 3, 3)
+	p = g[..., :3, 3]  # (B, 3)
+
+	if len(g.size()) == len(a.size()):
+		b = torch.matmul(a, R.transpose(-1, -2)) + p[..., None, :]
+	else:
+		raise NotImplementedError
+		b = R.matmul(a.unsqueeze(-1)).squeeze(-1) + p  # No batch. Not checked
+
+	if normals is not None:
+		rotated_normals = normals @ R.transpose(-1, -2)
+		return b, rotated_normals
+
+	else:
+		return b
+
+def match_features(feat_src, feat_ref, metric='l2'):
+	""" Compute pairwise distance between features
+
+	Args:
+		feat_src: (B, J, C)
+		feat_ref: (B, K, C)
+		metric: either 'angle' or 'l2' (squared euclidean)
+
+	Returns:
+		Matching matrix (B, J, K). i'th row describes how well the i'th point
+		 in the src agrees with every point in the ref.
+	"""
+	if feat_src.shape[-1] != feat_ref.shape[-1]:
+		if feat_src.shape[-1] > feat_ref.shape[-1]:
+			feat_src = feat_src[:,:,:feat_ref.shape[-1]]
+		elif feat_src.shape[-1] < feat_ref.shape[-1]:
+			feat_ref = feat_ref[:,:,:feat_src.shape[-1]]
+
+	assert feat_src.shape[-1] == feat_ref.shape[-1]
+
+	if metric == 'l2':
+		dist_matrix = square_distance(feat_src, feat_ref)
+	elif metric == 'angle':
+		feat_src_norm = feat_src / (torch.norm(feat_src, dim=-1, keepdim=True) + _EPS)
+		feat_ref_norm = feat_ref / (torch.norm(feat_ref, dim=-1, keepdim=True) + _EPS)
+
+		dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
+	else:
+		raise NotImplementedError
+
+	return dist_matrix
+
+
+def sinkhorn(log_alpha, n_iters: int = 5, slack: bool = True, eps: float = -1) -> torch.Tensor:
+	""" Run sinkhorn iterations to generate a near doubly stochastic matrix, where each row or column sum to <=1
+
+	Args:
+		log_alpha: log of positive matrix to apply sinkhorn normalization (B, J, K)
+		n_iters (int): Number of normalization iterations
+		slack (bool): Whether to include slack row and column
+		eps: eps for early termination (Used only for handcrafted RPM). Set to negative to disable.
+
+	Returns:
+		log(perm_matrix): Doubly stochastic matrix (B, J, K)
+
+	Modified from original source taken from:
+		Learning Latent Permutations with Gumbel-Sinkhorn Networks
+		https://github.com/HeddaCohenIndelman/Learning-Gumbel-Sinkhorn-Permutations-w-Pytorch
+	"""
+
+	# Sinkhorn iterations
+	prev_alpha = None
+	if slack:
+		zero_pad = nn.ZeroPad2d((0, 1, 0, 1))
+		log_alpha_padded = zero_pad(log_alpha[:, None, :, :])
+
+		log_alpha_padded = torch.squeeze(log_alpha_padded, dim=1)
+
+		for i in range(n_iters):
+			# Row normalization
+			log_alpha_padded = torch.cat((
+					log_alpha_padded[:, :-1, :] - (torch.logsumexp(log_alpha_padded[:, :-1, :], dim=2, keepdim=True)),
+					log_alpha_padded[:, -1, None, :]),  # Don't normalize last row
+				dim=1)
+
+			# Column normalization
+			log_alpha_padded = torch.cat((
+					log_alpha_padded[:, :, :-1] - (torch.logsumexp(log_alpha_padded[:, :, :-1], dim=1, keepdim=True)),
+					log_alpha_padded[:, :, -1, None]),  # Don't normalize last column
+				dim=2)
+
+			if eps > 0:
+				if prev_alpha is not None:
+					abs_dev = torch.abs(torch.exp(log_alpha_padded[:, :-1, :-1]) - prev_alpha)
+					if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
+						break
+				prev_alpha = torch.exp(log_alpha_padded[:, :-1, :-1]).clone()
+
+		log_alpha = log_alpha_padded[:, :-1, :-1]
+	else:
+		for i in range(n_iters):
+			# Row normalization (i.e. each row sum to 1)
+			log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=2, keepdim=True))
+
+			# Column normalization (i.e. each column sum to 1)
+			log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=1, keepdim=True))
+
+			if eps > 0:
+				if prev_alpha is not None:
+					abs_dev = torch.abs(torch.exp(log_alpha) - prev_alpha)
+					if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
+						break
+				prev_alpha = torch.exp(log_alpha).clone()
+
+	return log_alpha
+
+
+def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
+	"""Compute rigid transforms between two point sets
+
+	Args:
+		a (torch.Tensor): (B, M, 3) points
+		b (torch.Tensor): (B, N, 3) points
+		weights (torch.Tensor): (B, M)
+
+	Returns:
+		Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
+	"""
+
+	weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
+	centroid_a = torch.sum(a * weights_normalized, dim=1) 
+	centroid_b = torch.sum(b * weights_normalized, dim=1) 
+	a_centered = a - centroid_a[:, None, :]  
+	b_centered = b - centroid_b[:, None, :]
+	cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)  
+
+	# Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
+	# and choose based on determinant to avoid flips
+	u, s, v = torch.svd(cov, some=False, compute_uv=True) 
+	rot_mat_pos = v @ u.transpose(-1, -2)  
+	v_neg = v.clone()
+	v_neg[:, :, 2] *= -1
+	rot_mat_neg = v_neg @ u.transpose(-1, -2)
+	rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
+	assert torch.all(torch.det(rot_mat) > 0)
+
+	# Compute translation (uncenter centroid)
+	translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]   
+
+	transform = torch.cat((rot_mat, translation), dim=2)  
+	return transform
+
+class R3PMNet(nn.Module):
+	def __init__(self, feature_model):
+
+		super().__init__()
+
+		self.add_slack = True
+		self.num_sk_iter = 5
+
+		self.weights_net = ParameterPredictionNet(weights_dim=[0])
+		self.feat_extractor = feature_model
+
+	def compute_affinity(self, beta, feat_distance, alpha=0.5):
+		"""Compute logarithm of Initial match matrix values, i.e. log(m_jk)"""
+		if isinstance(alpha, float):
+			hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha)
+		else:
+			hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha[:, None, None])
+		return hybrid_affinity
+
+	@staticmethod
+	def split_normals(data):
+		if data.shape[2] == 6:
+			xyz, normals = data[:, :, :3], data[:, :, 3:6]
+		elif data.shape[2] == 3:
+			xyz, normals = data, torch.zeros(data.shape).to(data.device)
+		return xyz, normals
+
+	def spam(self, xyz_template, norm_template, xyz_source, norm_source):
+		self.beta, self.alpha = self.weights_net([xyz_source, xyz_template])
+
+		try: # R3PMNET feature extractor
+			self.feat_source = self.feat_extractor(xyz_source)
+			self.feat_template = self.feat_extractor(xyz_template)
+		except: 
+			self.feat_source = self.feat_extractor(xyz_source, norm_source)
+			self.feat_template = self.feat_extractor(xyz_template, norm_template)
+
+		feat_distance = match_features(self.feat_source, self.feat_template)  
+		self.affinity = self.compute_affinity(self.beta, feat_distance, alpha=self.alpha) 
+
+		# Compute weighted coordinates
+		log_perm_matrix = sinkhorn(self.affinity, n_iters=self.num_sk_iter, slack=self.add_slack) 
+		self.perm_matrix = torch.exp(log_perm_matrix) 
+
+		try: # R3PMNET features
+			weighted_template = self.perm_matrix @ xyz_template[:,:self.perm_matrix.shape[1]] / (torch.sum(self.perm_matrix, dim=2, keepdim=True) + _EPS)
+		except: 
+			weighted_template = self.perm_matrix @ xyz_template / (torch.sum(self.perm_matrix, dim=2, keepdim=True) + _EPS) 
+		return weighted_template
+
+	def forward(self, template, source, max_iterations: int = 1):
+		"""Forward pass for R3PM-Net
+
+		Args:
+			data: Dict containing the following fields:
+					'points_src': Source points (B, J, 6)
+					'points_ref': Reference points (B, K, 6)
+			num_iter (int): Number of iterations. Recommended to be 2 for training
+
+		Returns:
+			transform: Transform to apply to source points such that they align to reference
+			src_transformed: Transformed source points
+		"""
+
+		xyz_template, norm_template = self.split_normals(template)
+		xyz_source, norm_source = self.split_normals(source)
+
+		xyz_source_t, norm_source_t = xyz_source, norm_source  # a copy of source to apply transformation to 
+
+		transforms = []
+		all_gamma, all_perm_matrices, all_weighted_template = [], [], []
+		all_beta, all_alpha = [], []
+
+		for i in range(max_iterations):
+			weighted_template = self.spam(xyz_template, norm_template, xyz_source_t, norm_source_t)		# Finding better correspondences after each iteration.
+			
+			# Compute transform and transform points
+			try: # R3PMNET features
+				transform = compute_rigid_transform(xyz_source[:,:weighted_template.shape[1]], weighted_template, weights=torch.sum(self.perm_matrix, dim=2))
+				xyz_source_t, norm_source_t = se3_transform(transform.detach(), xyz_source[:,:weighted_template.shape[1]], norm_source)			# Apply transformation to original source.	
+			except:
+				transform = compute_rigid_transform(xyz_source_t, weighted_template, weights=torch.sum(self.perm_matrix, dim=2))
+				xyz_source_t, norm_source_t = se3_transform(transform.detach(), xyz_source, norm_source)			# Apply transformation to original source.
+			
+
+			transforms.append(transform)
+			all_gamma.append(torch.exp(self.affinity))
+			all_perm_matrices.append(self.perm_matrix)
+			all_weighted_template.append(weighted_template)
+			all_beta.append(to_numpy(self.beta))
+			all_alpha.append(to_numpy(self.alpha))
+
+		est_T = convert2transformation(transforms[max_iterations-1][:, :3, :3], transforms[max_iterations-1][:, :3, 3])
+		transformed_source = torch.bmm(est_T[:, :3, :3], source[:,:,:3].permute(0, 2, 1)).permute(0, 2, 1) + est_T[:, :3, 3].unsqueeze(1)
+
+		try:  # for training
+			result = {'est_R': est_T[:, :3, :3],	# source -> template
+					'est_t': est_T[:, :3,  3],		# source -> template
+					'est_T': est_T,			# source -> template
+					'r': self.feat_template - self.feat_source,
+					'transformed_source': transformed_source}
+		except RuntimeError: 
+			result = {'est_R': est_T[:, :3, :3],	# source -> template
+					'est_t': est_T[:, :3,  3],		# source -> template
+					'est_T': est_T,			# source -> template
+					'transformed_source': transformed_source}
+			
+		result['perm_matrices_init'] = all_gamma
+		result['perm_matrices'] = all_perm_matrices
+		result['weighted_template'] = all_weighted_template
+		result['beta'] = np.stack(all_beta, axis=0)
+		result['alpha'] = np.stack(all_alpha, axis=0)
+		result['transforms'] = transforms
+
+		return result
+
+
+if __name__ == '__main__':
+	template, source = torch.rand(10,1024,6), torch.rand(10,1024,6)
+	
+	net = R3PMNet()
+	result = net(template, source)
+	import ipdb; ipdb.set_trace()

r3pm_net/paths.py

@@ -0,0 +1,11 @@
+"""Repository-root resolution for portable paths."""
+
+from pathlib import Path
+
+# r3pm_net/paths.py -> parents[1] is the repository root
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def repo_path(*parts: str) -> str:
+    """Join path segments relative to the repository root."""
+    return str(REPO_ROOT.joinpath(*parts))

scripts/eval_modelnet40.py

@@ -0,0 +1,335 @@
+import os
+import copy
+import sys
+from pathlib import Path
+from typing import Any
+
+# Repository root on PYTHONPATH (run: python scripts/test_modelnet40.py from repo root).
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+import argparse
+import random
+
+import numpy as np
+import open3d as o3d
+import torch
+from tqdm import tqdm
+
+from tools import augmentation, data, l3d_helper, print_results, transformations
+from tools import l3d_registration_and_evaluation, predator_registration_and_evaluation, geotransformer_registration_and_evaluation, logdesc_registration_and_evaluation, regtr_registration_and_evaluation
+from r3pm_net.config_loader import get_method_paths,get_modelnet40_paths, get_pretrained_rpmnet_dir
+
+'''
+This script evaluates the performance on the ModelNet40 test dataset.
+The results are averaged ovet the dataset with 2468 samples.
+All the point clouds are normalized to a sphere of radius 1.
+
+Augmentations:
+- Transformation = Random rotation (0 - 45) and translation (-0.5 to 0.5)
+- Noise = Gaussian noise with mean 0 and std deviation of 0.01 [optional]
+- Outliers = with level 1 which means 2% of the points are outliers (PC size = 2040) [optional]
+- Occlusion = 90000 radius which means 0.7% of the points are occluded (PC size = 1986) [optional]
+'''
+def set_seed(seed: int) -> None:
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" 
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(True)
+
+# arguments
+parser = argparse.ArgumentParser(description="ModelNet40 R3PM-Net evaluation")
+parser.add_argument("--seed", type=int, default=42, help="random seed (default: 42)")
+
+args = parser.parse_args()
+set_seed(args.seed)
+method_paths = get_method_paths()
+
+pretrained_base_dir = get_pretrained_rpmnet_dir()
+_path_zs = os.path.join(pretrained_base_dir, "clean-trained.pth")
+_path_ft = os.path.join(pretrained_base_dir, "best_model_PointNet.t7") #TODO: CHANGE 
+
+def fix_off_file(file_path):
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+    
+    if lines[0].startswith("OFF") and len(lines[0].strip().split()) > 1:
+        header = lines[0].strip()
+        new_header = "OFF\n" + header[3:] + "\n"
+        lines = [new_header] + lines[1:]
+        
+        with open(file_path, 'w') as f:
+            f.writelines(lines)
+        print(f"Fixed: {file_path}")
+
+def load_modelnet40_test_data(dataset_path, num_points=2000):
+    test_data = []
+    test_labels = []
+    categories = os.listdir(dataset_path)
+    for label, category in enumerate(tqdm(categories, desc="Loading Data")):
+        test_dir = os.path.join(dataset_path, category, 'test')
+        if not os.path.exists(test_dir):
+            continue
+        for file in tqdm(os.listdir(test_dir), desc=f"Processing {category} Category", leave=False):
+            if file.endswith('.off'):
+                file_path = os.path.join(test_dir, file)
+                mesh = o3d.io.read_triangle_mesh(file_path)
+                point_cloud = mesh.sample_points_poisson_disk(number_of_points=num_points)
+                test_data.append(point_cloud)
+                test_labels.append(label)
+    
+    return test_data, test_labels, categories
+
+# download from http://modelnet.cs.princeton.edu/ModelNet40.zip unzip and put the path in the config/eval.yaml
+dataset_path, save_dir = get_modelnet40_paths()
+test_data_path = os.path.join(save_dir, "test_data.npy")
+test_labels_path = os.path.join(save_dir, "test_labels.npy")
+categories_path = os.path.join(save_dir, "categories.npy")
+
+os.makedirs(save_dir, exist_ok=True)
+
+# Check if data already exists
+if os.path.exists(test_data_path) and os.path.exists(test_labels_path) and os.path.exists(categories_path):
+    print("Loading existing test data...")
+    test_data_np = np.load(test_data_path, allow_pickle=True)
+    test_labels = np.load(test_labels_path)
+    categories = np.load(categories_path)
+    print("Done! Testing the models...")
+else:
+    print("Loading and processing ModelNet40 test data...")
+    # Fix all .OFF files in the dataset
+    for root, _, files in os.walk(dataset_path):
+        for file in files:
+            if file.endswith(".off"):
+                fix_off_file(os.path.join(root, file))
+
+    test_data, test_labels, categories = load_modelnet40_test_data(dataset_path)
+
+    test_data_np = [data.normalize_pc(pc, return_as_np = True) for pc in test_data]
+
+    np.save(test_data_path, test_data_np)
+    np.save(test_labels_path, test_labels)
+    np.save(categories_path, categories)
+    print("Test data saved!")
+
+# Initialize arrays to store results
+rpm_results_all = []
+predator_results_all = []
+geotransformer_results_all = []
+logdesc_results_all = []
+regtr_results_all = []
+r3pm_net_results_all = []
+tuned_r3pm_net_results_all = []
+
+rpm_reg_results_all = []
+predator_reg_results_all = []
+geotransformer_reg_results_all = []
+logdesc_reg_results_all = []
+regtr_reg_results_all = []
+r3pm_net_reg_results_all = []
+tuned_r3pm_net_reg_results_all = []
+
+all_sources = []
+all_targets = []
+all_angles ={}
+
+# Reconstruct Open3D PointCloud objects from saved npy arrays
+test_data = [o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points)) for points in test_data_np]
+
+noise_level = 0
+outlier_level = 0
+outlier_lowerbound = -0.5
+outlier_upperbound = 0.5
+# occlusion_level = 90000 # Higher value means less occlusion
+occlusion_level = 0 # Higher value means less occlusion
+
+
+# set arguments for models
+rpm_args = l3d_helper.options(modelName="RPMNet")
+rpm_args.pretrained = _path_zs
+
+# OverlapPredator (used by Predator runner)
+predator_cfg = method_paths.get("predator", {})
+predator_root = predator_cfg.get("root")
+predator_config_path = predator_cfg.get("config_path")
+predator_weights_path = predator_cfg.get("weights_path")
+
+# GeoTransformer 
+geo_cfg = method_paths.get("geotransformer", {})
+geotransformer_root = geo_cfg.get("root")
+geotransformer_exp_subdir = geo_cfg.get("exp_subdir")
+geotransformer_weights_path = geo_cfg.get("weights_path")
+
+# LoGDesc
+logdesc_cfg = method_paths.get("logdesc", {})
+logdesc_root = logdesc_cfg.get("root")
+logdesc_weights_path = logdesc_cfg.get("weights_path")
+
+# RegTR
+regtr_cfg = method_paths.get("regtr", {})
+regtr_root = regtr_cfg.get("root")
+regtr_ckpt_path = regtr_cfg.get("ckpt_path")
+regtr_config_path = regtr_cfg.get("config_path")
+
+# R3PM-Net (ours) - ZS - no training
+r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+r3pm_net_args.pretrained = _path_zs
+
+# R3PM-Net (ours) - FT
+tuned_r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+tuned_r3pm_net_args.pretrained = _path_ft
+
+for i, item in enumerate(tqdm(test_data, desc="Testing methods")):
+    
+   # Simulate data
+    x_angle = int(random.uniform(0, 45))
+    y_angle = int(random.uniform(0, 45))
+    z_angle = int(random.uniform(0, 45))
+    translation_range = (-0.5, 0.5)
+    gt_transformation = transformations.create_transformation(x_angle, y_angle, z_angle, translation_range)
+    source = copy.deepcopy(item)
+
+    target = copy.deepcopy(item).transform(gt_transformation)
+
+    # Apply augmentations
+    noisy_source = copy.deepcopy(source)    
+    if noise_level != 0:
+        noisy_source = augmentation.apply_noise(noisy_source, noise_level)
+    if outlier_level != 0:
+        noisy_source = augmentation.add_outliers(noisy_source, outlier_level, outlier_lowerbound, outlier_upperbound)
+    if occlusion_level != 0:
+        noisy_source, _ = augmentation.apply_occlusion(noisy_source, occlusion_level) 
+    if len(noisy_source.points) < 1024: # cannot be smaller than embedding dims in config/default.yaml
+        noisy_source = copy.deepcopy(source)
+        noisy_source = augmentation.apply_noise(noisy_source, noise_level)
+        noisy_source, _ = augmentation.apply_occlusion(noisy_source, occlusion_level * 100) 
+    assert len(noisy_source.points) >= 1024, "Noisy source point cloud has less than 1024 points."
+    
+    # RPMNet
+    rpm_results_pc, rpm_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        noisy_source, target, 'rpmnet', gt_transformation, rpm_args)
+    rpm_results_all.append(rpm_results)
+    rpm_reg_results_all.append(rpm_results_pc)
+
+    # OverlapPredator
+    predator_results_pc, predator_results = predator_registration_and_evaluation.predator_reg_and_eval(
+        noisy_source,
+        target,
+        gt_transformation=gt_transformation,
+        predator_root=predator_root,
+        config_path=predator_config_path,
+        weights_path=predator_weights_path,
+        ransac_n_points=1000,
+        ransac_distance_threshold=0.05,
+        ransac_n=3,
+        sampling="prob",
+        mutual=False,
+        input_num_points=1024,
+    )
+    predator_results_all.append(predator_results)
+    predator_reg_results_all.append(predator_results_pc)
+
+    # GeoTransformer (ModelNet)
+    geotransformer_results_pc, geotransformer_results = geotransformer_registration_and_evaluation.geotransformer_reg_and_eval(
+        noisy_source,
+        target,
+        gt_transformation=gt_transformation,
+        geotransformer_root=geotransformer_root,
+        exp_subdir=geotransformer_exp_subdir,
+        weights_path=geotransformer_weights_path,
+    )
+    geotransformer_results_all.append(geotransformer_results)
+    geotransformer_reg_results_all.append(geotransformer_results_pc)
+
+    # LoGDesc
+    logdesc_results_pc, logdesc_results = logdesc_registration_and_evaluation.logdesc_reg_and_eval(
+        noisy_source,
+        target,
+        gt_transformation=gt_transformation,
+        logdesc_root=logdesc_root,
+        weights_path=logdesc_weights_path,
+        max_keypoints=768,
+        num_points_per_sample=128,
+        sample_radius=0.3,
+        topk_matches=128,
+        use_kpt=False,
+    )
+    logdesc_results_all.append(logdesc_results)
+    logdesc_reg_results_all.append(logdesc_results_pc)
+
+    # RegTR (ModelNet)
+    regtr_results_pc, regtr_results = regtr_registration_and_evaluation.regtr_reg_and_eval(
+        noisy_source,
+        target,
+        gt_transformation=gt_transformation,
+        regtr_root=regtr_root,
+        ckpt_path=regtr_ckpt_path,
+        config_path=regtr_config_path,
+    )
+    regtr_results_all.append(regtr_results)
+    regtr_reg_results_all.append(regtr_results_pc)
+
+    # R3PM-Net (ours) - no training
+    r3pm_net_results_pc, r3pm_net_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        noisy_source, target, 'r3pmnet', gt_transformation, r3pm_net_args)
+    r3pm_net_results_all.append(r3pm_net_results)
+    r3pm_net_reg_results_all.append(r3pm_net_results_pc)
+
+    # R3PM-Net (ours) (Tuned on 4 sioux data)
+    tuned_r3pm_net_results_pc, tuned_r3pm_net_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        noisy_source, target, 'r3pmnet', gt_transformation, tuned_r3pm_net_args)
+    tuned_r3pm_net_results_all.append(tuned_r3pm_net_results)
+    tuned_r3pm_net_reg_results_all.append(tuned_r3pm_net_results_pc)
+
+
+    all_sources.append(noisy_source)
+    all_targets.append(target)
+    all_angles[i] = {
+        "x_angle": x_angle,
+        "y_angle": y_angle,
+        "z_angle": z_angle,
+        "translation": gt_transformation[:3, 3]
+    }
+    
+# Convert results to numpy arrays for easier manipulation
+rpm_results_all = np.array(rpm_results_all)
+predator_results_all = np.array(predator_results_all)
+geotransformer_results_all = np.array(geotransformer_results_all)
+logdesc_results_all = np.array(logdesc_results_all)
+regtr_results_all = np.array(regtr_results_all)
+r3pm_net_results_all = np.array(r3pm_net_results_all)
+tuned_r3pm_net_results_all = np.array(tuned_r3pm_net_results_all)
+
+rpm_mean_results = np.mean(rpm_results_all, axis=0)
+predator_mean_results = np.mean(predator_results_all, axis=0)
+geotransformer_mean_results = np.mean(geotransformer_results_all, axis=0)
+logdesc_mean_results = np.mean(logdesc_results_all, axis=0)
+regtr_mean_results = np.mean(regtr_results_all, axis=0)
+r3pm_net_mean_results = np.mean(r3pm_net_results_all, axis=0)
+tuned_r3pm_net_mean_results = np.mean(tuned_r3pm_net_results_all, axis=0)
+
+# Print the results
+metric_names = ['mean_rmse', 'mean_rotation_error', 'mean_translation_error',
+                'mean_computation_time', 'mean_cd', 'mean_error',
+                'mean_fitness', 'mean_inlier_rmse']
+
+reports = {
+    "RPMNet": dict(zip(metric_names, rpm_mean_results)),
+    "Predator": dict(zip(metric_names, predator_mean_results)),
+    "GeoTransformer": dict(zip(metric_names, geotransformer_mean_results)),
+    "LoGDesc": dict(zip(metric_names, logdesc_mean_results)),
+    "RegTR": dict(zip(metric_names, regtr_mean_results)),
+    "R3PM-Net (ours) (ZS)": dict(zip(metric_names, r3pm_net_mean_results)),
+    "R3PM-Net (ours) (FT)": dict(zip(metric_names, tuned_r3pm_net_mean_results)),
+}
+
+# Print the table
+print_results.print_table(reports)

scripts/eval_sioux_cranfield.py

@@ -0,0 +1,302 @@
+import os
+import copy
+import open3d as o3d
+import numpy as np
+from tqdm import tqdm
+import sys
+from pathlib import Path
+import torch
+import random
+import argparse
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tools import augmentation, data, l3d_helper, print_results, transformations
+from tools import l3d_registration_and_evaluation, predator_registration_and_evaluation, geotransformer_registration_and_evaluation, logdesc_registration_and_evaluation, regtr_registration_and_evaluation
+from r3pm_net.config_loader import get_method_paths, get_pretrained_rpmnet_dir, get_sioux_data_root, get_sioux_paths
+'''
+This script evaluates the performance on a Sioux-Cranfield dataset 
+Cranfield dataset from: https://github.com/Menthy-Denayer/PCR_CAD_Model_Alignment_Comparison/tree/main/datasets
+'''
+def set_seed(seed: int) -> None:
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" 
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(True)
+
+# arguments
+parser = argparse.ArgumentParser(description="Sioux-Cranfield R3PM-Net evaluation")
+parser.add_argument("--seed", type=int, default=42, help="random seed (default: 42)")
+args = parser.parse_args()
+set_seed(args.seed)
+
+base_dir = get_sioux_data_root()
+sioux_cfg = get_sioux_paths()
+method_paths = get_method_paths()
+
+pretrained_base_dir = get_pretrained_rpmnet_dir()
+_path_zs = os.path.join(pretrained_base_dir, "clean-trained.pth")
+_path_ft = os.path.join(pretrained_base_dir, "best_model_PointNet.t7") #TODO: CHANGE 
+
+# Paths to the CAD models
+cad_dir_made = os.path.join(base_dir, 'sioux_cranfield') 
+
+cad_paths = [os.path.join(cad_dir_made, 'Base-Top_Plate.stl'),
+             os.path.join(cad_dir_made, 'Pendulum.stl'),
+             os.path.join(cad_dir_made, 'Round-Peg.stl'),
+             os.path.join(cad_dir_made, 'Separator.stl'),
+             os.path.join(cad_dir_made, 'Shaft-New.stl'),
+             os.path.join(cad_dir_made, 'Square-Peg.stl'),
+             os.path.join(cad_dir_made, 'elephant.stl'),
+             os.path.join(cad_dir_made, 'house.stl'),
+             os.path.join(cad_dir_made, 'shoe.stl')]
+
+# Test parameters
+num_tests = 25
+angles = list(range(0, 45))
+translation_range = (-0.5, 0.5)
+np.random.seed(42)
+
+# Augmentation parameters
+noise_level = 0
+outlier_level = 0
+outlier_lowerbound = -0.5
+outlier_upperbound = 0.5
+# occlusion_level = 9000 # Higher value means less occlusion
+occ_level = 0
+
+# Make dataset
+sources = []
+targets = []
+x_angles = []
+y_angles = []
+z_angles = []
+gt_transformations = []
+
+for cadPath in tqdm (cad_paths, desc="Preparing Sioux-Cranfield Dataset", total=len(cad_paths)):
+
+        num_points = 2000
+        # Load the data
+        mesh = o3d.io.read_triangle_mesh(cadPath)
+        cad = mesh.sample_points_poisson_disk(number_of_points=num_points)  # modify to a suitable number of points
+        normalized_point_cloud = data.normalize_pc(cad)
+        source = copy.deepcopy(normalized_point_cloud)
+
+        for test in range(num_tests):
+            # Data simulation
+            x_angle= np.random.uniform(angles[0], angles[-1], size=1)
+            y_angle= np.random.uniform(angles[0], angles[-1], size=1)
+            z_angle= np.random.uniform(angles[0], angles[-1], size=1)
+            gt_transformation = transformations.create_transformation(x_angle, y_angle, z_angle, translation_range)
+            target = copy.deepcopy(normalized_point_cloud).transform(gt_transformation)
+
+            # Data augmentation
+            if occ_level == 0 and noise_level == 0 and outlier_level == 0:
+                noisy_source = copy.deepcopy(source)
+               
+            # Noise + Occlusion
+            elif occ_level != 0 and noise_level != 0:
+                noisy_source_noise = augmentation.apply_noise(source, noise_level)
+                noisy_source, _ = augmentation.apply_occlusion(noisy_source_noise, occ_level)
+                if len(noisy_source.points) < 1024:  # Handle excessive occlusion
+                    source = copy.deepcopy(target).transform(gt_transformation)
+                    noisy_source_noise = augmentation.apply_noise(source, noise_level)
+                    noisy_source, _ = augmentation.apply_occlusion(noisy_source_noise, occ_level * 1.5)
+
+            # Noise + Outlier
+            elif noise_level != 0 and outlier_level != 0:
+                noisy_source_noise = augmentation.apply_noise(source, noise_level)
+                noisy_source = augmentation.add_outliers(noisy_source_noise, outlier_level, outlier_lowerbound=-0.5, outlier_upperbound=0.5)
+
+            # Noise + Outlier + Occlusion
+            elif occ_level != 0 and noise_level != 0 and outlier_level != 0:       
+                noisy_source_noise = augmentation.apply_noise(source, noise_level)
+                noisy_source, _ = augmentation.apply_occlusion(noisy_source_noise, occ_level)
+                if len(noisy_source.points) < 1024:  # Handle excessive occlusion
+                    source = copy.deepcopy(target).transform(gt_transformation)
+                    noisy_source_noise = augmentation.apply_noise(source, noise_level)
+                    noisy_source, _ = augmentation.apply_occlusion(noisy_source_noise, occ_level * 1.5)
+                noisy_source = augmentation.add_outliers(noisy_source, outlier_level, outlier_lowerbound=-0.5, outlier_upperbound=0.5)
+
+            # collect dataset in lists
+            sources.append(noisy_source)
+            targets.append(target)
+            x_angles.append(x_angle)
+            y_angles.append(y_angle)
+            z_angles.append(z_angle)
+            gt_transformations.append(gt_transformation)
+
+# Initialize arrays to store results
+rpm_results_all = []
+predator_results_all = []
+geotransformer_results_all = []
+logdesc_results_all = []
+regtr_results_all = []
+r3pm_net_results_all = []
+tuned_r3pm_net_results_all = []
+
+rpm_reg_results_all = []
+predator_reg_results_all = []
+geotransformer_reg_results_all = []
+logdesc_reg_results_all = []
+regtr_reg_results_all = []
+r3pm_net_reg_results_all = []
+tuned_r3pm_net_reg_results_all = []
+
+# set arguments for models
+rpm_args = l3d_helper.options(modelName="RPMNet")
+rpm_args.pretrained = _path_zs
+
+# OverlapPredator (used by Predator runner)
+predator_cfg = method_paths.get("predator", {})
+predator_root = predator_cfg.get("root")
+predator_config_path = predator_cfg.get("config_path")
+predator_weights_path = predator_cfg.get("weights_path")
+
+# GeoTransformer 
+geo_cfg = method_paths.get("geotransformer", {})
+geotransformer_root = geo_cfg.get("root")
+geotransformer_exp_subdir = geo_cfg.get("exp_subdir")
+geotransformer_weights_path = geo_cfg.get("weights_path")
+
+# LoGDesc
+logdesc_cfg = method_paths.get("logdesc", {})
+logdesc_root = logdesc_cfg.get("root")
+logdesc_weights_path = logdesc_cfg.get("weights_path")
+
+# RegTR
+regtr_cfg = method_paths.get("regtr", {})
+regtr_root = regtr_cfg.get("root")
+regtr_ckpt_path = regtr_cfg.get("ckpt_path")
+regtr_config_path = regtr_cfg.get("config_path")
+
+# R3PM-Net (ours) - ZS - no training
+r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+r3pm_net_args.pretrained = _path_zs
+
+# R3PM-Net (ours) - FT
+tuned_r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+tuned_r3pm_net_args.pretrained = _path_ft
+
+
+for i, item in enumerate(tqdm(zip(sources, targets, gt_transformations), desc="Testing methods", total=len(sources))):
+    
+    # RPMNet
+    rpm_results_pc, rpm_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        sources[i], targets[i], 'rpmnet', gt_transformations[i], rpm_args)
+    rpm_results_all.append(rpm_results)
+    rpm_reg_results_all.append(rpm_results_pc)
+
+    # OverlapPredator
+    predator_results_pc, predator_results = predator_registration_and_evaluation.predator_reg_and_eval(
+        sources[i],
+        targets[i],
+        gt_transformation=gt_transformations[i],
+        predator_root=predator_root,
+        config_path=predator_config_path,
+        weights_path=predator_weights_path,
+        ransac_n_points=1000,
+        ransac_distance_threshold=0.05,
+        ransac_n=3,
+        sampling="prob",
+        mutual=False,
+        input_num_points=1024,
+    )
+    predator_results_all.append(predator_results)
+    predator_reg_results_all.append(predator_results_pc)
+
+    # GeoTransformer (ModelNet)
+    geotransformer_results_pc, geotransformer_results = geotransformer_registration_and_evaluation.geotransformer_reg_and_eval(
+        sources[i],
+        targets[i],
+        gt_transformation=gt_transformations[i],
+        geotransformer_root=geotransformer_root,
+        exp_subdir=geotransformer_exp_subdir,
+        weights_path=geotransformer_weights_path,
+    )
+    geotransformer_results_all.append(geotransformer_results)
+    geotransformer_reg_results_all.append(geotransformer_results_pc)
+
+        # LoGDesc
+    logdesc_results_pc, logdesc_results = logdesc_registration_and_evaluation.logdesc_reg_and_eval(
+        sources[i],
+        targets[i],
+        gt_transformation=gt_transformations[i],
+        logdesc_root=logdesc_root,
+        weights_path=logdesc_weights_path,
+        max_keypoints=768,
+        num_points_per_sample=128,
+        sample_radius=0.3,
+        topk_matches=128,
+        use_kpt=False,
+    )
+    logdesc_results_all.append(logdesc_results)
+    logdesc_reg_results_all.append(logdesc_results_pc)
+
+    # RegTR (ModelNet)
+    regtr_results_pc, regtr_results = regtr_registration_and_evaluation.regtr_reg_and_eval(
+        sources[i],
+        targets[i],
+        gt_transformation=gt_transformations[i],
+        regtr_root=regtr_root,
+        ckpt_path=regtr_ckpt_path,
+        config_path=regtr_config_path,
+    )
+    regtr_results_all.append(regtr_results)
+    regtr_reg_results_all.append(regtr_results_pc)
+
+    # R3PM-Net (ours) - ZS - no training
+    r3pm_net_results_pc, r3pm_net_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        sources[i], targets[i], 'r3pmnet', gt_transformations[i], r3pm_net_args)
+    r3pm_net_results_all.append(r3pm_net_results)
+    r3pm_net_reg_results_all.append(r3pm_net_results_pc)
+
+    # R3PM-Net (ours) - FT
+    tuned_r3pm_net_results_pc, tuned_r3pm_net_results = l3d_registration_and_evaluation.l3d_reg_and_eval(
+        sources[i], targets[i], 'r3pmnet', gt_transformations[i], tuned_r3pm_net_args)
+    tuned_r3pm_net_results_all.append(tuned_r3pm_net_results)
+    tuned_r3pm_net_reg_results_all.append(tuned_r3pm_net_results_pc)
+
+ 
+# Convert results to numpy arrays for easier manipulation
+rpm_results_all = np.array(rpm_results_all)
+predator_results_all = np.array(predator_results_all)
+geotransformer_results_all = np.array(geotransformer_results_all)
+logdesc_results_all = np.array(logdesc_results_all)
+regtr_results_all = np.array(regtr_results_all)
+r3pm_net_results_all = np.array(r3pm_net_results_all)
+tuned_r3pm_net_results_all = np.array(tuned_r3pm_net_results_all)
+
+rpm_mean_results = np.mean(rpm_results_all, axis=0)
+predator_mean_results = np.mean(predator_results_all, axis=0)
+geotransformer_mean_results = np.mean(geotransformer_results_all, axis=0)
+logdesc_mean_results = np.mean(logdesc_results_all, axis=0)
+regtr_mean_results = np.mean(regtr_results_all, axis=0)
+r3pm_net_mean_results = np.mean(r3pm_net_results_all, axis=0)
+tuned_r3pm_net_mean_results = np.mean(tuned_r3pm_net_results_all, axis=0)
+
+# Print the results
+metric_names = ['mean_rmse', 'mean_rotation_error', 'mean_translation_error',
+                'mean_computation_time', 'mean_cd', 'mean_error',
+                'mean_fitness', 'mean_inlier_rmse']
+
+reports = {
+    "RPMNet": dict(zip(metric_names, rpm_mean_results)),
+    "Predator": dict(zip(metric_names, predator_mean_results)),
+    "GeoTransformer": dict(zip(metric_names, geotransformer_mean_results)),
+    "LoGDesc": dict(zip(metric_names, logdesc_mean_results)),
+    "RegTR": dict(zip(metric_names, regtr_mean_results)),
+    "R3PM-Net (ours) (ZS)": dict(zip(metric_names, r3pm_net_mean_results)),
+    "R3PM-Net (ours) (FT)": dict(zip(metric_names, tuned_r3pm_net_mean_results)),}
+
+# Print the table
+print_results.print_table(reports)

scripts/eval_sioux_scans.py

@@ -0,0 +1,341 @@
+import os
+import copy
+import argparse
+import numpy as np
+import random
+import torch
+from tabulate import tabulate
+from tqdm import tqdm
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tools import data, l3d_helper, visualization
+from tools import icp_registration_and_evaluation, l3d_registration_and_evaluation, predator_registration_and_evaluation, geotransformer_registration_and_evaluation, logdesc_registration_and_evaluation, regtr_registration_and_evaluation
+from r3pm_net.config_loader import get_pretrained_rpmnet_dir, get_sioux_data_root, get_method_paths
+
+'''
+This script is used to evaluate the performance of the pipeline with R3PM-Net as global and GICP as local registeration.
+
+The script takes the following arguments:
+--local_reg: the local registration method to be used.  
+--seed: random seed for python/numpy/torch. The default is 42.
+--verbose: if set to True, the results will be printed in a table format. The default is False.
+'''
+def set_seed(seed: int) -> None:
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(True)
+
+
+# arguments
+parser = argparse.ArgumentParser(description="Choosing local registration method")
+parser.add_argument(
+    "--local_reg", type=str, default="gicp", help="local registration: gicp or freg"
+)
+parser.add_argument("--seed", type=int, default=42, help="random seed (default: 42)")
+
+args = parser.parse_args()
+set_seed(args.seed)
+print(f"Using {args.local_reg} for local registration")
+
+def analyze_results(results: dict, recall_threshold = 1, rmse_threshold = 0.053, verbose = False):  # change the default values to your needs
+    table = []
+    fail_count = 0
+    success_count = 0
+    for object, values in results.items():
+        row = [object] + list(values)
+        if round(row[2], 3) < recall_threshold or round(row[3], 3) > rmse_threshold:
+            status = 'failed'
+            fail_count += 1
+            print(f'No match for {object}! Try a different method. If the issue persists, please check the data.')
+        else:
+            status = 'success'
+            success_count += 1
+            print(f'Found match for {object}!')
+        row.append(status)
+        table.append(row)
+
+    if verbose:
+        print(tabulate(table, headers=['Object', 'Chamfer Distance', 'Reg. Recall', 'Inlier RMSE', 'Computation Time', 'Status'], tablefmt='grid'))
+        print(f"Success rate: {success_count / (success_count + fail_count) * 100:.2f}%")
+
+    return table
+
+def show_successful_resutls(table, sources, targets, pc_results, method_name = None):
+    for i in range (len(table)):
+        if table[i][-1] == 'success':
+            # visualization.plot_point_cloud(sources[i], targets[i], list(pc_results.values())[i])   # uncomment if below visualization does not work 
+            visualization.draw_registration_result(targets[i], list(pc_results.values())[i], np.eye(4), method_name)
+
+def main():
+    base_dir = get_sioux_data_root()
+    scan_dir = os.path.join(base_dir, 'sioux_scans')
+    cad_dir = os.path.join(base_dir, 'sioux_cranfield')
+
+    pcd_paths = [ os.path.join(scan_dir,'teeth_clean.ply'),
+                    os.path.join(scan_dir,'lime_clean.ply'),
+                    os.path.join(scan_dir,'cube_clean.ply'),  
+                    os.path.join(scan_dir,'lego_clean.ply'),
+                    os.path.join(scan_dir,'elephant_clean.ply'), 
+                    os.path.join(scan_dir,'house_clean.ply'),
+                    os.path.join(scan_dir,'shoe_clean.ply')]
+
+    cad_paths = [ os.path.join(cad_dir,'teeth.stl'),
+                    os.path.join(cad_dir,'lime.stl'),
+                    os.path.join(cad_dir,'cube.stl'),  
+                    os.path.join(cad_dir,'lego.stl'),
+                    os.path.join(cad_dir,'elephant.stl'), 
+                    os.path.join(cad_dir,'house.stl'),
+                    os.path.join(cad_dir,'shoe.stl')]
+
+    # Initialize lists and dictionaries to store results
+    rpm_net_results = {}
+    rpm_net_pc_results = {}
+    predator_results = {}
+    predator_pc_results = {}
+    geotransformer_results = {}
+    geotransformer_pc_results = {}
+    logdesc_results = {}
+    logdesc_pc_results = {}
+    regtr_results = {}
+    regtr_pc_results = {}
+    r3pm_net_results = {}
+    r3pm_net_pc_results ={}
+    tuned_r3pm_net_results = {}
+    tuned_r3pm_net_pc_results = {}
+    subset_tuned_r3pm_net_results = {}
+    subset_tuned_r3pm_net_pc_results = {}
+
+    sources = []
+    targets = []
+
+    pretrained_base_dir = get_pretrained_rpmnet_dir()
+    method_paths = get_method_paths()
+    _path_zs = os.path.join(pretrained_base_dir, "clean-trained.pth")
+    _path_ft = os.path.join(pretrained_base_dir, "best_model_PointNet2.t7") #TODO: CHANGE 
+    _path_ft_sub = os.path.join(pretrained_base_dir, "best_model_PointNet_subset.t7") #TODO: CHANGE 
+
+    # set arguments for models
+    rpm_args = l3d_helper.options(modelName="RPMNet")
+    rpm_args.pretrained = _path_zs
+
+    # OverlapPredator (used by Predator runner)
+    predator_cfg = method_paths.get("predator", {})
+    predator_root = predator_cfg.get("root")
+    predator_config_path = predator_cfg.get("config_path")
+    predator_weights_path = predator_cfg.get("weights_path")
+
+    # GeoTransformer 
+    geo_cfg = method_paths.get("geotransformer", {})
+    geotransformer_root = geo_cfg.get("root")
+    geotransformer_exp_subdir = geo_cfg.get("exp_subdir")
+    geotransformer_weights_path = geo_cfg.get("weights_path")
+
+    # LoGDesc
+    logdesc_cfg = method_paths.get("logdesc", {})
+    logdesc_root = logdesc_cfg.get("root")
+    logdesc_weights_path = logdesc_cfg.get("weights_path")
+
+    # RegTR
+    regtr_cfg = method_paths.get("regtr", {})
+    regtr_root = regtr_cfg.get("root")
+    regtr_ckpt_path = regtr_cfg.get("ckpt_path")
+    regtr_config_path = regtr_cfg.get("config_path")
+
+    # R3PM-Net (ours) - no training
+    r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+    r3pm_net_args.pretrained = _path_zs
+
+    # R3PM-Net (ours) (FT)
+    tuned_r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+    tuned_r3pm_net_args.pretrained = _path_ft
+
+    # R3PM-Net (ours) (FT) (Subset)
+    subset_tuned_r3pm_net_args = l3d_helper.options(modelName="R3PMNet")
+    subset_tuned_r3pm_net_args.pretrained = _path_ft_sub
+
+    for pcdPath, cadPath in tqdm(zip(pcd_paths, cad_paths), desc="Registering objects", total=len(pcd_paths)):
+        # Define the number of points to sample from the CAD model (change this based on your data)
+        if 'teeth' in pcdPath:
+            every_k_points = 100
+            key = 'teeth'
+        elif'lime' in pcdPath:
+            every_k_points = 100
+            key = 'lime'
+        elif 'cube' in pcdPath:
+            every_k_points = 1
+            key = 'cube'
+        elif 'lego' in pcdPath:
+            every_k_points = 10
+            key = 'lego'
+        elif 'elephant' in pcdPath:
+            every_k_points = 30
+            key = 'elephant'
+        elif 'house' in pcdPath:
+            every_k_points = 25
+            key = 'house'
+        elif 'shoe' in pcdPath:
+            every_k_points = 15
+            key = 'shoe'
+        else:
+            print("Unknown object type, using default every_k_points = 1")
+            every_k_points = 1
+        
+        # Load the data
+        pcd, cad = data.load_data(pcdPath, cadPath, every_k_points=every_k_points)
+        source = copy.deepcopy(pcd)
+        target = copy.deepcopy(cad)
+        
+        # Normalize the point clouds
+        source = data.normalize_pc(source)
+        target = data.normalize_pc(target)
+
+        sources.append(source)
+        targets.append(target)
+
+        gt_transformation = None
+
+        # Perform the registration
+
+        # RPMNet
+        rpm_pc_result, _ = l3d_registration_and_evaluation.l3d_reg_and_eval(
+            source, target, 'rpmnet', gt_transformation, rpm_args)
+        if args.local_reg == 'gicp':
+            final_rpm_net_pc_result, final_rpm_net_results = icp_registration_and_evaluation.icp_reg_and_eval(rpm_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        rpm_net_results[key] = final_rpm_net_results
+        rpm_net_pc_results[key] = final_rpm_net_pc_result
+
+        # OverlapPredator
+        predator_results_pc, _ = predator_registration_and_evaluation.predator_reg_and_eval(
+            source,
+            target,
+            gt_transformation=gt_transformation,
+            predator_root=predator_root,
+            config_path=predator_config_path,
+            weights_path=predator_weights_path,
+            ransac_n_points=1000,
+            ransac_distance_threshold=0.05,
+            ransac_n=3,
+            sampling="prob",
+            mutual=False,
+            input_num_points=1024,
+        )
+        if args.local_reg == 'gicp':
+            final_predator_pc_result, final_predator_results = icp_registration_and_evaluation.icp_reg_and_eval(predator_results_pc, target, 'gicp', 1, np.identity(4), gt_transformation)
+        predator_results[key] = final_predator_results
+        predator_pc_results[key] = final_predator_pc_result
+
+        # GeoTransformer (ModelNet)
+        geotransformer_pc_result, _ = geotransformer_registration_and_evaluation.geotransformer_reg_and_eval(
+            source,
+            target,
+            gt_transformation=gt_transformation,
+            geotransformer_root=geotransformer_root,
+            exp_subdir=geotransformer_exp_subdir,
+            weights_path=geotransformer_weights_path,
+        )
+        if args.local_reg == 'gicp':
+            final_geotransformer_pc_result, final_geotransformer_results = icp_registration_and_evaluation.icp_reg_and_eval(geotransformer_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        geotransformer_results[key] = final_geotransformer_results
+        geotransformer_pc_results[key] = final_geotransformer_pc_result
+
+        # LoGDesc
+        logdesc_pc_result, _ = logdesc_registration_and_evaluation.logdesc_reg_and_eval(
+            source,
+            target,
+            gt_transformation=gt_transformation,
+            logdesc_root=logdesc_root,
+            weights_path=logdesc_weights_path,
+            max_keypoints=768,
+            num_points_per_sample=128,
+            sample_radius=0.3,
+            topk_matches=128,
+            use_kpt=False,
+        )
+        if args.local_reg == 'gicp':
+            final_logdesc_pc_result, final_logdesc_results = icp_registration_and_evaluation.icp_reg_and_eval(logdesc_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        logdesc_results[key] = final_logdesc_results
+        logdesc_pc_results[key] = final_logdesc_pc_result
+
+        # RegTR (ModelNet)
+        regtr_pc_result, _ = regtr_registration_and_evaluation.regtr_reg_and_eval(
+            source,
+            target,
+            gt_transformation=gt_transformation,
+            regtr_root=regtr_root,
+            ckpt_path=regtr_ckpt_path,
+            config_path=regtr_config_path,
+        )
+        if args.local_reg == 'gicp':
+            final_regtr_pc_result, final_regtr_results = icp_registration_and_evaluation.icp_reg_and_eval(regtr_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        regtr_results[key] = final_regtr_results
+        regtr_pc_results[key] = final_regtr_pc_result
+
+        # R3PM-Net (ours) (ZS)
+        r3pm_net_pc_result, _ = l3d_registration_and_evaluation.l3d_reg_and_eval(source, target, 'r3pmnet', gt_transformation, r3pm_net_args)
+        if args.local_reg == 'gicp':
+            final_r3pm_net_pc_result, final_r3pm_net_results = icp_registration_and_evaluation.icp_reg_and_eval(r3pm_net_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        r3pm_net_results[key] = final_r3pm_net_results
+        r3pm_net_pc_results[key] = final_r3pm_net_pc_result
+
+        # R3PM-Net (ours) (FT)
+        tuned_r3pm_net_pc_result, _ = l3d_registration_and_evaluation.l3d_reg_and_eval(source, target, 'r3pmnet', gt_transformation, tuned_r3pm_net_args)
+        if args.local_reg == 'gicp':
+            final_tuned_r3pm_net_pc_result, final_tuned_r3pm_net_results = icp_registration_and_evaluation.icp_reg_and_eval(tuned_r3pm_net_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        tuned_r3pm_net_results[key] = final_tuned_r3pm_net_results
+        tuned_r3pm_net_pc_results[key] = final_tuned_r3pm_net_pc_result
+
+        # R3PM-Net (ours) (FT) (Subset)
+        subset_tuned_r3pm_net_pc_result, _ = l3d_registration_and_evaluation.l3d_reg_and_eval(source, target, 'r3pmnet', gt_transformation, subset_tuned_r3pm_net_args)
+        if args.local_reg == 'gicp':
+            final_subset_tuned_r3pm_net_pc_result, final_subset_tuned_r3pm_net_results = icp_registration_and_evaluation.icp_reg_and_eval(subset_tuned_r3pm_net_pc_result, target, 'gicp', 1, np.identity(4), gt_transformation)
+        subset_tuned_r3pm_net_results[key] = final_subset_tuned_r3pm_net_results
+        subset_tuned_r3pm_net_pc_results[key] = final_subset_tuned_r3pm_net_pc_result
+
+    # Print the results
+    print("----- RPMNet: -----")
+    rpm_net_table = analyze_results(rpm_net_results, verbose=True)
+    show_successful_resutls(rpm_net_table, sources, targets, rpm_net_pc_results, 'RPMNet')
+
+    print("----- Predator: -----")
+    predator_table = analyze_results(predator_results, verbose=True)
+    show_successful_resutls(predator_table, sources, targets, predator_pc_results, 'Predator')
+
+    print("----- GeoTransformer: -----")
+    geotransformer_table = analyze_results(geotransformer_results, verbose=True)
+    show_successful_resutls(geotransformer_table, sources, targets, geotransformer_pc_results, 'GeoTransformer')
+
+    print("----- LoGDesc: -----")
+    logdesc_table = analyze_results(logdesc_results, verbose=True)
+    show_successful_resutls(logdesc_table, sources, targets, logdesc_pc_results, 'LoGDesc')
+
+    print("----- RegTR: -----")
+    regtr_table = analyze_results(regtr_results, verbose=True)
+    show_successful_resutls(regtr_table, sources, targets, regtr_pc_results, 'RegTR')
+
+    print("----- R3PM-Net (ours) (ZS): -----")
+    r3pm_net_table = analyze_results(r3pm_net_results, verbose=True)
+    show_successful_resutls(r3pm_net_table, sources, targets, r3pm_net_pc_results, 'R3PM-Net (ours) (ZS)')
+
+    print("----- R3PM-Net (ours) (FT): ----- ")
+    tuned_r3pm_net_table = analyze_results(tuned_r3pm_net_results, verbose=True)
+    show_successful_resutls(tuned_r3pm_net_table, sources, targets, tuned_r3pm_net_pc_results, 'R3PM-Net (ours) (FT)')
+
+    print("----- R3PM-Net (ours) (FT) (Subset): ----- ")
+    subset_tuned_r3pm_net_table = analyze_results(subset_tuned_r3pm_net_results, verbose=True)
+    show_successful_resutls(subset_tuned_r3pm_net_table, sources, targets, subset_tuned_r3pm_net_pc_results, 'R3PM-Net (ours) (FT) (Subset)')
+
+if __name__ == "__main__":
+    main()

scripts/modelnet40.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+#SBATCH --partition=gpu_h100
+#SBATCH --gpus=1
+#SBATCH --job-name=modelnet40
+#SBATCH --ntasks=1
+#SBATCH --time=09:00:00
+#SBATCH --output=modelnet40_output_%A.txt
+#SBATCH --error=modelnet40_error_%A.txt
+
+# Load necessary modules (adjust based on your environment)
+module purge
+module load 2023
+module load CUDA/12.1.1
+
+# my miniconda3 path
+export PATH="$HOME/miniconda3/bin:$PATH"
+unset -f conda 2>/dev/null      
+source "$HOME/miniconda3/etc/profile.d/conda.sh"
+
+# Activate the conda environment
+conda activate r3pm_net
+
+if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then
+  REPO_ROOT="$(cd "${SLURM_SUBMIT_DIR}" && pwd)"
+else
+  REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+fi
+cd "$REPO_ROOT" || { echo "ERROR: cannot cd to REPO_ROOT=${REPO_ROOT}" >&2; exit 1; }
+if [[ ! -f "${REPO_ROOT}/pyproject.toml" ]]; then
+  echo "ERROR: REPO_ROOT=${REPO_ROOT} is not the r3pm_net tree (missing pyproject.toml)." >&2
+  echo "Run: cd /path/to/r3pm_net && sbatch scripts/modelnet40.sh" >&2
+  exit 1
+fi
+
+LOGDIR="${REPO_ROOT}/logs/slurm"
+mkdir -p "$LOGDIR"
+JOB_ID="${SLURM_JOB_ID:-local}"
+
+# seeds=(42 61 92 114 123 456 789)
+seeds=(42)
+
+for seed in "${seeds[@]}"; do
+  srun python scripts/eval_modelnet40.py --seed "${seed}" \
+    >"${LOGDIR}/modelnet40_job${JOB_ID}_seed${seed}.log" 2>&1
+done

scripts/sioux_cranfield.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+#SBATCH --partition=gpu_h100
+#SBATCH --gpus=1
+#SBATCH --job-name=sioux_cranfield
+#SBATCH --ntasks=1  
+#SBATCH --time=04:00:00
+#SBATCH --output=sioux_cranfield_output_%A.txt
+#SBATCH --error=sioux_cranfield_error_%A.txt
+
+# Load necessary modules (adjust based on your environment)
+module purge
+module load 2023
+module load CUDA/12.1.1
+
+# my miniconda3 path
+export PATH="$HOME/miniconda3/bin:$PATH"
+unset -f conda 2>/dev/null      
+source "$HOME/miniconda3/etc/profile.d/conda.sh"
+
+# Activate the conda environment
+conda activate r3pm_net
+
+
+if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then
+  REPO_ROOT="$(cd "${SLURM_SUBMIT_DIR}" && pwd)"
+else
+  REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+fi
+cd "$REPO_ROOT" || { echo "ERROR: cannot cd to REPO_ROOT=${REPO_ROOT}" >&2; exit 1; }
+if [[ ! -f "${REPO_ROOT}/pyproject.toml" ]]; then
+  echo "ERROR: REPO_ROOT=${REPO_ROOT} is not the r3pm_net tree (missing pyproject.toml)." >&2
+  echo "Run: cd /path/to/r3pm_net && sbatch scripts/sioux_cranfield.sh" >&2
+  exit 1
+fi
+
+LOGDIR="${REPO_ROOT}/logs/slurm"
+mkdir -p "$LOGDIR"
+JOB_ID="${SLURM_JOB_ID:-local}"
+
+# seeds=(42 61 92 114 123 456 789)
+seeds=(42)
+
+for seed in "${seeds[@]}"; do
+  srun python scripts/eval_sioux_cranfield.py --seed "${seed}" \
+    >"${LOGDIR}/sioux_cranfield_job${JOB_ID}_seed${seed}.log" 2>&1
+done

scripts/sioux_scans.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+#SBATCH --partition=gpu_h100
+#SBATCH --gpus=1
+#SBATCH --job-name=sioux_scans
+#SBATCH --ntasks=1
+#SBATCH --time=01:00:00
+#SBATCH --output=sioux_scans_output_%A.txt
+#SBATCH --error=sioux_scans_error_%A.txt
+
+# Load necessary modules (adjust based on your environment)
+module purge
+module load 2023
+module load CUDA/12.1.1
+
+# my miniconda3 path
+export PATH="$HOME/miniconda3/bin:$PATH"
+unset -f conda 2>/dev/null      
+source "$HOME/miniconda3/etc/profile.d/conda.sh"
+
+# Activate the conda environment
+conda activate r3pm_net
+
+if [[ -n "${SLURM_SUBMIT_DIR:-}" ]]; then
+  REPO_ROOT="$(cd "${SLURM_SUBMIT_DIR}" && pwd)"
+else
+  REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+fi
+cd "$REPO_ROOT" || { echo "ERROR: cannot cd to REPO_ROOT=${REPO_ROOT}" >&2; exit 1; }
+if [[ ! -f "${REPO_ROOT}/pyproject.toml" ]]; then
+  echo "ERROR: REPO_ROOT=${REPO_ROOT} is not the r3pm_net tree (missing pyproject.toml)." >&2
+  echo "Run: cd /path/to/r3pm_net && sbatch scripts/sioux_scans.sh" >&2
+  exit 1
+fi
+
+LOGDIR="${REPO_ROOT}/logs/slurm"
+mkdir -p "$LOGDIR"
+JOB_ID="${SLURM_JOB_ID:-local}"
+
+# seeds=(42 61 92 114 123 456 789)
+seeds=(42)
+
+for seed in "${seeds[@]}"; do
+  srun python scripts/eval_sioux_scans.py --seed "${seed}" \
+    >"${LOGDIR}/sioux_scans_job${JOB_ID}_seed${seed}.log" 2>&1
+done

src/train.py

@@ -0,0 +1,366 @@
+import argparse
+import os
+import pickle
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from tensorboardX import SummaryWriter
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+# Repository root on PYTHONPATH (for `python src/train.py` or srun).
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from r3pm_net.model import R3PMNet
+from r3pm_net.config_loader import parse_train_args, resolve_path_args
+from r3pm_net.paths import REPO_ROOT
+from thirdparty.learning3d.losses import FrobeniusNormLoss, RMSEFeaturesLoss
+from dataloader.user_data import UserData
+from r3pm_net.feature_extractor import feature_extractor  # import your feature extractor here
+
+def _init_(args):
+    Path(args.save_dir).mkdir(parents=True, exist_ok=True)
+    (REPO_ROOT / "checkpoints" / args.exp_name).mkdir(parents=True, exist_ok=True)
+
+    if os.path.isfile("main.py"):
+        os.system("cp main.py checkpoints" + "/" + args.exp_name + "/" + "main.py.backup")
+    if os.path.isfile("model.py"):
+        os.system("cp model.py checkpoints" + "/" + args.exp_name + "/" + "model.py.backup")
+
+
+class IOStream:
+    def __init__(self, path):
+        self.f = open(path, "a")
+
+    def cprint(self, text):
+        print(text)
+        self.f.write(text + "\n")
+        self.f.flush()
+
+    def close(self):
+        self.f.close()
+
+
+def test_one_epoch(device, model, test_loader):
+    model.eval()
+    test_loss = 0.0
+    count = 0
+    for i, data in enumerate(tqdm(test_loader)):
+        template, source, igt = data
+
+        template = template.to(device)
+        source = source.to(device)
+        igt = igt.to(device)
+
+        output = model(template, source)
+        loss_val = FrobeniusNormLoss()(output["est_T"], igt) + RMSEFeaturesLoss()(output["r"])
+
+        test_loss += loss_val.item()
+        count += 1
+
+    test_loss = float(test_loss) / count
+    return test_loss
+
+
+def test(args, model, test_loader, textio):
+    test_loss = test_one_epoch(args.device, model, test_loader)
+    textio.cprint("Validation Loss: %f" % (test_loss))
+
+
+def train_one_epoch(device, model, train_loader, optimizer):
+    model.train()
+    train_loss = 0.0
+    count = 0
+    for i, data in enumerate(tqdm(train_loader)):
+        template, source, igt = data
+
+        template = template.to(device)
+        source = source.to(device)
+        igt = igt.to(device)
+
+        output = model(template, source)
+        loss_val = FrobeniusNormLoss()(output["est_T"], igt) + RMSEFeaturesLoss()(output["r"])
+
+        optimizer.zero_grad()
+        loss_val.backward()
+        optimizer.step()
+
+        train_loss += loss_val.item()
+        count += 1
+
+    train_loss = float(train_loss) / count
+    return train_loss
+
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+    Path(args.save_dir).mkdir(parents=True, exist_ok=True)
+
+    learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+    if args.optimizer == "Adam":
+        optimizer = torch.optim.Adam(learnable_params)
+    else:
+        optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+    if checkpoint is not None:
+        optimizer.load_state_dict(checkpoint["optimizer"])
+
+    best_test_loss = np.inf
+
+    for epoch in range(args.start_epoch, args.epochs):
+        train_loss = train_one_epoch(args.device, model, train_loader, optimizer)
+        test_loss = test_one_epoch(args.device, model, test_loader)
+
+        snap = {
+            "epoch": epoch + 1,
+            "model": model.state_dict(),
+            "min_loss": test_loss,
+            "optimizer": optimizer.state_dict(),
+        }
+
+        if test_loss < best_test_loss:
+            best_test_loss = test_loss
+            best_snap_path = os.path.join(
+                args.save_dir, "best_model_snap.t7")
+            best_model_path = os.path.join(
+                args.save_dir, "best_model.t7")
+
+            torch.save(snap, best_snap_path)
+            torch.save(model.state_dict(), best_model_path)
+
+        torch.save(snap, os.path.join(args.save_dir, "model_snap.t7"))
+        torch.save(model.state_dict(), os.path.join(args.save_dir, "model.t7"))
+
+        boardio.add_scalar("Train Loss", train_loss, epoch + 1)
+        boardio.add_scalar("Test Loss", test_loss, epoch + 1)
+        boardio.add_scalar("Best Test Loss", best_test_loss, epoch + 1)
+
+        textio.cprint(
+            "EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f"
+            % (epoch + 1, train_loss, test_loss, best_test_loss)
+        )
+
+
+def build_parser(default_config_path: str):
+    parser = argparse.ArgumentParser(description="Point Cloud Registration")
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=default_config_path,
+        help="YAML file with defaults (see config/default.yaml); can be overridden on the command line",
+    )
+    parser.add_argument(
+        "--exp_name",
+        type=str,
+        default="exp_r3pmnet",
+        metavar="N",
+        help="Name of the experiment",
+    )
+    parser.add_argument("--eval", action="store_true", help="Run evaluation only (no training).")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="",
+        help="Directory to save model checkpoints (default: checkpoints/<exp_name>/models)",
+    )
+
+    parser.add_argument(
+        "--num_points",
+        default=1024,
+        type=int,
+        metavar="N",
+        help="points in point-cloud (default: 1024)",
+    )
+
+    parser.add_argument(
+        "--fine_tune_feature_extractor",
+        default="tune",
+        type=str,
+        choices=["fixed", "tune"],
+        help="train feature extractor (default: tune)",
+    )
+    parser.add_argument(
+        "--transfer_weights",
+        default="",
+        type=str,
+        metavar="PATH",
+        help="optional path to feature extractor checkpoint",
+    )
+    parser.add_argument(
+        "--symfn",
+        default="max",
+        choices=["max", "avg"],
+        help="symmetric function (default: max)",
+    )
+
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "-j",
+        "--workers",
+        default=4,
+        type=int,
+        metavar="N",
+        help="number of data loading workers (default: 4)",
+    )
+    parser.add_argument(
+        "-b",
+        "--batch_size",
+        default=5,
+        type=int,
+        metavar="N",
+        help="mini-batch size (default: 5)",
+    )
+    parser.add_argument(
+        "--epochs",
+        default=50,
+        type=int,
+        metavar="N",
+        help="number of total epochs to run",
+    )
+    parser.add_argument(
+        "--start_epoch",
+        default=0,
+        type=int,
+        metavar="N",
+        help="manual epoch number (useful on restarts)",
+    )
+    parser.add_argument(
+        "--optimizer",
+        default="Adam",
+        choices=["Adam", "SGD"],
+        metavar="METHOD",
+        help="name of an optimizer (default: Adam)",
+    )
+    parser.add_argument(
+        "--resume",
+        default="",
+        type=str,
+        metavar="PATH",
+        help="path to latest checkpoint (default: none)",
+    )
+    parser.add_argument(
+        "--pretrained",
+        default="",
+        type=str,
+        metavar="PATH",
+        help="path to pretrained full model (default: none)",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda:0",
+        type=str,
+        metavar="DEVICE",
+        help="use CUDA if available",
+    )
+
+    parser.add_argument(
+        "--train_dict_path",
+        type=str,
+        default="data/simulators/data_dict_train.pkl",
+        help="Pickled training data_dict",
+    )
+    parser.add_argument(
+        "--test_dict_path",
+        type=str,
+        default="data/simulators/data_dict_test.pkl",
+        help="Pickled test data_dict",
+    )
+
+    return parser
+
+
+def _torch_load(path, map_location):
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
+def main():
+    args = parse_train_args(sys.argv[1:], build_parser)
+
+    resolve_path_args(
+        args,
+        (
+            "save_dir",
+            "train_dict_path",
+            "test_dict_path",
+            "resume",
+            "pretrained",
+            "transfer_weights",
+        ),
+    )
+
+    if not args.save_dir:
+        args.save_dir = str(REPO_ROOT / "checkpoints" / args.exp_name / "models")
+
+    torch.backends.cudnn.deterministic = True
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    np.random.seed(args.seed)
+
+    ckpt_dir = REPO_ROOT / "checkpoints" / args.exp_name
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    boardio = SummaryWriter(log_dir=str(ckpt_dir))
+    _init_(args)
+
+    textio = IOStream(str(ckpt_dir / "run.log"))
+    textio.cprint(str(args))
+
+    if not os.path.isfile(args.train_dict_path):
+        raise FileNotFoundError(f"Training dict not found: {args.train_dict_path}")
+    if not os.path.isfile(args.test_dict_path):
+        raise FileNotFoundError(f"Test dict not found: {args.test_dict_path}")
+
+    with open(args.train_dict_path, "rb") as f:
+        data_dict_train = pickle.load(f)
+    with open(args.test_dict_path, "rb") as f:
+        data_dict_test = pickle.load(f)
+
+    trainset = UserData("registration", data_dict_train)
+    testset = UserData("registration", data_dict_test)
+    train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=False, drop_last=True, num_workers=args.workers)
+    test_loader = DataLoader(testset, batch_size=5, shuffle=False, drop_last=False, num_workers=args.workers)
+
+    if not torch.cuda.is_available():
+        args.device = "cpu"
+    args.device = torch.device(args.device)
+
+    # feature extractor model
+    FEATURE_MODEL = feature_extractor
+    model = R3PMNet(feature_model=FEATURE_MODEL)
+    model = model.to(args.device)
+
+    if args.transfer_weights and os.path.isfile(args.transfer_weights):
+        feat_model_dict = _torch_load(args.transfer_weights, args.device)
+        model.feat_extractor.load_state_dict(feat_model_dict)
+
+    checkpoint = None
+    if args.resume:
+        assert os.path.isfile(args.resume)
+        checkpoint = _torch_load(args.resume, args.device)
+        args.start_epoch = checkpoint["epoch"]
+        model.load_state_dict(checkpoint["model"])
+
+    if args.pretrained:
+        assert os.path.isfile(args.pretrained)
+        try:
+            model.load_state_dict(_torch_load(args.pretrained, "cpu"))
+        except RuntimeError:
+            model_data = _torch_load(args.pretrained, "cpu")
+            state_dict = model_data["state_dict"]
+            model.load_state_dict(state_dict)
+    model.to(args.device)
+
+    Path(args.save_dir).mkdir(parents=True, exist_ok=True)
+
+    if args.eval:
+        test(args, model, test_loader, textio)
+    else:
+        train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == "__main__":
+    main()

thirdparty/__init__.py

@@ -0,0 +1 @@
+# Namespace for vendored thirdparty.learning3d

thirdparty/learning3d/data_utils/__init__.py

@@ -0,0 +1,4 @@
+from .dataloaders import ModelNet40Data
+from .dataloaders import ClassificationData, RegistrationData, SegmentationData, FlowData, SceneflowDataset
+from .dataloaders import download_modelnet40, deg_to_rad, create_random_transform
+from .user_data import UserData

thirdparty/learning3d/data_utils/dataloaders.py

@@ -0,0 +1,454 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F 
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+import numpy as np
+import os
+import h5py
+import subprocess
+import shlex
+import json
+import glob
+from .. ops import transform_functions, se3
+from sklearn.neighbors import NearestNeighbors
+from scipy.spatial.distance import minkowski
+from scipy.spatial import cKDTree
+from torch.utils.data import Dataset
+
+def download_modelnet40():
+	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+	if not os.path.exists(DATA_DIR):
+		os.mkdir(DATA_DIR)
+	if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
+		www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip'
+		zipfile = os.path.basename(www)
+		os.system('wget --no-check-certificate %s; unzip %s' % (www, zipfile))
+		os.system('mv %s %s' % (zipfile[:-4], DATA_DIR))
+		os.system('rm %s' % (zipfile))
+
+def load_data(train, use_normals):
+	if train: partition = 'train'
+	else: partition = 'test'
+	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+	all_data = []
+	all_label = []
+	for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5' % partition)):
+		f = h5py.File(h5_name)
+		if use_normals: data = np.concatenate([f['data'][:], f['normal'][:]], axis=-1).astype('float32')
+		else: data = f['data'][:].astype('float32')
+		label = f['label'][:].astype('int64')
+		f.close()
+		all_data.append(data)
+		all_label.append(label)
+	all_data = np.concatenate(all_data, axis=0)
+	all_label = np.concatenate(all_label, axis=0)
+	return all_data, all_label
+
+def deg_to_rad(deg):
+	return np.pi / 180 * deg
+
+def create_random_transform(dtype, max_rotation_deg, max_translation):
+	max_rotation = deg_to_rad(max_rotation_deg)
+	rot = np.random.uniform(-max_rotation, max_rotation, [1, 3])
+	trans = np.random.uniform(-max_translation, max_translation, [1, 3])
+	quat = transform_functions.euler_to_quaternion(rot, "xyz")
+
+	vec = np.concatenate([quat, trans], axis=1)
+	vec = torch.tensor(vec, dtype=dtype)
+	return vec
+
+def jitter_pointcloud(pointcloud, sigma=0.04, clip=0.05):
+	# N, C = pointcloud.shape
+	sigma = 0.04*np.random.random_sample()
+	pointcloud += torch.empty(pointcloud.shape).normal_(mean=0, std=sigma).clamp(-clip, clip)
+	return pointcloud
+
+def farthest_subsample_points(pointcloud1, num_subsampled_points=768):
+	pointcloud1 = pointcloud1
+	num_points = pointcloud1.shape[0]
+	nbrs1 = NearestNeighbors(n_neighbors=num_subsampled_points, algorithm='auto',
+							 metric=lambda x, y: minkowski(x, y)).fit(pointcloud1[:, :3])
+	random_p1 = np.random.random(size=(1, 3)) + np.array([[500, 500, 500]]) * np.random.choice([1, -1, 1, -1])
+	idx1 = nbrs1.kneighbors(random_p1, return_distance=False).reshape((num_subsampled_points,))
+	gt_mask = torch.zeros(num_points).scatter_(0, torch.tensor(idx1), 1)
+	return pointcloud1[idx1, :], gt_mask
+
+def uniform_2_sphere(num: int = None):
+	"""Uniform sampling on a 2-sphere
+
+	Source: https://gist.github.com/andrewbolster/10274979
+
+	Args:
+		num: Number of vectors to sample (or None if single)
+
+	Returns:
+		Random Vector (np.ndarray) of size (num, 3) with norm 1.
+		If num is None returned value will have size (3,)
+
+	"""
+	if num is not None:
+		phi = np.random.uniform(0.0, 2 * np.pi, num)
+		cos_theta = np.random.uniform(-1.0, 1.0, num)
+	else:
+		phi = np.random.uniform(0.0, 2 * np.pi)
+		cos_theta = np.random.uniform(-1.0, 1.0)
+
+	theta = np.arccos(cos_theta)
+	x = np.sin(theta) * np.cos(phi)
+	y = np.sin(theta) * np.sin(phi)
+	z = np.cos(theta)
+
+	return np.stack((x, y, z), axis=-1)
+
+def planar_crop(points, p_keep= 0.7):
+	p_keep = np.array(p_keep, dtype=np.float32)
+
+	rand_xyz = uniform_2_sphere()
+	pts = points.numpy()
+	centroid = np.mean(pts[:, :3], axis=0)
+	points_centered = pts[:, :3] - centroid
+
+	dist_from_plane = np.dot(points_centered, rand_xyz)
+
+	mask = dist_from_plane > np.percentile(dist_from_plane, (1.0 - p_keep) * 100)
+	idx_x = torch.Tensor(np.nonzero(mask))
+
+	return torch.Tensor(pts[mask, :3]), idx_x
+
+def knn_idx(pts, k):
+	kdt = cKDTree(pts) 
+	_, idx = kdt.query(pts, k=k+1)
+	return idx[:, 1:]
+
+def get_rri(pts, k):
+	# pts: N x 3, original points
+	# q: N x K x 3, nearest neighbors
+	q = pts[knn_idx(pts, k)]
+	p = np.repeat(pts[:, None], k, axis=1)
+	# rp, rq: N x K x 1, norms
+	rp = np.linalg.norm(p, axis=-1, keepdims=True)
+	rq = np.linalg.norm(q, axis=-1, keepdims=True)
+	pn = p / rp
+	qn = q / rq
+	dot = np.sum(pn * qn, -1, keepdims=True)
+	# theta: N x K x 1, angles
+	theta = np.arccos(np.clip(dot, -1, 1))
+	T_q = q - dot * p
+	sin_psi = np.sum(np.cross(T_q[:, None], T_q[:, :, None]) * pn[:, None], -1)
+	cos_psi = np.sum(T_q[:, None] * T_q[:, :, None], -1)
+	psi = np.arctan2(sin_psi, cos_psi) % (2*np.pi)
+	idx = np.argpartition(psi, 1)[:, :, 1:2]
+	# phi: N x K x 1, projection angles
+	phi = np.take_along_axis(psi, idx, axis=-1)
+	feat = np.concatenate([rp, rq, theta, phi], axis=-1)
+	return feat.reshape(-1, k * 4)
+
+def get_rri_cuda(pts, k, npts_per_block=1):
+	try:
+		import pycuda.autoinit
+		from pycuda import gpuarray
+		from pycuda.compiler import SourceModule
+	except Exception as e:
+		print("Error raised in pycuda modules! pycuda only works with GPU, ", e)
+		raise
+
+	mod_rri = SourceModule(open('rri.cu').read() % (k, npts_per_block))
+	rri_cuda = mod_rri.get_function('get_rri_feature')
+
+	N = len(pts)
+	pts_gpu = gpuarray.to_gpu(pts.astype(np.float32).ravel())
+	k_idx = knn_idx(pts, k)
+	k_idx_gpu = gpuarray.to_gpu(k_idx.astype(np.int32).ravel())
+	feat_gpu = gpuarray.GPUArray((N * k * 4,), np.float32)
+
+	rri_cuda(pts_gpu, np.int32(N), k_idx_gpu, feat_gpu,
+				grid=(((N-1) // npts_per_block)+1, 1),
+				block=(npts_per_block, k, 1))
+
+	feat = feat_gpu.get().reshape(N, k * 4).astype(np.float32)
+	return feat
+
+
+class UnknownDataTypeError(Exception):
+	def __init__(self, *args):
+		if args: self.message = args[0]
+		else: self.message = 'Datatype not understood for dataset.'
+
+	def __str__(self):
+		return self.message
+
+
+class ModelNet40Data(Dataset):
+	def __init__(
+		self,
+		train=True,
+		num_points=1024,
+		download=True,
+		randomize_data=False,
+		use_normals=False
+	):
+		super(ModelNet40Data, self).__init__()
+		if download: download_modelnet40()
+		self.data, self.labels = load_data(train, use_normals)
+		if not train: self.shapes = self.read_classes_ModelNet40()
+		self.num_points = num_points
+		self.randomize_data = randomize_data
+
+	def __getitem__(self, idx):
+		if self.randomize_data: current_points = self.randomize(idx)
+		else: current_points = self.data[idx].copy()
+
+		current_points = torch.from_numpy(current_points[:self.num_points, :]).float()
+		label = torch.from_numpy(self.labels[idx]).type(torch.LongTensor)
+
+		return current_points, label
+
+	def __len__(self):
+		return self.data.shape[0]
+
+	def randomize(self, idx):
+		pt_idxs = np.arange(0, self.num_points)
+		np.random.shuffle(pt_idxs)
+		return self.data[idx, pt_idxs].copy()
+
+	def get_shape(self, label):
+		return self.shapes[label]
+
+	def read_classes_ModelNet40(self):
+		BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+		DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+		file = open(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'shape_names.txt'), 'r')
+		shape_names = file.read()
+		shape_names = np.array(shape_names.split('\n')[:-1])
+		return shape_names
+
+
+class ClassificationData(Dataset):
+	def __init__(self, data_class=ModelNet40Data()):
+		super(ClassificationData, self).__init__()
+		self.set_class(data_class)
+
+	def __len__(self):
+		return len(self.data_class)
+
+	def set_class(self, data_class):
+		self.data_class = data_class
+
+	def get_shape(self, label):
+		try:
+			return self.data_class.get_shape(label)
+		except:
+			return -1
+
+	def __getitem__(self, index):
+		return self.data_class[index]
+
+
+class RegistrationData(Dataset):
+	def __init__(self, algorithm, data_class=ModelNet40Data(), partial_source=False, partial_template=False, noise=False, additional_params={}):
+		super(RegistrationData, self).__init__()
+		available_algorithms = ['PCRNet', 'PointNetLK', 'DCP', 'PRNet', 'iPCRNet', 'RPMNet', 'DeepGMR']
+		if algorithm in available_algorithms: self.algorithm = algorithm
+		else: raise Exception("Algorithm not available for registration.")
+		
+		self.set_class(data_class)
+		self.partial_template = partial_template
+		self.partial_source = partial_source
+		self.noise = noise
+		self.additional_params = additional_params
+		self.use_rri = False
+
+		if self.algorithm == 'PCRNet' or self.algorithm == 'iPCRNet':
+			from .. ops.transform_functions import PCRNetTransform
+			self.transforms = PCRNetTransform(len(data_class), angle_range=45, translation_range=1)
+		if self.algorithm == 'PointNetLK':
+			from .. ops.transform_functions import PNLKTransform
+			self.transforms = PNLKTransform(0.8, True)
+		if self.algorithm == 'RPMNet':
+			from .. ops.transform_functions import RPMNetTransform
+			self.transforms = RPMNetTransform(0.8, True)
+		if self.algorithm == 'DCP' or self.algorithm == 'PRNet':
+			from .. ops.transform_functions import DCPTransform
+			self.transforms = DCPTransform(angle_range=45, translation_range=1)
+		if self.algorithm == 'DeepGMR':
+			self.get_rri = get_rri_cuda if torch.cuda.is_available() else get_rri
+			from .. ops.transform_functions import DeepGMRTransform
+			self.transforms = DeepGMRTransform(angle_range=90, translation_range=1)
+			if 'nearest_neighbors' in self.additional_params.keys() and self.additional_params['nearest_neighbors'] > 0:
+				self.use_rri = True
+				self.nearest_neighbors = self.additional_params['nearest_neighbors']
+
+	def __len__(self):
+		return len(self.data_class)
+
+	def set_class(self, data_class):
+		self.data_class = data_class
+
+	def __getitem__(self, index):
+		template, label = self.data_class[index]
+		self.transforms.index = index				# for fixed transformations in PCRNet.
+		source = self.transforms(template)
+
+		# Check for Partial Data.
+		if self.additional_params.get('partial_point_cloud_method', None) == 'planar_crop':
+			source, gt_idx_source = planar_crop(source)
+			template, gt_idx_template = planar_crop(template)
+			intersect_mask, intersect_x, intersect_y  = np.intersect1d(gt_idx_source, gt_idx_template, return_indices=True)
+
+			self.template_mask = torch.zeros(template.shape[0])
+			self.source_mask = torch.zeros(source.shape[0])
+			self.template_mask[intersect_y]  = 1
+			self.source_mask[intersect_x]  = 1
+		else:
+			if self.partial_source: source, self.source_mask = farthest_subsample_points(source)
+			if self.partial_template: template, self.template_mask = farthest_subsample_points(template)
+
+
+
+		# Check for Noise in Source Data.
+		if self.noise: source = jitter_pointcloud(source)
+
+		if self.use_rri:
+			template, source = template.numpy(), source.numpy()
+			template = np.concatenate([template, self.get_rri(template - template.mean(axis=0), self.nearest_neighbors)], axis=1)
+			source = np.concatenate([source, self.get_rri(source - source.mean(axis=0), self.nearest_neighbors)], axis=1)
+			template, source = torch.tensor(template).float(), torch.tensor(source).float()
+
+		igt = self.transforms.igt
+		
+		if self.additional_params.get('use_masknet', False):
+			if self.partial_source and self.partial_template:
+				return template, source, igt, self.template_mask, self.source_mask
+			elif self.partial_source:
+				return template, source, igt, self.source_mask
+			elif self.partial_template:
+				return template, source, igt, self.template_mask
+		else:
+			return template, source, igt
+
+
+class SegmentationData(Dataset):
+	def __init__(self):
+		super(SegmentationData, self).__init__()
+
+	def __len__(self):
+		pass
+
+	def __getitem__(self, index):
+		pass
+
+
+class FlowData(Dataset):
+	def __init__(self):
+		super(FlowData, self).__init__()
+		self.pc1, self.pc2, self.flow = self.read_data()
+
+	def __len__(self):
+		if isinstance(self.pc1, np.ndarray):
+			return self.pc1.shape[0]
+		elif isinstance(self.pc1, list):
+			return len(self.pc1)
+		else:
+			raise UnknownDataTypeError
+
+	def read_data(self):
+		pass
+
+	def __getitem__(self, index):
+		return self.pc1[index], self.pc2[index], self.flow[index]
+
+
+class SceneflowDataset(Dataset):
+	def __init__(self, npoints=1024, root='', partition='train'):
+		if root == '':
+			BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+			DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+			root = os.path.join(DATA_DIR, 'data_processed_maxcut_35_20k_2k_8192')
+			if not os.path.exists(root): 
+				print("To download dataset, click here: https://drive.google.com/file/d/1CMaxdt-Tg1Wct8v8eGNwuT7qRSIyJPY-/view")
+				exit()
+			else:
+				print("SceneflowDataset Found Successfully!")
+
+		self.npoints = npoints
+		self.partition = partition
+		self.root = root
+		if self.partition=='train':
+			self.datapath = glob.glob(os.path.join(self.root, 'TRAIN*.npz'))
+		else:
+			self.datapath = glob.glob(os.path.join(self.root, 'TEST*.npz'))
+		self.cache = {}
+		self.cache_size = 30000
+
+		###### deal with one bad datapoint with nan value
+		self.datapath = [d for d in self.datapath if 'TRAIN_C_0140_left_0006-0' not in d]
+		######
+		print(self.partition, ': ',len(self.datapath))
+
+	def __getitem__(self, index):
+		if index in self.cache:
+			pos1, pos2, color1, color2, flow, mask1 = self.cache[index]
+		else:
+			fn = self.datapath[index]
+			with open(fn, 'rb') as fp:
+				data = np.load(fp)
+				pos1 = data['points1'].astype('float32')
+				pos2 = data['points2'].astype('float32')
+				color1 = data['color1'].astype('float32')
+				color2 = data['color2'].astype('float32')
+				flow = data['flow'].astype('float32')
+				mask1 = data['valid_mask1']
+
+			if len(self.cache) < self.cache_size:
+				self.cache[index] = (pos1, pos2, color1, color2, flow, mask1)
+
+		if self.partition == 'train':
+			n1 = pos1.shape[0]
+			sample_idx1 = np.random.choice(n1, self.npoints, replace=False)
+			n2 = pos2.shape[0]
+			sample_idx2 = np.random.choice(n2, self.npoints, replace=False)
+
+			pos1 = pos1[sample_idx1, :]
+			pos2 = pos2[sample_idx2, :]
+			color1 = color1[sample_idx1, :]
+			color2 = color2[sample_idx2, :]
+			flow = flow[sample_idx1, :]
+			mask1 = mask1[sample_idx1]
+		else:
+			pos1 = pos1[:self.npoints, :]
+			pos2 = pos2[:self.npoints, :]
+			color1 = color1[:self.npoints, :]
+			color2 = color2[:self.npoints, :]
+			flow = flow[:self.npoints, :]
+			mask1 = mask1[:self.npoints]
+
+		pos1_center = np.mean(pos1, 0)
+		pos1 -= pos1_center
+		pos2 -= pos1_center
+
+		return pos1, pos2, color1, color2, flow, mask1
+
+	def __len__(self):
+		return len(self.datapath)
+
+
+if __name__ == '__main__':
+	class Data():
+		def __init__(self):
+			super(Data, self).__init__()
+			self.data, self.label = self.read_data()
+
+		def read_data(self):
+			return [4,5,6], [4,5,6]
+
+		def __len__(self):
+			return len(self.data)
+
+		def __getitem__(self, idx):
+			return self.data[idx], self.label[idx]
+
+	cd = RegistrationData('abc')
+	import ipdb; ipdb.set_trace()

thirdparty/learning3d/data_utils/user_data.py

@@ -0,0 +1,119 @@
+import os
+import numpy as np
+import torch
+
+class ClassificationData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.pcs = self.find_attribute('pcs')
+		self.labels = self.find_attribute('labels')
+		self.check_data()
+
+	def find_attribute(self, attribute):
+		try:
+			attribute_data = self.data_dict[attribute]
+		except:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.pcs.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.pcs.shape)
+		assert 0 < len(self.labels.shape) < 3, "Error in dimension of labels! Given data dimension: {}".format(self.labels.shape)
+		
+		if len(self.pcs.shape)==2: self.pcs = self.pcs.reshape(1, -1, 3)
+		if len(self.labels.shape) == 1: self.labels = self.labels.reshape(1, -1)
+
+		assert self.pcs.shape[0] == self.labels.shape[0], "Inconsistency in the number of point clouds and number of ground truth labels!"
+
+
+	def __len__(self):
+		return self.pcs.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.pcs[index]).float(), torch.from_numpy(self.labels[idx]).type(torch.LongTensor)
+
+
+class RegistrationData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.template = self.find_attribute('template')
+		self.source = self.find_attribute('source')
+		self.transformation = self.find_attribute('transformation')
+		self.check_data()
+
+	def find_attribute(self, attribute):
+		try:
+			attribute_data = self.data[attribute]
+		except:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.template.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.template.shape)
+		assert 1 < len(self.source.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.source.shape)
+		assert 1 < len(self.transformation.shape) < 4, "Error in dimension of transformations! Given data dimension: {}".format(self.transformation.shape)
+
+		if len(self.template.shape)==2: self.template = self.template.reshape(1, -1, 3)
+		if len(self.source.shape)==2: self.source = self.source.reshape(1, -1, 3)
+		if len(self.transformation.shape) == 2: self.transformation = self.transformation.reshape(1, 4, 4)
+
+		assert self.template.shape[0] == self.source.shape[0], "Inconsistency in the number of template and source point clouds!"
+		assert self.source.shape[0] == self.transformation.shape[0], "Inconsistency in the number of transformation and source point clouds!"
+
+	def __len__(self):
+		return self.template.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.template[index]).float(), torch.tensor(self.source[index]).float(), torch.tensor(self.transformation[index]).float()
+
+
+class FlowData:
+	def __init__(self, data_dict):
+		self.data_dict = data_dict
+		self.frame1 = self.find_attribute('frame1')
+		self.frame2 = self.find_attribute('frame2')
+		self.flow = self.find_attribute('flow')
+		self.check_data()
+
+	def find_attribute(self, attribute):
+		try:
+			attribute_data = self.data[attribute]
+		except:
+			print("Given data directory has no key attribute \"{}\"".format(attribute))
+		return attribute_data
+
+	def check_data(self):
+		assert 1 < len(self.frame1.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame1.shape)
+		assert 1 < len(self.frame2.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame2.shape)
+		assert 1 < len(self.flow.shape) < 4, "Error in dimension of flow! Given data dimension: {}".format(self.flow.shape)
+
+		if len(self.frame1.shape)==2: self.frame1 = self.frame1.reshape(1, -1, 3)
+		if len(self.frame2.shape)==2: self.frame2 = self.frame2.reshape(1, -1, 3)
+		if len(self.flow.shape) == 2: self.flow = self.flow.reshape(1, -1, 3)
+
+		assert self.frame1.shape[0] == self.frame2.shape[0], "Inconsistency in the number of frame1 and frame2 point clouds!"
+		assert self.frame2.shape[0] == self.flow.shape[0], "Inconsistency in the number of flow and frame2 point clouds!"
+
+	def __len__(self):
+		return self.frame1.shape[0]
+
+	def __getitem__(self, index):
+		return torch.tensor(self.frame1[index]).float(), torch.tensor(self.frame2[index]).float(), torch.tensor(self.flow[index]).float()
+
+
+class UserData:
+	def __init__(self, application, data_dict):
+		self.application = application
+
+		if self.application == 'classification':
+			self.data_class = ClassificationData(data_dict)
+		elif self.application == 'registration':
+			self.data_class = RegistrationData(data_dict)
+		elif self.application == 'flow_estimation':
+			self.data_class = FlowData(data_dict)
+
+	def __len__(self):
+		return len(self.data_class)
+
+	def __getitem__(self, index):
+		return self.data_class[index]

thirdparty/learning3d/examples/test_curvenet.py

@@ -0,0 +1,118 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import CurveNet
+from learning3d.data_utils import ClassificationData, ModelNet40Data
+
+def display_open3d(template):
+	template_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	# template_.paint_uniform_color([1, 0, 0])
+	o3d.visualization.draw_geometries([template_])
+
+def test_one_epoch(device, model, test_loader, testset):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		points, target = data
+		target = target[:,0]
+
+		points = points.to(device)
+		target = target.to(device)
+
+		output = model(points)
+		loss_val = torch.nn.functional.nll_loss(
+			torch.nn.functional.log_softmax(output, dim=1), target, size_average=False)
+		print("Ground Truth Label: ", testset.get_shape(target[0].item()))
+		print("Predicted Label:    ", testset.get_shape(torch.argmax(output[0]).item()))
+		display_open3d(points.detach().cpu().numpy()[0])
+
+		test_loss += loss_val.item()
+		count += output.size(0)
+
+		_, pred1 = output.max(dim=1)
+		ag = (pred1 == target)
+		am = ag.sum()
+		pred += am.item()
+
+	test_loss = float(test_loss)/count
+	accuracy = float(pred)/count
+	return test_loss, accuracy
+
+def test(args, model, test_loader, testset):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader, testset)
+	print("Accuracy: ", test_accuracy*100)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for CurveNet
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--num_classes', default=40, type=int,
+						metavar='K', help='number of classes to be predicted')
+
+	# settings for on training
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_curvenet/models/model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	args.dataset_path = os.path.join(os.getcwd(), os.pardir, os.pardir, 'ModelNet40', 'ModelNet40')
+	
+	testset = ClassificationData(ModelNet40Data(train=False))
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	model = CurveNet(num_classes=args.num_classes, k=20)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		weights = torch.load(args.pretrained, map_location='cpu')
+		weights = {k[7:]: v for k, v in weights.items()}
+		model.load_state_dict(weights)
+	model.to(args.device)
+
+	test(args, model, test_loader, testset)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_dcp.py

@@ -0,0 +1,139 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import DGCNN, DCP
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def get_transformations(igt):
+	R_ba = igt[:, 0:3, 0:3]                             # Ps = R_ba * Pt
+	translation_ba = igt[:, 0:3, 3].unsqueeze(2)        # Ps = Pt + t_ba
+	R_ab = R_ba.permute(0, 2, 1)                        # Pt = R_ab * Ps
+	translation_ab = -torch.bmm(R_ab, translation_ba)   # Pt = Ps + t_ab
+	return R_ab, translation_ab, R_ba, translation_ba
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+		transformations = get_transformations(igt)
+		transformations = [t.to(device) for t in transformations]
+		R_ab, translation_ab, R_ba, translation_ba = transformations
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		display_open3d(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], output['transformed_source'].detach().cpu().numpy()[0])
+
+		identity = torch.eye(3).cuda().unsqueeze(0).repeat(template.shape[0], 1, 1)
+		loss_val = torch.nn.functional.mse_loss(torch.matmul(output['est_R'].transpose(2, 1), R_ab), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t'], translation_ab[:,:,0])
+
+		cycle_loss = torch.nn.functional.mse_loss(torch.matmul(output['est_R_'].transpose(2, 1), R_ba), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t_'], translation_ba[:,:,0])
+		loss_val = loss_val + cycle_loss * 0.1
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_ipcrnet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('--emb_dims', default=512, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=2, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_dcp/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+	
+	trainset = RegistrationData('DCP', ModelNet40Data(train=True))
+	testset = RegistrationData('DCP', ModelNet40Data(train=False))
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	dgcnn = DGCNN(emb_dims=args.emb_dims)
+	model = DCP(feature_model=dgcnn, cycle=True)
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'), strict=False)
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_deepgmr.py

@@ -0,0 +1,144 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import DeepGMR
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def rotation_error(R, R_gt):
+	cos_theta = (torch.einsum('bij,bij->b', R, R_gt) - 1) / 2
+	cos_theta = torch.clamp(cos_theta, -1, 1)
+	return torch.acos(cos_theta) * 180 / math.pi
+
+def translation_error(t, t_gt):
+	return torch.norm(t - t_gt, dim=1)
+
+def rmse(pts, T, T_gt):
+	pts_pred = pts @ T[:, :3, :3].transpose(1, 2) + T[:, :3, 3].unsqueeze(1)
+	pts_gt = pts @ T_gt[:, :3, :3].transpose(1, 2) + T_gt[:, :3, 3].unsqueeze(1)
+	return torch.norm(pts_pred - pts_gt, dim=2).mean(dim=1)
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	rotation_errors, translation_errors, rmses = [], [], []
+
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		display_open3d(template.detach().cpu().numpy()[0, :, :3], source.detach().cpu().numpy()[0, :, :3], output['transformed_source'].detach().cpu().numpy()[0])
+
+		eye = torch.eye(4).expand_as(igt).to(igt.device)
+		mse1 = F.mse_loss(output['est_T_inverse'] @ torch.inverse(igt), eye)
+		mse2 = F.mse_loss(output['est_T'] @ igt, eye)
+		loss = mse1 + mse2
+
+		r_err = rotation_error(est_T_inverse[:, :3, :3], igt[:, :3, :3])
+		t_err = translation_error(est_T_inverse[:, :3, 3], igt[:, :3, 3])
+		rmse_val = rmse(template[:, :100], est_T_inverse, igt)
+		rotation_errors.append(r_err)
+		translation_errors.append(t_err)
+		rmses.append(rmse_val)
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	print("Mean rotation error: {}, Mean translation error: {} and Mean RMSE: {}".format(np.mean(rotation_errors), np.mean(translation_errors), np.mean(rmses)))
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss = test_one_epoch(args.device, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_deepgmr', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	parser.add_argument('--nearest_neighbors', default=20, type=int,
+						metavar='K', help='No of nearest neighbors to be estimated.')
+	parser.add_argument('--use_rri', default=True, type=bool,
+						help='Find nearest neighbors to estimate features from PointNet.')
+
+	# settings for on training
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=2, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_deepgmr/models/best_model.pth', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+	
+	trainset = RegistrationData('DeepGMR', ModelNet40Data(train=True))
+	testset = RegistrationData('DeepGMR', ModelNet40Data(train=False))
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	model = DeepGMR(use_rri=args.use_rri, nearest_neighbors=args.nearest_neighbors)
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'), strict=False)
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_flownet.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+
+import open3d as o3d
+import os
+import gc
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.optim.lr_scheduler import MultiStepLR
+from learning3d.models import FlowNet3D
+from learning3d.data_utils import SceneflowDataset
+import numpy as np
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+def display_open3d(template, source, transformed_source):
+    template_ = o3d.geometry.PointCloud()
+    source_ = o3d.geometry.PointCloud()
+    transformed_source_ = o3d.geometry.PointCloud()
+    template_.points = o3d.utility.Vector3dVector(template)
+    source_.points = o3d.utility.Vector3dVector(source + np.array([0,0.5,0.5]))
+    transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+    template_.paint_uniform_color([1, 0, 0])
+    source_.paint_uniform_color([0, 1, 0])
+    transformed_source_.paint_uniform_color([0, 0, 1])
+    o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(args, net, test_loader):
+    net.eval()
+
+    total_loss = 0
+    num_examples = 0
+    for i, data in enumerate(tqdm(test_loader)):
+        data = [d.to(args.device) for d in data]
+        pc1, pc2, color1, color2, flow, mask1 = data
+        pc1 = pc1.transpose(2,1).contiguous()
+        pc2 = pc2.transpose(2,1).contiguous()
+        color1 = color1.transpose(2,1).contiguous()
+        color2 = color2.transpose(2,1).contiguous()
+        flow = flow
+        mask1 = mask1.float()
+
+        batch_size = pc1.size(0)
+        num_examples += batch_size
+        flow_pred = net(pc1, pc2, color1, color2).permute(0,2,1)
+        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) * (flow_pred - flow), -1) / 2.0)
+
+        pc1, pc2 = pc1.permute(0,2,1), pc2.permute(0,2,1)
+        pc1_ = pc1 - flow_pred
+        print("Loss: ", loss_1)
+        display_open3d(pc1.detach().cpu().numpy()[0], pc2.detach().cpu().numpy()[0], pc1_.detach().cpu().numpy()[0])
+        total_loss += loss_1.item() * batch_size        
+
+    return total_loss * 1.0 / num_examples
+
+
+def test(args, net, test_loader):
+    test_loss = test_one_epoch(args, net, test_loader)
+
+def main():
+    parser = argparse.ArgumentParser(description='Point Cloud Registration')
+    parser.add_argument('--model', type=str, default='flownet', metavar='N',
+                        choices=['flownet'], help='Model to use, [flownet]')
+    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
+                        help='Dimension of embeddings')
+    parser.add_argument('--num_points', type=int, default=2048,
+                        help='Point Number [default: 2048]')
+    parser.add_argument('--test_batch_size', type=int, default=1, metavar='batch_size',
+                        help='Size of batch)')
+    
+    parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N',
+                        help='Wheter to add gaussian noise')
+    parser.add_argument('--unseen', type=bool, default=False, metavar='N',
+                        help='Whether to test on unseen category')
+    parser.add_argument('--dataset', type=str, default='SceneflowDataset',
+                        choices=['SceneflowDataset'], metavar='N',
+                        help='dataset to use')
+    parser.add_argument('--dataset_path', type=str, default='data_processed_maxcut_35_20k_2k_8192', metavar='N',
+                        help='dataset to use')
+    parser.add_argument('--pretrained', type=str, default='learning3d/pretrained/exp_flownet/models/model.best.t7', metavar='N',
+                        help='Pretrained model path')
+    parser.add_argument('--device', default='cuda:0', type=str,
+                        metavar='DEVICE', help='use CUDA if available')
+
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        args.device = torch.device('cpu')
+    else:
+        args.device = torch.device('cuda')
+
+    if args.dataset == 'SceneflowDataset':
+        test_loader = DataLoader(
+            SceneflowDataset(npoints=args.num_points, partition='test'),
+            batch_size=args.test_batch_size, shuffle=False, drop_last=False)
+    else:
+        raise Exception("not implemented")
+
+    net = FlowNet3D()
+    assert os.path.exists(args.pretrained), "Pretrained Model Doesn't Exists!"
+    net.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+    net = net.to(args.device)
+        
+    test(args, net, test_loader)
+    print('FINISH')
+
+
+if __name__ == '__main__':
+    main()

thirdparty/learning3d/examples/test_masknet.py

@@ -0,0 +1,159 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import MaskNet
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def pc2open3d(data):
+	if torch.is_tensor(data): data = data.detach().cpu().numpy()
+	if len(data.shape) == 2:
+		pc = o3d.geometry.PointCloud()
+		pc.points = o3d.utility.Vector3dVector(data)
+		return pc
+	else:
+		print("Error in the shape of data given to Open3D!, Shape is ", data.shape)
+
+def display_results(template, source, masked_template):
+	template = pc2open3d(template)
+	source = pc2open3d(source)
+	masked_template = pc2open3d(masked_template)
+	
+	template.paint_uniform_color([1, 0, 0])
+	source.paint_uniform_color([0, 1, 0])
+	masked_template.paint_uniform_color([0, 0, 1])
+
+	o3d.visualization.draw_geometries([template, source])
+	o3d.visualization.draw_geometries([masked_template, source])
+
+def evaluate_metrics(TP, FP, FN, TN, gt_mask):
+	# TP, FP, FN, TN: 		True +ve, False +ve, False -ve, True -ve
+	# gt_mask:				Ground Truth mask [Nt, 1]
+	
+	accuracy = (TP + TN)/gt_mask.shape[1]
+	misclassification_rate = (FN + FP)/gt_mask.shape[1]
+	# Precision: (What portion of positive identifications are actually correct?)
+	precision = TP / (TP + FP)
+	# Recall: (What portion of actual positives are identified correctly?)
+	recall = TP / (TP + FN)
+
+	fscore = (2*precision*recall) / (precision + recall)
+	return accuracy, precision, recall, fscore
+
+# Function used to evaluate the predicted mask with ground truth mask.
+def evaluate_mask(gt_mask, predicted_mask, predicted_mask_idx):
+	# gt_mask:					Ground Truth Mask [Nt, 1]
+	# predicted_mask:			Mask predicted by network [Nt, 1]
+	# predicted_mask_idx:		Point indices chosen by network [Ns, 1]
+
+	if torch.is_tensor(gt_mask): gt_mask = gt_mask.detach().cpu().numpy()
+	if torch.is_tensor(gt_mask): predicted_mask = predicted_mask.detach().cpu().numpy()
+	if torch.is_tensor(predicted_mask_idx): predicted_mask_idx = predicted_mask_idx.detach().cpu().numpy()
+	gt_mask, predicted_mask, predicted_mask_idx = gt_mask.reshape(1,-1), predicted_mask.reshape(1,-1), predicted_mask_idx.reshape(1,-1)
+	
+	gt_idx = np.where(gt_mask == 1)[1].reshape(1,-1) 				# Find indices of points which are actually in source.
+
+	# TP + FP = number of source points.
+	TP = np.intersect1d(predicted_mask_idx[0], gt_idx[0]).shape[0]			# is inliner and predicted as inlier (True Positive) 		(Find common indices in predicted_mask_idx, gt_idx)
+	FP = len([x for x in predicted_mask_idx[0] if x not in gt_idx])			# isn't inlier but predicted as inlier (False Positive)
+	FN = FP															# is inlier but predicted as outlier (False Negative) (due to binary classification)
+	TN = gt_mask.shape[1] - gt_idx.shape[1] - FN 					# is outlier and predicted as outlier (True Negative)
+	return evaluate_metrics(TP, FP, FN, TN, gt_mask)
+
+def test_one_epoch(args, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	precision_list = []
+
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt, gt_mask = data
+
+		template = template.to(args.device)
+		source = source.to(args.device)
+		igt = igt.to(args.device)						# [source] = [igt]*[template]
+		gt_mask = gt_mask.to(args.device)
+
+		masked_template, predicted_mask = model(template, source)
+		
+		# Evaluate mask based on classification metrics.
+		accuracy, precision, recall, fscore = evaluate_mask(gt_mask, predicted_mask, predicted_mask_idx = model.mask_idx)
+		precision_list.append(precision)
+		
+		# Different ways to visualize results.
+		display_results(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], masked_template.detach().cpu().numpy()[0])
+
+	print("Mean Precision: ", np.mean(precision_list))
+
+def test(args, model, test_loader):
+	test_one_epoch(args, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='MaskNet: A Fully-Convolutional Network For Inlier Estimation (Testing)')
+
+	# settings for input data
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+	parser.add_argument('--partial_source', default=True, type=bool,
+						help='create partial source point cloud in dataset.')
+	parser.add_argument('--noise', default=False, type=bool,
+						help='Add noise in source point clouds.')
+	parser.add_argument('--outliers', default=False, type=bool,
+						help='Add outliers to template point cloud.')
+
+	# settings for on testing
+	parser.add_argument('-j', '--workers', default=1, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--test_batch_size', default=1, type=int,
+						metavar='N', help='test-mini-batch size (default: 1)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_masknet/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+	parser.add_argument('--unseen', default=False, type=bool,
+						help='Use first 20 categories for training and last 20 for testing')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+
+	testset = RegistrationData('PointNetLK', ModelNet40Data(train=False, num_points=args.num_points),
+									partial_source=args.partial_source, noise=args.noise,
+									additional_params={'use_masknet': True})
+	test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Load Pretrained MaskNet.
+	model = MaskNet()
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model = model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_masknet2.py

@@ -0,0 +1,162 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import MaskNet2
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def pc2open3d(data):
+	if torch.is_tensor(data): data = data.detach().cpu().numpy()
+	if len(data.shape) == 2:
+		pc = o3d.geometry.PointCloud()
+		pc.points = o3d.utility.Vector3dVector(data)
+		return pc
+	else:
+		print("Error in the shape of data given to Open3D!, Shape is ", data.shape)
+
+def display_results(template, source, masked_template, masked_source):
+	template = pc2open3d(template)
+	source = pc2open3d(source)
+	masked_template = pc2open3d(masked_template)
+	masked_source = pc2open3d(masked_source)
+	
+	template.paint_uniform_color([1, 0, 0])
+	source.paint_uniform_color([0, 1, 0])
+	# masked_template.paint_uniform_color([0, 0, 1])
+	masked_template.paint_uniform_color([1, 0, 0])
+	masked_source.paint_uniform_color([0, 1, 0])
+
+	o3d.visualization.draw_geometries([template, source])
+	o3d.visualization.draw_geometries([masked_template, masked_source])
+
+def evaluate_metrics(TP, FP, FN, TN, gt_mask):
+	# TP, FP, FN, TN: 		True +ve, False +ve, False -ve, True -ve
+	# gt_mask:				Ground Truth mask [Nt, 1]
+	
+	accuracy = (TP + TN)/gt_mask.shape[1]
+	misclassification_rate = (FN + FP)/gt_mask.shape[1]
+	# Precision: (What portion of positive identifications are actually correct?)
+	precision = TP / (TP + FP)
+	# Recall: (What portion of actual positives are identified correctly?)
+	recall = TP / (TP + FN)
+
+	fscore = (2*precision*recall) / (precision + recall)
+	return accuracy, precision, recall, fscore
+
+# Function used to evaluate the predicted mask with ground truth mask.
+def evaluate_mask(gt_mask, predicted_mask, predicted_mask_idx):
+	# gt_mask:					Ground Truth Mask [Nt, 1]
+	# predicted_mask:			Mask predicted by network [Nt, 1]
+	# predicted_mask_idx:		Point indices chosen by network [Ns, 1]
+
+	if torch.is_tensor(gt_mask): gt_mask = gt_mask.detach().cpu().numpy()
+	if torch.is_tensor(gt_mask): predicted_mask = predicted_mask.detach().cpu().numpy()
+	if torch.is_tensor(predicted_mask_idx): predicted_mask_idx = predicted_mask_idx.detach().cpu().numpy()
+	gt_mask, predicted_mask, predicted_mask_idx = gt_mask.reshape(1,-1), predicted_mask.reshape(1,-1), predicted_mask_idx.reshape(1,-1)
+	
+	gt_idx = np.where(gt_mask == 1)[1].reshape(1,-1) 				# Find indices of points which are actually in source.
+
+	# TP + FP = number of source points.
+	TP = np.intersect1d(predicted_mask_idx[0], gt_idx[0]).shape[0]			# is inliner and predicted as inlier (True Positive) 		(Find common indices in predicted_mask_idx, gt_idx)
+	FP = len([x for x in predicted_mask_idx[0] if x not in gt_idx])			# isn't inlier but predicted as inlier (False Positive)
+	FN = FP															# is inlier but predicted as outlier (False Negative) (due to binary classification)
+	TN = gt_mask.shape[1] - gt_idx.shape[1] - FN 					# is outlier and predicted as outlier (True Negative)
+	return evaluate_metrics(TP, FP, FN, TN, gt_mask)
+
+def test_one_epoch(args, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt, gt_template_mask, gt_source_mask = data
+
+		template = template.to(args.device)
+		source = source.to(args.device)
+		igt = igt.to(args.device)						# [source] = [igt]*[template]
+		gt_template_mask = gt_template_mask.to(args.device)
+		gt_source_mask = gt_source_mask.to(args.device)
+
+		masked_template, masked_source, template_mask, source_mask = model(template, source)
+		
+		# TODO: Implement evaluation strategy.
+		'''
+		Evaluate mask based on classification metrics.
+		accuracy, precision, recall, fscore = evaluate_mask(gt_template_mask, template_mask, predicted_mask_idx = model.mask_idx)
+		precision_list.append(precision)
+		'''
+		
+		# Different ways to visualize results.
+		display_results(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], masked_template.detach().cpu().numpy()[0], masked_source.detach().cpu().numpy()[0])
+
+def test(args, model, test_loader):
+	test_one_epoch(args, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='MaskNet: A Fully-Convolutional Network For Inlier Estimation (Testing)')
+
+	# settings for input data
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+	parser.add_argument('--partial_source', default=True, type=bool,
+						help='create partial source point cloud in dataset.')
+	parser.add_argument('--partial_template', default=True, type=bool,
+						help='create partial source point cloud in dataset.')
+	parser.add_argument('--noise', default=False, type=bool,
+						help='Add noise in source point clouds.')
+	parser.add_argument('--outliers', default=False, type=bool,
+						help='Add outliers to template point cloud.')
+
+	# settings for on testing
+	parser.add_argument('-j', '--workers', default=1, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--test_batch_size', default=1, type=int,
+						metavar='N', help='test-mini-batch size (default: 1)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_masknet2/models/best_model_0.7.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+	parser.add_argument('--unseen', default=False, type=bool,
+						help='Use first 20 categories for training and last 20 for testing')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+
+	testset = RegistrationData('PointNetLK', ModelNet40Data(train=False, num_points=args.num_points),
+									partial_template=args.partial_template, partial_source=args.partial_source, 
+									noise=args.noise, additional_params={'use_masknet': True, 'partial_point_cloud_method': 'planar_crop'})
+	test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Load Pretrained MaskNet.
+	model = MaskNet2()
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model = model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_pcn.py

@@ -0,0 +1,118 @@
+# author: Vinit Sarode (vinitsarode5@gmail.com) 03/23/2020
+
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import PCN
+from learning3d.data_utils import ModelNet40Data, ClassificationData
+from learning3d.losses import ChamferDistanceLoss
+
+def display_open3d(input_pc, output):
+	input_pc_ = o3d.geometry.PointCloud()
+	output_ = o3d.geometry.PointCloud()
+	input_pc_.points = o3d.utility.Vector3dVector(input_pc)
+	output_.points = o3d.utility.Vector3dVector(output + np.array([1,0,0]))
+	input_pc_.paint_uniform_color([1, 0, 0])
+	output_.paint_uniform_color([0, 1, 0])
+	o3d.visualization.draw_geometries([input_pc_, output_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		points, _ = data
+
+		points = points.to(device)
+
+		output = model(points)
+		loss_val = ChamferDistanceLoss()(points, output['coarse_output'])
+		print("Loss Val: ", loss_val)
+		display_open3d(points[0].detach().cpu().numpy(), output['coarse_output'][0].detach().cpu().numpy())
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss = test_one_epoch(args.device, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Completion Network')
+	parser.add_argument('--exp_name', type=str, default='exp_pcn', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PCN
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--detailed_output', default=False, type=bool,
+						help='Coarse + Fine Output')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_pcn/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	args.dataset_path = os.path.join(os.getcwd(), os.pardir, os.pardir, 'ModelNet40', 'ModelNet40')
+
+	trainset = ClassificationData(ModelNet40Data(train=True))
+	testset = ClassificationData(ModelNet40Data(train=False))
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	model = PCN(emb_dims=args.emb_dims, detailed_output=args.detailed_output)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_pcrnet.py

@@ -0,0 +1,120 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+
+from learning3d.models import PointNet, iPCRNet
+from learning3d.losses import ChamferDistanceLoss
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		display_open3d(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], output['transformed_source'].detach().cpu().numpy()[0])		
+		loss_val = ChamferDistanceLoss()(template, output['transformed_source'])
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_ipcrnet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=20, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_ipcrnet/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	testset = RegistrationData('PCRNet', ModelNet40Data(train=False))
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	ptnet = PointNet(emb_dims=args.emb_dims)
+	model = iPCRNet(feature_model=ptnet)
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_pnlk.py

@@ -0,0 +1,121 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+
+from learning3d.models import PointNet, PointNetLK
+from learning3d.losses import FrobeniusNormLoss, RMSEFeaturesLoss
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+
+		display_open3d(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], output['transformed_source'].detach().cpu().numpy()[0])
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt) + RMSEFeaturesLoss()(output['r'])
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+	
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_pnlk_v1', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=10, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_pnlk/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	testset = RegistrationData('PointNetLK', ModelNet40Data(train=False))
+	test_loader = DataLoader(testset, batch_size=8, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	ptnet = PointNet(emb_dims=args.emb_dims, use_bn=True)
+	model = PointNetLK(feature_model=ptnet)
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_pointconv.py

@@ -0,0 +1,126 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import create_pointconv
+from learning3d.models import Classifier
+from learning3d.data_utils import ClassificationData, ModelNet40Data
+
+def display_open3d(template):
+	template_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	# template_.paint_uniform_color([1, 0, 0])
+	o3d.visualization.draw_geometries([template_])
+
+def test_one_epoch(device, model, test_loader, testset):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		points, target = data
+		target = target[:,0]
+
+		points = points.to(device)
+		target = target.to(device)
+
+		output = model(points)
+		loss_val = torch.nn.functional.nll_loss(
+			torch.nn.functional.log_softmax(output, dim=1), target, size_average=False)
+		print("Ground Truth Label: ", testset.get_shape(target[0].item()))
+		print("Predicted Label:    ", testset.get_shape(torch.argmax(output[0]).item()))
+		display_open3d(points.detach().cpu().numpy()[0])
+
+		test_loss += loss_val.item()
+		count += output.size(0)
+
+		_, pred1 = output.max(dim=1)
+		ag = (pred1 == target)
+		am = ag.sum()
+		pred += am.item()
+
+	test_loss = float(test_loss)/count
+	accuracy = float(pred)/count
+	return test_loss, accuracy
+
+def test(args, model, test_loader, testset):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader, testset)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_classifier/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	args.dataset_path = os.path.join(os.getcwd(), os.pardir, os.pardir, 'ModelNet40', 'ModelNet40')
+	
+	testset = ClassificationData(ModelNet40Data(train=False))
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# To use pretrained model provided by authors.
+	# PointConv = create_pointconv(classifier=True, pretrained='path of pretrained model.')
+	# model = PointConv(emb_dims=args.emb_dims, classifier=True, pretrained='path of pretrained model.')
+
+	# To use your own pretrained model.
+	PointConv = create_pointconv(classifier=False, pretrained=None)
+	ptconv = PointConv(emb_dims=args.emb_dims, classifier=True, pretrained=None)
+	model = Classifier(feature_model=ptconv)	
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	test(args, model, test_loader, testset)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_pointnet.py

@@ -0,0 +1,121 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import PointNet
+from learning3d.models import Classifier
+from learning3d.data_utils import ClassificationData, ModelNet40Data
+
+def display_open3d(template):
+	template_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	# template_.paint_uniform_color([1, 0, 0])
+	o3d.visualization.draw_geometries([template_])
+
+def test_one_epoch(device, model, test_loader, testset):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		points, target = data
+		target = target[:,0]
+
+		points = points.to(device)
+		target = target.to(device)
+
+		output = model(points)
+		loss_val = torch.nn.functional.nll_loss(
+			torch.nn.functional.log_softmax(output, dim=1), target, size_average=False)
+		print("Ground Truth Label: ", testset.get_shape(target[0].item()))
+		print("Predicted Label:    ", testset.get_shape(torch.argmax(output[0]).item()))
+		display_open3d(points.detach().cpu().numpy()[0])
+
+		test_loss += loss_val.item()
+		count += output.size(0)
+
+		_, pred1 = output.max(dim=1)
+		ag = (pred1 == target)
+		am = ag.sum()
+		pred += am.item()
+
+	test_loss = float(test_loss)/count
+	accuracy = float(pred)/count
+	return test_loss, accuracy
+
+def test(args, model, test_loader, testset):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader, testset)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_classifier/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	args.dataset_path = os.path.join(os.getcwd(), os.pardir, os.pardir, 'ModelNet40', 'ModelNet40')
+	
+	testset = ClassificationData(ModelNet40Data(train=False))
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	ptnet = PointNet(emb_dims=args.emb_dims, use_bn=True)
+	model = Classifier(feature_model=ptnet)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	test(args, model, test_loader, testset)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_prnet.py

@@ -0,0 +1,126 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import PRNet
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def get_transformations(igt):
+	R_ba = igt[:, 0:3, 0:3]                             # Ps = R_ba * Pt
+	translation_ba = igt[:, 0:3, 3].unsqueeze(2)        # Ps = Pt + t_ba
+	R_ab = R_ba.permute(0, 2, 1)                        # Pt = R_ab * Ps
+	translation_ab = -torch.bmm(R_ab, translation_ba)   # Pt = Ps + t_ab
+	return R_ab, translation_ab, R_ba, translation_ba
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		transformations = get_transformations(igt)
+		transformations = [t.to(device) for t in transformations]
+		R_ab, translation_ab, R_ba, translation_ba = transformations
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source, R_ab, translation_ab.squeeze(2))
+		display_open3d(template.detach().cpu().numpy()[0], source.detach().cpu().numpy()[0], output['transformed_source'].detach().cpu().numpy()[0])
+
+		test_loss += output['loss'].item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss = test_one_epoch(args.device, model, test_loader)
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_prnet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+
+	# settings for PointNet
+	parser.add_argument('--emb_dims', default=512, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--num_iterations', default=3, type=int,
+						help='Number of Iterations')
+
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=1, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_prnet/models/best_model.t7', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+	
+	trainset = RegistrationData('PRNet', ModelNet40Data(train=True), partial_source=True, partial_template=True)
+	testset = RegistrationData('PRNet', ModelNet40Data(train=False), partial_source=True, partial_template=True)
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	model = PRNet(emb_dims=args.emb_dims, num_iters=args.num_iterations)
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'), strict=False)
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/test_rpmnet.py

@@ -0,0 +1,120 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+
+from learning3d.models import RPMNet, PPFNet
+from learning3d.losses import FrobeniusNormLoss, RMSEFeaturesLoss
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+
+		display_open3d(template.detach().cpu().numpy()[0,:,:3], source.detach().cpu().numpy()[0,:,:3], output['transformed_source'].detach().cpu().numpy()[0])
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt)
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+	
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_rpmnet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=10, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='learning3d/pretrained/exp_rpmnet/models/partial-trained.pth', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	testset = RegistrationData('RPMNet', ModelNet40Data(train=False, num_points=args.num_points, use_normals=True), partial_source=True, partial_template=False)
+	test_loader = DataLoader(testset, batch_size=1, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create RPMNet Model.
+	model = RPMNet(feature_model=PPFNet())
+	model = model.to(args.device)
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu')['state_dict'])
+	model.to(args.device)
+
+	test(args, model, test_loader)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/train_PointNetLK.py

@@ -0,0 +1,240 @@
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import PointNet
+from learning3d.models import PointNetLK
+from learning3d.losses import FrobeniusNormLoss, RMSEFeaturesLoss
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def _init_(args):
+	if not os.path.exists('checkpoints'):
+		os.makedirs('checkpoints')
+	if not os.path.exists('checkpoints/' + args.exp_name):
+		os.makedirs('checkpoints/' + args.exp_name)
+	if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
+		os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
+	os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' + 'main.py.backup')
+	os.system('cp model.py checkpoints' + '/' + args.exp_name + '/' + 'model.py.backup')
+
+
+class IOStream:
+	def __init__(self, path):
+		self.f = open(path, 'a')
+
+	def cprint(self, text):
+		print(text)
+		self.f.write(text + '\n')
+		self.f.flush()
+
+	def close(self):
+		self.f.close()
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt) + RMSEFeaturesLoss()(output['r'])
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader, textio):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+	textio.cprint('Validation Loss: %f & Validation Accuracy: %f'%(test_loss, test_accuracy))
+
+def train_one_epoch(device, model, train_loader, optimizer):
+	model.train()
+	train_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(train_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt) + RMSEFeaturesLoss()(output['r'])
+		# print(loss_val.item())
+
+		# forward + backward + optimize
+		optimizer.zero_grad()
+		loss_val.backward()
+		optimizer.step()
+
+		train_loss += loss_val.item()
+		count += 1
+
+	train_loss = float(train_loss)/count
+	return train_loss
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+	learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+	if args.optimizer == 'Adam':
+		optimizer = torch.optim.Adam(learnable_params)
+	else:
+		optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+	if checkpoint is not None:
+		min_loss = checkpoint['min_loss']
+		optimizer.load_state_dict(checkpoint['optimizer'])
+
+	best_test_loss = np.inf
+
+	for epoch in range(args.start_epoch, args.epochs):
+		train_loss = train_one_epoch(args.device, model, train_loader, optimizer)
+		test_loss = test_one_epoch(args.device, model, test_loader)
+
+		if test_loss<best_test_loss:
+			best_test_loss = test_loss
+			snap = {'epoch': epoch + 1,
+					'model': model.state_dict(),
+					'min_loss': best_test_loss,
+					'optimizer' : optimizer.state_dict(),}
+			torch.save(snap, 'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
+			torch.save(model.state_dict(), 'checkpoints/%s/models/best_model.t7' % (args.exp_name))
+			torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/best_ptnet_model.t7' % (args.exp_name))
+
+		torch.save(snap, 'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
+		torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % (args.exp_name))
+		torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/ptnet_model.t7' % (args.exp_name))
+		
+		boardio.add_scalar('Train Loss', train_loss, epoch+1)
+		boardio.add_scalar('Test Loss', test_loss, epoch+1)
+		boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
+
+		textio.cprint('EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f'%(epoch+1, train_loss, test_loss, best_test_loss))
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_pnlk', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--fine_tune_pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('--transfer_ptnet_weights', default='./checkpoints/exp_classifier/models/best_ptnet_model.t7', type=str,
+						metavar='PATH', help='path to pointnet features file')
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=10, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--epochs', default=200, type=int,
+						metavar='N', help='number of total epochs to run')
+	parser.add_argument('--start_epoch', default=0, type=int,
+						metavar='N', help='manual epoch number (useful on restarts)')
+	parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'],
+						metavar='METHOD', help='name of an optimizer (default: Adam)')
+	parser.add_argument('--resume', default='', type=str,
+						metavar='PATH', help='path to latest checkpoint (default: null (no-use))')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	torch.backends.cudnn.deterministic = True
+	torch.manual_seed(args.seed)
+	torch.cuda.manual_seed_all(args.seed)
+	np.random.seed(args.seed)
+
+	boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+	_init_(args)
+
+	textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+	textio.cprint(str(args))
+
+	
+	trainset = RegistrationData('PointNetLK', ModelNet40Data(train=True))
+	testset = RegistrationData('PointNetLK', ModelNet40Data(train=False))
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=8, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	ptnet = PointNet(emb_dims=args.emb_dims, use_bn=True)
+	
+	if args.transfer_ptnet_weights and os.path.isfile(args.transfer_ptnet_weights):
+		ptnet.load_state_dict(torch.load(args.transfer_ptnet_weights, map_location='cpu'))
+
+	if args.fine_tune_pointnet == 'tune':
+		pass
+	elif args.fine_tune_pointnet == 'fixed':
+		for param in ptnet.parameters():
+			param.requires_grad_(False)
+
+	model = PointNetLK(feature_model=ptnet)
+	model = model.to(args.device)
+
+	checkpoint = None
+	if args.resume:
+		assert os.path.isfile(args.resume)
+		checkpoint = torch.load(args.resume)
+		args.start_epoch = checkpoint['epoch']
+		model.load_state_dict(checkpoint['model'])
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	if args.eval:
+		test(args, model, test_loader, textio)
+	else:
+		train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/train_dcp.py

@@ -0,0 +1,249 @@
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import DGCNN, DCP
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def _init_(args):
+	if not os.path.exists('checkpoints'):
+		os.makedirs('checkpoints')
+	if not os.path.exists('checkpoints/' + args.exp_name):
+		os.makedirs('checkpoints/' + args.exp_name)
+	if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
+		os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
+	os.system('cp train_dcp.py checkpoints' + '/' + args.exp_name + '/' + 'train.py.backup')
+
+class IOStream:
+	def __init__(self, path):
+		self.f = open(path, 'a')
+
+	def cprint(self, text):
+		print(text)
+		self.f.write(text + '\n')
+		self.f.flush()
+
+	def close(self):
+		self.f.close()
+
+def get_transformations(igt):
+	R_ba = igt[:, 0:3, 0:3]								# Ps = R_ba * Pt
+	translation_ba = igt[:, 0:3, 3].unsqueeze(2)		# Ps = Pt + t_ba
+	R_ab = R_ba.permute(0, 2, 1)						# Pt = R_ab * Ps
+	translation_ab = -torch.bmm(R_ab, translation_ba)	# Pt = Ps + t_ab
+	return R_ab, translation_ab, R_ba, translation_ba
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+		transformations = get_transformations(igt)
+		transformations = [t.to(device) for t in transformations]
+		R_ab, translation_ab, R_ba, translation_ba = transformations
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		identity = torch.eye(3).cuda().unsqueeze(0).repeat(template.shape[0], 1, 1)
+		loss_val = torch.nn.functional.mse_loss(torch.matmul(output['est_R'].transpose(2, 1), R_ab), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t'], translation_ab[:,:,0])
+
+		cycle_loss = torch.nn.functional.mse_loss(torch.matmul(output['est_R_'].transpose(2, 1), R_ba), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t_'], translation_ba[:,:,0])
+		loss_val = loss_val + cycle_loss * 0.1
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader, textio):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+	textio.cprint('Validation Loss: %f & Validation Accuracy: %f'%(test_loss, test_accuracy))
+
+def train_one_epoch(device, model, train_loader, optimizer):
+	model.train()
+	train_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(train_loader)):
+		template, source, igt = data
+		transformations = get_transformations(igt)
+		transformations = [t.to(device) for t in transformations]
+		R_ab, translation_ab, R_ba, translation_ba = transformations
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		identity = torch.eye(3).cuda().unsqueeze(0).repeat(template.shape[0], 1, 1)
+		loss_val = torch.nn.functional.mse_loss(torch.matmul(output['est_R'].transpose(2, 1), R_ab), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t'], translation_ab[:,:,0])
+
+		cycle_loss = torch.nn.functional.mse_loss(torch.matmul(output['est_R_'].transpose(2, 1), R_ba), identity) \
+			   + torch.nn.functional.mse_loss(output['est_t_'], translation_ba[:,:,0])
+		loss_val = loss_val + cycle_loss * 0.1
+		# print(loss_val.item())
+
+		# forward + backward + optimize
+		optimizer.zero_grad()
+		loss_val.backward()
+		optimizer.step()
+
+		train_loss += loss_val.item()
+		count += 1
+
+	train_loss = float(train_loss)/count
+	return train_loss
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+	learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+	if args.optimizer == 'Adam':
+		optimizer = torch.optim.Adam(learnable_params)
+	else:
+		optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+	if checkpoint is not None:
+		min_loss = checkpoint['min_loss']
+		optimizer.load_state_dict(checkpoint['optimizer'])
+
+	best_test_loss = np.inf
+
+	for epoch in range(args.start_epoch, args.epochs):
+		train_loss = train_one_epoch(args.device, model, train_loader, optimizer)
+		test_loss = test_one_epoch(args.device, model, test_loader)
+
+		if test_loss<best_test_loss:
+			best_test_loss = test_loss
+			snap = {'epoch': epoch + 1,
+					'model': model.state_dict(),
+					'min_loss': best_test_loss,
+					'optimizer' : optimizer.state_dict(),}
+			torch.save(snap, 'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
+			torch.save(model.state_dict(), 'checkpoints/%s/models/best_model.t7' % (args.exp_name))
+			torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/best_ptnet_model.t7' % (args.exp_name))
+
+		torch.save(snap, 'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
+		torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % (args.exp_name))
+		torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/ptnet_model.t7' % (args.exp_name))
+		
+		boardio.add_scalar('Train Loss', train_loss, epoch+1)
+		boardio.add_scalar('Test Loss', test_loss, epoch+1)
+		boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
+
+		textio.cprint('EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f'%(epoch+1, train_loss, test_loss, best_test_loss))
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_dcp', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--epochs', default=200, type=int,
+						metavar='N', help='number of total epochs to run')
+	parser.add_argument('--start_epoch', default=0, type=int,
+						metavar='N', help='manual epoch number (useful on restarts)')
+	parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'],
+						metavar='METHOD', help='name of an optimizer (default: Adam)')
+	parser.add_argument('--resume', default='', type=str,
+						metavar='PATH', help='path to latest checkpoint (default: null (no-use))')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	torch.backends.cudnn.deterministic = True
+	torch.manual_seed(args.seed)
+	torch.cuda.manual_seed_all(args.seed)
+	np.random.seed(args.seed)
+
+	boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+	_init_(args)
+
+	textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+	textio.cprint(str(args))
+
+	
+	trainset = RegistrationData('DCP', ModelNet40Data(train=True))
+	testset = RegistrationData('DCP', ModelNet40Data(train=False))
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create PointNet Model.
+	dgcnn = DGCNN(emb_dims=args.emb_dims)
+	model = DCP(feature_model=dgcnn, cycle=True)
+	model = model.to(args.device)
+
+	checkpoint = None
+	if args.resume:
+		assert os.path.isfile(args.resume)
+		checkpoint = torch.load(args.resume)
+		args.start_epoch = checkpoint['epoch']
+		model.load_state_dict(checkpoint['model'])
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	if args.eval:
+		test(args, model, test_loader, textio)
+	else:
+		train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == '__main__':
+	main()

thirdparty/learning3d/examples/train_deepgmr.py

@@ -0,0 +1,244 @@
+import open3d as o3d
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import DeepGMR
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def display_open3d(template, source, transformed_source):
+	template_ = o3d.geometry.PointCloud()
+	source_ = o3d.geometry.PointCloud()
+	transformed_source_ = o3d.geometry.PointCloud()
+	template_.points = o3d.utility.Vector3dVector(template)
+	source_.points = o3d.utility.Vector3dVector(source + np.array([0,0,0]))
+	transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
+	template_.paint_uniform_color([1, 0, 0])
+	source_.paint_uniform_color([0, 1, 0])
+	transformed_source_.paint_uniform_color([0, 0, 1])
+	o3d.visualization.draw_geometries([template_, source_, transformed_source_])
+
+def rotation_error(R, R_gt):
+	cos_theta = (torch.einsum('bij,bij->b', R, R_gt) - 1) / 2
+	cos_theta = torch.clamp(cos_theta, -1, 1)
+	return torch.acos(cos_theta) * 180 / math.pi
+
+def translation_error(t, t_gt):
+	return torch.norm(t - t_gt, dim=1)
+
+def rmse(pts, T, T_gt):
+	pts_pred = pts @ T[:, :3, :3].transpose(1, 2) + T[:, :3, 3].unsqueeze(1)
+	pts_gt = pts @ T_gt[:, :3, :3].transpose(1, 2) + T_gt[:, :3, 3].unsqueeze(1)
+	return torch.norm(pts_pred - pts_gt, dim=2).mean(dim=1)
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	rotation_errors, translation_errors, rmses = [], [], []
+
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+
+		eye = torch.eye(4).expand_as(igt).to(igt.device)
+		mse1 = F.mse_loss(output['est_T_inverse'] @ torch.inverse(igt), eye)
+		mse2 = F.mse_loss(output['est_T'] @ igt, eye)
+		loss = mse1 + mse2
+
+		r_err = rotation_error(est_T_inverse[:, :3, :3], igt[:, :3, :3])
+		t_err = translation_error(est_T_inverse[:, :3, 3], igt[:, :3, 3])
+		rmse_val = rmse(template[:, :100], est_T_inverse, igt)
+		rotation_errors.append(r_err)
+		translation_errors.append(t_err)
+		rmses.append(rmse_val)
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	print("Mean rotation error: {}, Mean translation error: {} and Mean RMSE: {}".format(np.mean(rotation_errors), np.mean(translation_errors), np.mean(rmses)))
+	return test_loss
+
+def test(args, model, test_loader, textio):
+	test_loss = test_one_epoch(args.device, model, test_loader)
+	textio.cprint('Validation Loss: %f'%(test_loss))
+
+def train_one_epoch(device, model, train_loader, optimizer):
+	model.train()
+	train_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(train_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		
+		eye = torch.eye(4).expand_as(igt).to(igt.device)
+		mse1 = F.mse_loss(output['est_T_inverse'] @ torch.inverse(igt), eye)
+		mse2 = F.mse_loss(output['est_T'] @ igt, eye)
+		loss = mse1 + mse2
+
+		# forward + backward + optimize
+		optimizer.zero_grad()
+		loss_val.backward()
+		optimizer.step()
+
+		train_loss += loss_val.item()
+		count += 1
+
+	train_loss = float(train_loss)/count
+	return train_loss
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+	learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+	if args.optimizer == 'Adam':
+		optimizer = torch.optim.Adam(learnable_params)
+	else:
+		optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+	if checkpoint is not None:
+		min_loss = checkpoint['min_loss']
+		optimizer.load_state_dict(checkpoint['optimizer'])
+
+	best_test_loss = np.inf
+
+	for epoch in range(args.start_epoch, args.epochs):
+		train_loss = train_one_epoch(args.device, model, train_loader, optimizer)
+		test_loss = test_one_epoch(args.device, model, test_loader)
+
+		if test_loss<best_test_loss:
+			best_test_loss = test_loss
+			snap = {'epoch': epoch + 1,
+					'model': model.state_dict(),
+					'min_loss': best_test_loss,
+					'optimizer' : optimizer.state_dict(),}
+			torch.save(snap, 'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
+			torch.save(model.state_dict(), 'checkpoints/%s/models/best_model.t7' % (args.exp_name))
+			torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/best_ptnet_model.t7' % (args.exp_name))
+
+		torch.save(snap, 'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
+		torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % (args.exp_name))
+		torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/ptnet_model.t7' % (args.exp_name))
+		
+		boardio.add_scalar('Train Loss', train_loss, epoch+1)
+		boardio.add_scalar('Test Loss', test_loss, epoch+1)
+		boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
+
+		textio.cprint('EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f'%(epoch+1, train_loss, test_loss, best_test_loss))
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_deepgmr', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	parser.add_argument('--nearest_neighbors', default=20, type=int,
+						metavar='K', help='No of nearest neighbors to be estimated.')
+	parser.add_argument('--use_rri', default=True, type=bool,
+						help='Find nearest neighbors to estimate features from PointNet.')
+
+	# settings for on training
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+	parser.add_argument('--epochs', default=200, type=int,
+						metavar='N', help='number of total epochs to run')
+	parser.add_argument('--start_epoch', default=0, type=int,
+						metavar='N', help='manual epoch number (useful on restarts)')
+	parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'],
+						metavar='METHOD', help='name of an optimizer (default: Adam)')
+	parser.add_argument('--resume', default='', type=str,
+						metavar='PATH', help='path to latest checkpoint (default: null (no-use))')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	if args.nearest_neighbors > 0:
+		args.use_rri = True
+	return args
+
+def main():
+	args = options()
+	torch.backends.cudnn.deterministic = True
+	torch.manual_seed(args.seed)
+	torch.cuda.manual_seed_all(args.seed)
+	np.random.seed(args.seed)
+
+	boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+	_init_(args)
+
+	textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+	textio.cprint(str(args))
+	
+	trainset = RegistrationData('DeepGMR', ModelNet40Data(train=True), additional_params={'nearest_neighbors': args.nearest_neighbors})
+	testset = RegistrationData('DeepGMR', ModelNet40Data(train=False), additional_params={'nearest_neighbors': args.nearest_neighbors})
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	model = DeepGMR(use_rri=args.use_rri, nearest_neighbors=args.nearest_neighbors)
+	model = model.to(args.device)
+
+	checkpoint = None
+	if args.resume:
+		assert os.path.isfile(args.resume)
+		checkpoint = torch.load(args.resume)
+		args.start_epoch = checkpoint['epoch']
+		model.load_state_dict(checkpoint['model'])
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained), strict=False)
+	model.to(args.device)
+
+	if args.eval:
+		test(args, model, test_loader, textio)
+	else:
+		train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == '__main__':
+	main()