Point clouds are a set of data points in space to represent the 3D geometry of objects. A fundamental step in the processing is to achieve segmentation of the point cloud at different levels of detail. Within this context, hierarchical clustering (HC) breaks the point cloud down into coherent subsets to recognize the parts that make up the object. Along with classic approaches that build a hierarchical tree bottom-up using linkage criteria, recent developments exploit the tree-likeness of hyperbolic metric space, embedding data into the Poincaré Ball and capturing a hierarchical structure with low distortion. The main advantage of this kind of solution is the possibility to explore the space of discrete binary trees using continuous optimization. However, in this framework, a similarity function between points is assumed to be known, while this cannot always be granted for point cloud applications. In our method, we propose to use metric learning to fit at the same time the good similarity function and the optimal embedding into the hyperbolic space. Furthermore, when arbitrary rotations are applied to a 3D object, the pose should not influence the segmentation quality. Therefore, to avoid extensive data augmentation, we impose rotation invariance to ensure the uniqueness of the hierarchical segmentation of point clouds. We show the performance of our method on two datasets, ShapeNet and PartNet, at different levels of granularity. The results obtained are promising when compared to state-of-the-art flat segmentation.
Use the following command to create an Anaconda environment containing all the dependecies
conda env create -f hpcs-env.yaml
Before training/testing models you need to download ShapeNet dataset from HERE.
Please use the following structure to organize data
HPCS/
...
data/
...
ShapeNet/
raw/
...
...
Similarly, download PartNet data from HERE and use the following structure for PartNet data
HPCS/
...
data/
PartNet/
sem_seg_h5/
...
...
bash run_train_shapenet.sh
or
python train.py --dataset 'shapenet' \
--fixed_points 1024 \
--model 'vn_dgcnn_partseg' \
--eucl_embedding 32 \
--hyp_embedding 32 \
--k 20 \
--margin 0.35 \
--t_per_anchor 50 \
--temperature 0.05 \
--epochs 50 \
--lr 0.05 \
--accelerator 'gpu' \
--trade_off 0.10
To train on partnet, please choose one between the following categories:
- Bed
- Bottle
- Chair
- Clock
- Dishwasher
- Display
- Door
- Earphone
- Faucet
- Knife
- Lamp
- Microwave
- Refrigerator
- StorageFurniture
- Table
- TrashCan
- Vase
bash run_train_partnet.sh
or
python train.py --dataset 'partnet' \
--fixed_points 1024 \
--category 'Bed' \
--level 3 \
--model 'vn_dgcnn_partseg' \
--eucl_embedding 32 \
--hyp_embedding 32 \
--k 20 \
--margin 0.35 \
--t_per_anchor 50 \
--temperature 0.05 \
--epochs 50 \
--lr 0.05 \
--accelerator 'gpu' \
--trade_off 0.10
To infer on Shapenet data using the pretrained model, please
python infer.py shapenet --model_path checkpoints/shapenet/model.ckpt --test_batches 10 --plot
Please remove the flag --plot
to avoid visualization of results.
To test on partnet, please choose one between the following categories:
- Bed
- Bottle
- Chair
- Clock
- Dishwasher
- Display
- Door
- Earphone
- Faucet
- Knife
- Lamp
- Microwave
- Refrigerator
- StorageFurniture
- Table
- TrashCan
- Vase
And choose between level
from 1 to 3. For example to test on 2
examples from category=Bottle
at level=3
, please run
python infer.py partnet --model_path checkpoints/partnet/Bottle/model.ckpt --level 3 --category bottle --test_batches 2 --plot