scripts/eval_modelnet40.py

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)