I2-Multimedia-Lab / RLGrid

PyTorch implementation for RLGrid: Reinforcement Learning Controlled Grid Deformation for Coarse-to-Fine Point Could Completion.

RLGrid: Reinforcement Learning Controlled Grid Deformation for Coarse-to-Fine Point Could Completion

This repository contains PyTorch implementation for RLGrid: Reinforcement Learning Controlled Grid Deformation for Coarse-to-Fine Point Could Completion.

RLGrid is a network that leverages reinforcement learning for 2D Grid Scale selection and can be plug-and-play on any operation that uses 2D Grid for coarse-to-fine completion.

Network

Usage

Pytorch >= 1.7.0

python >= 3.7

CUDA >= 11.0

GCC >= 4.9

tensorboardX

open3d

pyntcloud

conda create --name RLGrid --file requirements.txt

Building Pytorch Extensions for Chamfer Distance, Earth Mover's Distance

NOTE: PyTorch >= 1.7 and GCC >= 4.9 are required.

#Chamfer Distance

cd utils/Chamfer_dist

python setup.py install

#Earth Mover's Distance

cd utils/EMD

python setup.py install

Dataset

The details of used datasets can be found in DATASET.md.

Training

To train a point cloud completion model from scratch, run:

#train AE
python train_AE.py

#train sc-GAN
python pretrain_treegan.py

#train RL Agent
python train_RL_Agent.py

Pretrained model

The pretrained model will be coming soon.

Visualization

Some Results on PCN Dataset

Some intermediate process (The green box represents that the better coarse output is generated by sc-GAN, while the blue box represents that the better coarse output is generated by AE.)