Awesome

3DD-TTA

Official PyTorch Implementation of the Paper "Test-Time Adaptation of 3D Point Clouds via Denoising Diffusion Models"

Test-time adaptation (TTA) of 3D point clouds is essential for addressing discrepancies between training and testing samples, particularly in corrupted point clouds like those from LiDAR data. Adapting models online to distribution shifts is crucial, as training for every variation is impractical. Existing methods often fine-tune models using self-supervised learning or pseudo-labeling, which can result in forgetting source domain knowledge.

We propose a training-free, online 3D TTA method called 3DD-TTA (3D Denoising Diffusion Test-Time Adaptation), which adapts input point clouds using a diffusion strategy while keeping the source model unchanged. A Variational Autoencoder (VAE) encodes corrupted point clouds into latent spaces, followed by a denoising diffusion process.

Key Features:

• No fine-tuning required: Keeps the source model intact.
• Diffusion strategy: Updates the corrupted latent points for alignment with the source domain.
• Strong performance: Achieves state-of-the-art results on ShapeNet, ModelNet40, and ScanObjectNN.

Install:

# Create a new Conda environment named "3dd_tta_env" with Python 3.8
conda create --name 3dd_tta_env python=3.8
conda activate 3dd_tta_env

# Install PyTorch, torchvision, and torchaudio with CUDA 12.1 support
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Install all dependencies from requirements.txt
pip install -r requirements.txt

# Compile and install the Earth Mover's Distance (EMD) extension (used for point cloud comparison)
cd ./extensions/emd
python setup.py install --user
cd ../..

# Install the PointNet++ operations library (required for point cloud processing)
cd Pointnet2_PyTorch/pointnet2_ops_lib
pip install .
cd ../..

# Install KNN_CUDA (GPU-accelerated k-nearest neighbor functionality)
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl

# Install OpenAI's CLIP model (used for vision-language tasks)
pip install git+https://github.com/openai/CLIP.git 

# Build and package the project
python build_pkg.py

Data Preparation

Our code supports three datasets: ModelNet40, ShapeNetCore, and ScanObjectNN. To prepare the data, please follow the instructions in the preparation guide.

Obtaining Pre-trained Source Models

You can download the source model (PointMAE) pretrained on Modelnet40, ShapeNet, and ScanObjectNN from here. Put the PointMAE checkpoints inside pointnet_ckpts folder. To download the pretrained diffusion model you can use the following link. Put the diffusion checkpoints inside lion_ckpts folder.

Qualitative Evaluation of 3dd_tta Model

To evaluate the model qualitatively, use the following command:

python demo_3dd_tta.py --diff_ckpt=./lion_ckpts/epoch_10999_iters_2100999.pt --denoising_step=35 --dataset_root=./data/modelnet40_c --corruption=background --sample_id=11

• Change the corruption argument to the OOD type you desire.
• For background noise, choose higher denoising steps (e.g., 35) using the denoising_step argument.
• For other corruption types, use lower denoising steps like 5 or 10.

This code will:

1. Read the sample determined by the sample_id and corrupted by the specified corruption type.
2. Adapt it back to the source domain.

The output will automatically be saved in the ./outputs/qualitative folder.

Quantitative Evaluation of 3dd_tta Model

To evaluate the model quantitatively, use the following command:

python main_3dd_tta.py --batch_size=32 --pointmae_ckpt=./pointnet_ckpts/modelnet_jt.pth --diff_ckpt=./lion_ckpts/epoch_10999_iters_2100999.pt --dataset_name=modelnet-c --dataset_root=./data/modelnet40_c

Quantitative Evaluation Instructions

• Use the --batch_size argument to modify the batch size according to your requirements.
• Use the --dataset_name argument to select the dataset. Options include:
  • modelnet-c
  • shapenet-c
  • scanobjectnn-c
• Ensure the --dataset_root argument is updated to the appropriate root directory of the selected dataset.
• The script will automatically evaluate the model across all corruption types and generate a report. The results will be saved in: ./outputs/quantitative

Results:

Our method demonstrates superior generalization across multiple datasets, including ShapeNet, ModelNet40 and ScanObjectNN.

Classification accuracies on ShapeNet-c

Classification accuracies on ModelNet40-c