PRS-Net

Introduction

A pytorch implementation of PRS-Net, a research work to detect planar reflective symmetry published on TVCG by Lin Gao et al.

Official link: http://geometrylearning.com/prs-net/

Author: D-Hank

Feel free to contact me at [email protected]. To use my code, please cite the link of this repository.

Dependencies

PCL 1.12.1: generate point cloud (poisson disk) {https://pointclouds.org/}

cuda_voxelizer: generate voxel (mesh split) {https://github.com/Forceflow/cuda_voxelizer}

Open3D: read pcd file in Python {http://www.open3d.org/}

Libigl: compute closest points on the mesh (barycentric coordinates) {https://github.com/libigl/libigl-python-bindings}

point_cloud_utils: read and write obj files {https://github.com/fwilliams/point-cloud-utils}

Mayavi / matplotlib: visualization {http://docs.enthought.com/mayavi/mayavi/}

Note: The library cuda voxelizer should be localized for our model. We include the revised version in the directory extern/

Directory

The project directory should be organized like this:

├── augment      # dataset after augmentation
├── data         # dataset after preprocess
├── shapenet     # original shapenet data
	├── 02691156
	└── ……
└── prsnet-repr  # main working directory
	├── checkpoint     # saved models
	└── ……

Architecture

We organize the project as below (arrow -> denotes the workflow).

  ----------------                                 --------------------------------------
  |  Preprocess  |  --------------------------->   |             PRSRunner              |
  ----------------                                 --------------------------------------
                                                         |                          |
                                                         |                          |
                                  ----------------------------------        ---------------------
                                  |            PRSModel            |   -->  |    Visualization  |
                                  ----------------------------------        ---------------------
                                   |          |         |         |
                             ---------  ---------  -----------  --------
                             |network|->|mlphead|->|transform|->| loss |
                             ---------  ---------  -----------  --------

Running Tips

Change your working directory to prsnet-repr. About 3 days and 80 GB free space are required. You can set the default options in settings.py.

To run this project from the start, first run python augment.py to generate augmented data. It takes one day to run on CPU.

Then run python preprocess.py to generate voxel, point cloud and closest points. We use 4 processes to run simultaneously, which takes around 2 days with CUDA acceleration. After that, run python check.py to check the completeness of data and generate bad model list.

Finally use python main.py to run the main program (train + test). It takes 0.5 hour to train. And for inference, the speed is around 2.3 ms per obj on laptop with entry-level GPU MX250.

If you'd like to use the pre-trained model in checkpoint/, then set CONTINUE in settings.py to be True and run main.py directly.

Results

For different categories in test set, we've achieved great results. Reflective plane (with coordinate axes on the left-buttom corner):

For generalized objects, the rotation axis:

Limitations and Improvement

• Position of the rotation axes

  Motivated by YuzhuoChen99's implementation, the model can only predict rotation axes near the original point. Even for the already-normalized shapenet dataset, the rotational center is not always near the origin. Therefore the model performs not so well (sometimes disturbed) in some categories.

  The solution is to introduce a shift vector or use the generalized 4×4 rotation matrix.

• Problems with axis-angle representation

  Basically, the network can learn to use tricks for better performance. That is, it can randomly pick three orthogonal axes and set the rotational angle to be $0$ or $2 \pi$. Then both the distance and regularized loss will be relatively low (it's truly a global minima, but there seems no mechanism to avoid this in the official release). So sometimes the training rotation loss looks like:

  

  After training for a long time, the network can get lazy for rotation. Our solution is to limit the rotation angle in a certain range, say ${[\pi / 6, \pi]}$. And there's a little revision in forward method of class MLPHead. We replace rot_out in line 89 of model.py by:

  lim_rot = torch.cat((torch.clamp(rot_out[ : , : , 0], self.min_cos, self.max_cos).unsqueeze(-1), rot_out[ : , : , 1 : ]), dim = -1)
  rot_out = functional.normalize(lim_rot, p = 2, dim = 2, eps = 1e-12)

  where we set self.min_cos = np.cos(np.pi / 2) and self.max_cos = np.cos(np.pi / 6 / 2) in MLPHead's constructor.

  Then the rotational sde won't collapse after several hours of training, after which we gain a more reasonable training loss (below)

  

  and more satisfying predictions for some categories than the original one (old on the left, new on the right).

  

Acknowledgement

Quaternion: tiny reimplementation of pytorch3D for Quaternion

Pairwise-Cosine: pytorch/pytorch#11202

Reference: He Yue's implementation {https://github.com/hysssb/PRS_net}

Reference: official release {https://github.com/IGLICT/PRS-Net}