This repository contains the implementation of O-CNN and Aadptive O-CNN introduced in our SIGGRAPH 2017 paper and SIGGRAPH Asia 2018 paper. The code is released under the MIT license.
We have released the TensorFlow-based implementation under the tf
branch, and our future development will be focused on this implementation.
If you would like to have a try with the beta version, please pull the code and run the following command: git checkout -b tf
.
-
O-CNN: Octree-based Convolutional Neural Networks
By Peng-Shuai Wang, Yang Liu, Yu-Xiao Guo, Chun-Yu Sun and Xin Tong
ACM Transactions on Graphics (SIGGRAPH), 36(4), 2017 -
Adaptive O-CNN: A Patch-based Deep Representation of 3D Shapes
By Peng-Shuai Wang, Chun-Yu Sun, Yang Liu and Xin Tong
ACM Transactions on Graphics (SIGGRAPH Asia), 37(6), 2018
If you use our code or models, please cite our paper.
@article {Wang-2017-OCNN,
title = {{O-CNN: Octree-based Convolutional Neural Networks for 3D Shape Analysis}},
author = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
journal = {ACM Transactions on Graphics (SIGGRAPH)},
volume = {36},
number = {4},
year = {2017},
}
@article {Wang-2018-AOCNN,
title = {{Adaptive O-CNN: A Patch-based Deep Representation of 3D Shapes}},
author = {Wang, Peng-Shuai and Sun, Chun-Yu and Liu, Yang and Tong, Xin},
journal = {ACM Transactions on Graphics (SIGGRAPH Asia)},
volume = {37},
number = {6},
year = {2018},
}
O-CNN is built upon the Caffe framework and it supports octree-based convolution, deconvolution, pooling, and unpooling. The code has been tested on the Windows 10 x64 (It can also be built on the Ubuntu 16.04). Its installation is as follows:
- Clone Caffe into the caffe directory with revision
6bfc5ca
:git clone https://github.com/BVLC/caffe.git && cd caffe && git checkout 6bfc5ca
. - Clone the code for O-CNN, then copy the code contained in the directory
O-CNN/caffe
into the caffe directory to override the original Caffe code. - Follow the installation instructions of Caffe to build the code to get the executive files
caffe
,convert_octree_data
andfeature_pooling
etc. Note that the code should be compiled with c++11 support, which is not enabled by default for some compilers. So please check your compiler settings and enable the compiler to support c++11. - Our O-CNN takes the octree representation of 3D objects as input. The code for converting a point cloud into octree representation is contained in the folder
O-CNN/ocnn/octree
, which can be built via cmake:cd O-CNN/ocnn/octree && mkdir build && cd build && cmake .. && cmake --build . --config Release
. - We also provide one tool to pre-process meshes from online dataset, which can be downloaded here.
After the building, you will get the executable files which is useful for conducting the experiments:
virtualscanner
- used to convert obj/off files to points filesoctree
- used to convert point files to octree filesconvert_octree_data
- used to convert octree files to lmdb filescaffe
- executable for training / evaluating modelsfeature_pooling
- pools features and outputs them to an lmdb
NOTE: Compared with the original code used in the experiments of the O-CNN paper, the code in this repository is refactored for the readability and maintainability, with the sacrifice of speed (it is about 10% slower, but it is more memory-efficient). If you want to try the original code or do some speed comparisons with our O-CNN
, feel free to drop me an email, we can share the original code with you.
A docker build file is provided to automatically build your environments so you don't have to worry about project dependencies. To get your environment up and running, execute the following:
cd docker
docker build --network=host --tag=ocnn .
docker run --runtime=nvidia --network=host --name=ocnn -it ocnn /bin/bash
You will now find yourself in a container environment where you can automatically prepare datasets and train/test an o-cnn.
Many thanks to David Pisani (@dapisani) for his contribution to the docker setup!
- Download and unzip the corresponding 3D model dataset (like the ModelNet40 dataset) into a folder.
- Convert all the models (in OBJ/OFF format) to dense point clouds with normals (in
POINTS
format) with the toolvirtualscanner
. This tool is specially designed for the 3D model from ModelNet40 and ShapeNet. If your mesh is of good quality, there are many more efficient methods to sample points.
Usage:
VirtualScanner <file_name> [nviews] [flags] [normalize]
file_name: the name of the file (*.obj; *.off) to be processed.
nviews: the number of views for scanning. Default: 6
flags: Indicate whether to output normal flipping flag. Default: 0
normalize: Indicate whether to normalize input mesh. Default: 0
Example:
VirtualScanner input.obj 14 // process the file input.obj
VirtualScanner D:\data\ 14 // process all the obj/off files under the folder D:\Data
- Run the tool
octree
to convert point clouds into the octree files.
Usage:
octree
--filenames <A file contains the absolute filenames of input POINTS each line>
[--adaptive <Build adaptive octree>=0]
[--adp_depth <The starting depth of adaptive octree>=4]
[--axis <The upright axis of the input model>=z]
[--depth <The maximum depth of the octree>=6]
[--full_depth <The full layer of the octree>=2]
[--key2xyz <Convert the key to xyz when serialization>=0]
[--node_dis <Output per-node displacement>=0]
[--node_feature <Compute per node feature>=0]
[--offset <The offset value for handing thin shapes>=0.55]
[--output_path <The output path>=.]
[--rot_num <Number of poses rotated along the upright axis>=12]
[--split_label <Compute per node splitting label>=0]
[--th_distance <The threshold for simplifying octree>=2.0f]
[--th_normal <The threshold for simplifying octree>=0.2f]
Example:
// process all the points into octrees of depth 5
octree --filenames filelist.txt --depth 5 --axis z
- Store all the octree files in a
leveldb
orlmdb
database by the toolconvert_octree_data
, which serves as the input of Caffe. For the ones who are new to the Caffe framework and do not know how to use the generated database to train and test the neural network, the tutorial on this page is highly recommanded.
Usage:
convert_octree_data <rootfolder> <listfile> <db_name>
rootfolder: base folder where db will be output and listed files are relative to
listfile: file which contains a list of each octree file to be added to the leveldb.
Each entry should be the [octree_file] [category_number]
db_name: Name of db to be outputted
Example:
convert_octree_data D:/octrees/ D:/octrees/list.txt D:/octrees_lmdb
The instruction to run the shape classification experiment:
- Download the ModelNet40 dataset, and convert it to a
lmdb
database as described above. Here we provide almdb
database with 5-depth octrees for convenience. Since we upgraded the octree format in this version of code, please run the following command to upgrade the lmdb:upgrade_octree_database.exe <input lmdb> <output lmdb>
. When manually generating the dataset, note that the upright direction of the 3D models in theModelNet40
isz
axis, so the octree command is:octree --filenames filelist.txt --depth 5 --axis z
- Download the
O-CNN
protocol buffer files, which are contained in the foldercaffe/examples/o-cnn
. - Configure the path of the database and run
caffe.exe
according to the instructions of Caffe. We also provide our pre-trained Caffe model incaffe/examples/o-cnn
.
The instruction to run the shape retrieval experiment:
-
Download the dataset from SHREC16, and convert it to a
lmdb
database as described above. Note that the upright direction of the 3D models in theShapeNet55
isY
axis, so the octree command is:octree --filenames filelist.txt --depth 5 --axis y
. Here we provide the lmdb databases with 5-depth octrees for convenience, just download the files prefixed withS55
and un-zip them. Since we upgraded the octree format in this version of code, please run the following command to upgrade the lmdb:upgrade_octree_database.exe <input lmdb> <output lmdb>
. -
Follow the same approach as the classification task to train the O-CNN with the
O-CNN
protocal filesS55_5.prototxt
andsolver_S55_5.prototxt
, which are contained in the foldercaffe/examples/o-cnn
. -
In the retrieval experiment, the
orientation pooling
is used to achieve better performance, which can be perfromed following the steps below.-
Generate feature for each object. For example, to generate the feature for the training data, open the file
S55_5.prototxt
, uncomment line 275~283, set thesource
in line 27 to thetraining lmdb
, set thebatch_size
in line 28 to 1, and run the following command.caffe.exe test --model=S55_5.prototxt --weights=S55_5.caffemodel --blob_prefix=feature/S55_5_train_ ^ --gpu=0 --save_seperately=false --iterations=[the training object number]
Similarly, the feature for the validation data and testing data can also be generated. Then we can get three binary files,
S55_5_train_feature.dat, S55_5_val_feature.dat and S55_5_test_feature.dat
, containing the features of the training, validation and testing data respectively.-
Pool the features of the same object. There are 12 features for each object since each object is rotated 12 times. We use max-pooling to merge these features.
feature_pooling.exe --feature=feature/S55_5_train_feature.dat --number=12 ^ --dbname=feature/S55_5_train_lmdb --data=[the data list file name]
Then we can get the feature of training, validation and testing data after pooling, contained in the
lmdb
databaseS55_5_train_lmdb
,S55_5_val_lmdb
andS55_5_test_lmdb
.-
Fine tune the
FC
layers of O-CNN, i.e. using thesolver_S55_5_finetune.prototxt
to re-train theFC
layers.caffe.exe train --solver=solver_S55_5_finetune.prototxt --weights=S55_5.caffemodel
-
Finally, dump the probabilities of each testing objects. Open the file
S55_5_finetune.prototxt
, uncomment the line 120 ~ 129, set thebatch_size
in line 27 to 1, change thesource
in line 26 tofeature/S55_5_test_lmdb
, and run the following command.caffe.exe test --model=S55_5_finetune.prototxt --weights=S55_5_finetune.caffemodel ^ --blob_prefix=feature/S55_test_ --gpu=0 --save_seperately=false --iterations=[...]
-
-
Use the matlab script
retrieval.m
, contained in the foldercaffe/examples/o-cnn
, to generate the final retrieval result. And evaluated it by the javascript code provided by SHREC16.
The instruction to run the segmentation experiment:
-
The original part annotation data is provided as the supplemental material of the work "A Scalable Active Framework for Region Annotation in 3D Shape Collections". As detailed in Section 5.3 of our paper, the point cloud in the original dataset is relatively sparse and the normal information is missing. We convert the sparse point clouds to dense points with normal information and correct part annotation. Here is one converted dataset for your convenience. Since we upgraded the octree format in this version of code, please run the following command to upgrade the lmdb:
upgrade_octree_database.exe --node_label <input lmdb> <output lmdb>
. And the dense point clouds withsegmentation labels
can be downloaded here. Again, before using them, upgrade with the following command:upgrade_points.exe --filenames <the txt file containing the filenames> --has_label 1
. -
Run the
octree.exe
to convert these point clouds to octree files. Note that you should set the parameterSegmentation
to 1 when running theoctree.exe
. Then you can get the octree files, which also contains the segmentation label. -
Convert the dataset to a
lmdb
database. Since the segmentation label is contained in each octree file, the object label for each octree file can be set to any desirable value. And the object label is just ignored in the segmentation task. -
Download the protocol buffer files, which are contained in the folder
caffe/examples/o-cnn
.NOTE:
as detailed in our paper, the training parameters are tuned and the pre-trained model from the retrieval task is used when the training dataset is relatively small. More details will be released soon. -
In the testing stage, the output label and probability of each finest leaf node can also be obtained. Specifically, open the file
segmentation_5.prototxt
, uncomment line 458~485, , set thebatch_size
in line 31 to 1, and run the following command to dump the result.caffe.exe test --model=segmentation_5.prototxt --weights=segmentation_5.caffemodel --gpu=0 --blob_prefix=feature/segmentation_5_test_ --binary_mode=false --save_seperately=true --iterations=[...]
-
For CRF refinement, please refer to the code provided here.
- Download the ModelNet40 dataset, and convert it to a
lmdb
database as described above. To generating the adaptive octree, the octree command changes to:octree --filenames filelist.txt --depth 5 --adaptive 1 --node_dis 1
.
- Download the protocol buffer files
cls_5.solver.prototxt
andcls_5.prototxt
, which are contained in the foldercaffe/examples/ao-cnn
. - Configure the path of the database and run
caffe.exe
according to the instructions of Caffe.
-
Download points with normals from this link, and download the rendered views from this link. Unzip them after downloading.
-
Run the script
O-CNN/caffe/examples/ao-cnn/dataset.py
to generate the lmdbs:oct_test_lmdb
,oct_train_aug_lmdb
,oct_train_lmdb
,img_test_lmdb
andimg_train_lmdb
, of whichoct_train_aug_lmdb
andoct_test_lmdb
are used in the autoencoder task. -
Download the protocol buffer files
O-CNN/caffe/examples/ao-cnn/ae_7_4.train.prototxt
andO-CNN/caffe/examples/ao-cnn/ae_7_4.solver.prototxt
. Configure the training and tesing lmdb files and training the network. -
After training, suppose the trained model is
autoencoder.caffemodel
, download the protocol buffer fileO-CNN/caffe/examples/ao-cnn/ae_7_4.test.prototxt
to test the network with the following command. And the output octrees are dumped into the folderae_output
. Then use the tools such asoctree2mesh
,octree2points
andpoints2ply
to convert the octree intoobj
files orply
files, which can be visualized via some 3D viewer such asMeshLab
.caffe.exe test --model=ae_7_4.test.prototxt --weights=autoencoder.caffemodel ^ --blob_prefix=ae_output/ae_test_ --gpu=0 --blob_header=false --iterations=[...]
- The lmdbs used in this experiment are
oct_test_lmdb
,oct_train_lmdb
,img_test_lmdb
andimg_train_lmdb
, which are generated by the scriptO-CNN/caffe/examples/ao-cnn/dataset.py
, as detailed in section 4.2. - Download the protocol buffer files
O-CNN/caffe/examples/ao-cnn/image2shape.train.prototxt
andO-CNN/caffe/examples/ao-cnn/image2shape.solver.prototxt
. Configure the training and tesing lmdb files and training the network. - After training, download the protocol buffer file
O-CNN/caffe/examples/ao-cnn/image2shape.test.prototxt
to testing the network. Follow the same procedure to dump and visualize the output as detailed in Section 4.2.
We thank the authors of ModelNet, ShapeNet and Region annotation dataset for sharing their 3D model datasets with the public.
Please contact us (Pengshuai Wang [email protected], Yang Liu [email protected] ) if you have any problem about our implementation or request to access all the datasets.