Skip to content
Code for PartNet: A Recursive Part Decomposition Network for Fine-grained and Hierarchical Shape Segmentation
Branch: master
Clone or download
Latest commit 5b3f8d5 May 15, 2019
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
data/airplane initial Apr 22, 2019
models update Apr 22, 2019
picture initial Apr 22, 2019
pytorch_ops Update build.py May 10, 2019
.gitignore Initial commit Apr 16, 2019
LICENSE Initial commit Apr 16, 2019
README.md Update README.md May 9, 2019
ap_evaluate.py initial Apr 22, 2019
dataloader.py initial Apr 22, 2019
part_color_mapping.json initial Apr 22, 2019
partnet.py initial Apr 22, 2019
pointnet2.py initial Apr 22, 2019
test_demo.py variable name update Apr 23, 2019
torchfoldext.py initial Apr 22, 2019
train.py initial Apr 22, 2019
util.py Update util.py May 10, 2019

README.md

PartNet

Code for PartNet: A Recursive Part Decomposition Network for Fine-grained and Hierarchical Shape Segmentation

Intruduction

Deep learning approaches to 3D shape segmentation are typically formulated as a multi-class labeling problem. Existing models are trained for a fixed set of labels, which greatly limits their flexibility and adaptivity. We opt for topdown recursive decomposition and develop the first deep learning model for hierarchical segmentation of 3D shapes, based on recursive neural networks. Starting from a full shape represented as a point cloud, our model performs recursive binary decomposition, where the decomposition network at all nodes in the hierarchy share weights. At each node, a node classifier is trained to determine the type (adjacency or symmetry) and stopping criteria of its decomposition. The features extracted in higher level nodes are recursively propagated to lower level ones. Thus, the meaningful decompositions in higher levels provide strong contextual cues constraining the segmentations in lower levels. Meanwhile, to increase the segmentation accuracy at each node, we enhance the recursive contextual feature with the shape feature extracted for the corresponding part. Our method segments a 3D shape in point cloud into an unfixed number of parts, depending on the shape complexity, showing strong generality and flexibility. It achieves the stateof- the-art performance, both for fine-grained and semantic segmentation, on the public benchmark and a new benchmark of fine-grained segmentation proposed in this work. We also demonstrate its application for fine-grained part refinements in image-to-shape reconstruction.

Dependencies

Requirements:

  • Python 3.5 with numpy, scipy, torchfold, tensorboard and etc.
  • PyTorch

Our code has been tested with Python 3.5, PyTorch 0.4.0, CUDA 8.0 on Ubuntu 16.04.

Datasets and Pre-trained weights

The input pointcloud and training hierarchical trees are on Here. Each category contains following folds:

  • models_2048_points_normals: normalized input shape pointcloud with normal
  • training_data_models_segment_2048_normals: GT part pointcloud with normal. (Note that we only store half parts of a shape, the other parts needs to be recovered by symmetric parameters.)
  • training_trees: hierarchical structure(ops.mat) with symmetric parameters(syms.mat). The labels(labels.mat) are for testing.

Our dataset has been extended and updated these days:

Datasets information

category_name chair airplane table sofa helicopter bike
number of shapes 999 630 583 630 100 155
number of parts 9697 5234 3207 4747 1415 1238
maximum parts per shape 25 14 17 27 21 9
minimum parts per shape 3 4 2 2 6 6

The Pre-trained model are on Here.

Usage: Demo

Require 3GB RAM on the GPU and 5sec to run.

This script takes as input a normalized 2048 pointcloud with normal (Sampled from ShapeNet model using virtualscanner).

Please download Pre-trained weights of airplane first and put it at ./models/airplane.

python test_demo.py

input

Usage: Training

Put data of each category in ./data/category_name(eg ./data/airplane)

Build extention for each op in ./pytorch_ops/*** (eg ./pytorch_ops/sampling/) using build.py

ps: torch 0.4 is required, and it won't work with later torch version.

python build.py

Then run training process

python train.py

More training and testing arguments are set in util.py

'--epochs' (number of epochs; default=1000)
'--batch_size' (batch size; default=10)
'--show_log_every' (show training log for every X frames; default=3)
'--no_cuda' (don't use cuda)
'--gpu' (device id of GPU to run cuda)
'--data_path' (dataset path, default='data')
'--save_path' (trained model path, default='models')
'--output_path' (segmented result path, default='results')
'--training' (training or testing, default=False)
'--split_num' (training data size for each category)
'--total_num' (full data size for each category, only for testing)
'--label_category' (semantic labels for each category, only for testing)

Usage: Testing

To evaluate AP correctly, you need to set label_category for each category in util.py

'--label_category' (semantic labels for each category, only for testing)

We measure AP(%) with IoU threshold being 0.25 and 0.5, respectively.

python ap_evaluate.py

Segmentation results and its corresponding GT can also be found in ./results/category_name(eg ./data/airplane).

AP information of our method on updated datasets and training/testing split.

category name chair airplane table sofa helicopter bike
IoU > 0.25 93.83 96.33 78.49 76.07 83.0 98.22
IoU > 0.5 84.23 88.41 63.2 55.76 69.4 97.60

PS: If you want to try more new shapes, please make sure that them are oriented and normalized as our shapes.

Citation

If you use this code, please cite the following paper.

@inproceedings{yu2019partnet,
    title = {PartNet: A Recursive Part Decomposition Network for Fine-grained and Hierarchical Shape Segmentation},
    author = {Fenggen Yu and Kun Liu and Yan Zhang and Chenyang Zhu and Kai Xu},
    booktitle = {CVPR},
    pages = {to appear},
    year = {2019}
}
You can’t perform that action at this time.