czq142857 / bae-net Goto Github PK

The code for paper "BAE-NET: Branched Autoencoder for Shape Co-Segmentation".

License: Other

Python 74.08% Batchfile 25.92%

bae-net's Introduction

BAE-NET

The tensorflow code for paper "BAE-NET: Branched Autoencoder for Shape Co-Segmentation", Zhiqin Chen, Kangxue Yin, Matthew Fisher, Siddhartha Chaudhuri, Hao (Richard) Zhang.

Paper | Supplementary Material

Update

Added a post-processing step to all testing functions, controlled by the flag use_post_processing. The default is set to True. The post-processing step detects query points that have close-to-zero outputs at all branches. Those points are not inside any predicted parts, therefore their labels should be considered as invalid ones. We assign the labels of those points by copying the labels of the closest points with valid labels.

Citation

If you find our work useful in your research, please consider citing:

@article{chen2019bae_net,
  title={BAE-NET: Branched Autoencoder for Shape Co-Segmentation},
  author={Zhiqin Chen and Kangxue Yin and Matthew Fisher and Siddhartha Chaudhuri and Hao Zhang},
  journal={Proceedings of International Conference on Computer Vision (ICCV)},
  year={2019}
}

Dependencies

Requirements:

Python 3.5 with numpy, scipy and h5py
Tensorflow
PyMCubes (optional, for marching cubes)

Our code has been tested with Python 3.5, TensorFlow 1.9.0, CUDA 9.1 and cuDNN 7.0 on Ubuntu 16.04.

It has also been tested on Windows 10 but something went wrong. If sigmoid is placed before reduce_max, sigmoid won't be executed in certain cases. The solution is to change the last few layers of the decoder from "linear - sigmoid - reduce_max" to "linear - reduce_max - sigmoid".

Datasets and weights

We use the same point sampling method as in IM-NET. For data preparation, please see directory point_sampling.

We provide the ready-to-use ShapeNet dataset, together with our trained weights for one-shot training with 1/2/3 exemplars. If you would like to try our 4-layer model, please uncomment "level 2_2" of the generator in "model.py".

Backup links:

Training

To perform one-shot training, use the following command:

python main.py --train --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot_1_3000 --supervision_list ref1.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

The above command will train the model (with L1 regularization) 200000 iterations after 3000 iterations of supervised-loss-only pretraining. It will do one supervised PASS (training all shapes in supervision_list using supervised loss) every 4 iterations of unsupervised training. The file specified by "--supervision_list" should include the exemplars for supervised loss. "--real_size 32 --points_per_shape 8192" means during training each target shape has 8192 sampled points from its 32^3 voxel model.

You can run the batch files "train_1shot.bat", "train_2shot.bat" and "train_3shot.bat" to train on all categories in ShapeNet.

sh train_1shot.bat

(or simply double click it if you are using Windows.)

To perform unsupervised training, please remove "--supervised":

python main.py --train --L1reg --iteration 200000 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_unsup --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

In default settings the decoder will have 8 branches.

Evaluation and Visualization

To get the mean IOU on test shapes, replace "--train" with "--iou":

python main.py --iou --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot --supervision_list ref1.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

The IOU will be printed on the command prompt and saved to a txt in the checkpoint directory.

To visualize reconstructed shapes with colored segmentation, replace "--train" with "--recon" and specify the target shapes using "--supervision_list":

python main.py --recon --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot --supervision_list 03001627_test_vox.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

To visualize given point clouds with colored segmentation, replace "--train" with "--pointcloud" and specify the target shapes using "--supervision_list":

python main.py --pointcloud --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot --supervision_list 03001627_test_vox.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

To visualize given meshes with colored segmentation, please change the directory of ShapeNet (containing .obj meshes) in function "test_obj" at "model.py", then replace "--train" with "--mesh" and specify the target shapes using "--supervision_list":

python main.py --mesh --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot --supervision_list 03001627_test_vox.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

All visualization results are saved to folder "samples". Please see "test_reconstruction_1shot.bat", "test_pointcloud_1shot.bat", "test_mesh_1shot.bat" and "test_iou_1shot.bat" for more examples.

License

This project is licensed under the terms of the MIT license (see LICENSE for details).

bae-net's People

Contributors

Stargazers

Watchers

Forkers

peterzhousz hiyyg cv-ip gs14iitbbs apllolulu yunxiangzhang supriya-gdptl mamunctg flyskypigg krishia zoomkinseok harris0322 gsamful

bae-net's Issues

Could the network be trained in class-agnostic manner?

Hi Zhiqi,
First, thank you for the amazing work. I find the models are trained with each class, unlike the full-supervised network that provides one network for all the classes. So I wonder could this network be trained in a class-agnostic manner?

A question about the released model

Dear Dr Chen,

Thanks for the great work. I have a question about the released model. It seems only work for chair category, but for other categories like airplane, the result is noisy.

Could you provide pretrained models on the other categories?

Thanks with regards,
Zhengzhe

Other scores of ShapeNet unsupervised

Hello,

I really enjoyed your work, in particular the unsupervised segmentation part! Do you have the IOU score available for the unsupervised ShapeNet chair (back, seat, leg, arm) category, when it is trained without the table joint dataset?

Thanks in advance for your answer.

Greetings!

About unsupervised training of BAE-NET

Hello @czq142857,

Thanks for sharing the code.
I noticed that the commands you mentioned for unsupervised training of BAE-NET use all models in one category during the training step. Shouldn't it be using only train-set models instead of all models?
If it is by design (and not a typo), how did you decide the number of iterations the model should be trained for? How do you know the model is not overfitting?

Another question I have related to unsupervised training is how does the model learn part segmentation with just the mean square loss, especially the parts which are not connected to each other directly e.g. arms of a chair or engines of an airplane? How come one branch learns implicit field for both arms of a chair rather than two separate branches each one learning one arm?

Thank you,
Supriya

Reproducing The Results

Hi Zhiqin,

I'm trying to reproduce your results on Shapenet Part Segmentation dataset.

I cloned your code and downloaded your data. Then I trained the network under the one-example setting with the following parameters:

python main.py --train --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot_1_3000 --supervision_list ref1.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

Then I used the following code to test for IoU:

python main.py --iou --L1reg --supervised --iteration 200000 --pretrain_iters 3000 --retrain_iters 4 --dataset 03001627_vox --data_dir ./data/03001627_chair/ --checkpoint_dir checkpoint_1shot_1_3000 --supervision_list ref1.txt --sample_dir samples/03001627_chair  --real_size 32 --points_per_shape 8192

The results shows an average IoU of 59.5%. Although it does not directly corresponds to a number in your paper, it seems deviated from the average IoU reported in figure 9 for the chair category (which looks like in between 75% to 80%).

Can you kindly confirm if I did anything wrong?

Thanks a lot!

Best,
Xiangru

Unsupervised segmentation results

Hi Zhiqin,

Thank you for the great work! I'm trying to reproduce your unsupervised segmentation results. It seems like I can't get the results of the bag and cap categories the paper reported. The best results I got are 75.6 for bag category and 77.6 for cap category. Can you provide the parameters you set for unsupervised training, also the branch number and output shape of fully-connected layers? I'd apperate it if you can help.

Thank you

Code for mod-IoU

Hello,

how do you calculate the mod-IoU score as listed in Table 1 of the paper? I didn't find it in the code.

Greetings!

How to test with unsupervised?

hi, great work!!! could you please tell me how to test with unsupervised? thanks a lot!

Visualizing output for unsupervised trained BAE network

I am facing issues while visualizing the output for unsupervised trained BAE network.

I have trained the model on 02958343_car data using the following command:
python main.py --train --L1reg --iteration 20000 --dataset 02958343_vox --data_dir ../RAdata/02958343_car --checkpoint_dir checkpoint_unsup --sample_dir ../RAdata/02958343_car --real_size 32 --points_per_shape 8192

In order to view the output of the model, I am using this:
python main.py --recon --L1reg --iteration 20000 --dataset 02958343_vox --data_dir ../RAdata/02958343_car --checkpoint_dir checkpoint_unsup --sample_dir ../RAdata/02958343_car --real_size 32 --points_per_shape 8192

But this gives me the following error:

   for t in range(min(len(self.ref_voxels),16)):
AttributeError: 'IMSEG' object has no attribute 'ref_voxels'

Could you please guide me on how to debug this?
or suggest another way to visualize the output for unsupervised trained model?

Please note that I am able to visualize for the supervised trained model.

about test error!

hello,i am interested in your contribution,when i run the test code,an error is meeting,try to figure it out,but it is not OK,can you give me some advice how to solve this issue?

InternalError (see above for traceback): Blas SGEMM launch failed : m=1, n=128, k=16384
[[Node: encoder_1/conv_5/Conv3D = Conv3D[T=DT_FLOAT, data_format="NDHWC", dilations=[1, 1, 1, 1, 1], padding="VALID", strides=[1, 1, 1, 1, 1], _device="/job:localhost/replica:0/task:0/device:GPU:0"](encoder_1/Maximum_3, encoder/conv_5/Matrix/read)]]****

Looking forward to your early reply,thanks!
All the best.****

Point cloud as input

Hi,

Thanks for the insightful work!
I know that the method accepts: (i) data points with (2) their association 0 or 1 value for the inside or outside label, and (3) data voxel. I wonder, if my dataset only has point clouds data (not a mesh model), thus we cannot build data (2) and (3). In this case, can we still use or modify BAENet to do co-segmentation on my dataset (point clouds)? Many thanks!

Error calculating IOU

Hi:
When I train in an unsupervised way, the training phase is normal, but I encounter an error when I test the IOU, which prompts "imseg object has no attribute test" and how to solve this problem
Thanks

about the iou of unsupervised co-segmentation

hello, I run the unsupervised co-segmentation code recently. but i just get a very small iou value( such as 37 , 26.8). The same configuration gives me an IOU value of 82 on PointNet. so i want to know did you have any method to improve it? Thanks

about visualizing unsupervised training

hi, I'm trying to visualize colored point cloud segmentation for car dataset from shapeNet and something went wrong.
I used
python main.py --train --L1reg --iteration 200000 --dataset 02958343_vox --data_dir ./data/02958343_car/ --checkpoint_dir checkpoint_unsup --sample_dir samples/02958343_car --real_size 32 --points_per_shape 8192
for unsupervised training, and it worked well.
than I tried to visualized it with point cloud. I read some closed issues, and I changed name of 'obj_list.txt' in car dataset to '02958343_vox.txt'. Than I used
python main.py --pointcloud --L1reg --iteration 200000 --dataset 02958343_vox --data_dir ./data/02958343_car --checkpoint_dir checkpoint_unsup --supervision_list 02958343_vox.txt --sample_dir ./data/02958343_car --real_size 32 --points_per_shape 8192
for visualizing, but error occured

AttributeError: 'IMSEG' object has no attribute 'ref_b_point_num'

by the way, supervised learning and point cloud visualizing with chair dataset was working. This error occurs only occurs when I tried to visualize unsupervised-learned dataset.

It will be very helpful for me to have some guide.