Update ReadMe

.gitignore

@@ -21,4 +21,5 @@ __pycache__
 *.fig
 log
 ae_output
-pytorch/projects/logs
+pytorch/projects/logs
+pytorch/projects/data

README.md

@@ -2,8 +2,7 @@
 
 <!-- ## Introduction <a name="introduction"></a> -->
 
-This repository contains the implementation of *O-CNN*  and  *Adaptive O-CNN* 
-introduced in our SIGGRAPH 2017 paper and SIGGRAPH Asia 2018 paper.  
+This repository contains the implementation of our papers related with *O-CNN*.  
 The code is released under the **MIT license**.
 
 - **[O-CNN: Octree-based Convolutional Neural Networks](https://wang-ps.github.io/O-CNN.html)**<br/>
@@ -31,8 +30,30 @@ AAAI Conference on Artificial Intelligence (AAAI), 2021. [Arxiv, 2020.08]<br/>
 If you use our code or models, please [cite](docs/citation.md) our paper.
 
 
+
+### Contents
+- [Installation](docs/installation.md)
+- [Data Preparation](docs/data_preparation.md)
+- [Shape Classification](docs/classification.md)
+- [Shape Retrieval](docs/retrieval.md)
+- [Shape Segmentation](docs/segmentation.md)
+- [Shape Autoencoder](docs/autoencoder.md)
+- [Shape Completion](docs/completion.md)
+- [Image2Shape](docs/image2shape.md)
+- [Unsupverised Pretraining](docs/unsupervised.md)
+- [ScanNet Segmentation](docs/scannet.md)
+
+
+
+
 ### What's New?
-- 2021.03.01: Update the code for pytorch-based O-CNN, including a ResNet and some important modules.
+- 2021.08.24: Update the code for pythorch-based O-CNN, including a UNet and
+  some other major components. Our vanilla implementation without any tricks on
+  [ScanNet](docs/scannet.md) dataset achieves 76.2 mIoU on the 
+  [ScanNet benchmark](http://kaldir.vc.in.tum.de/scannet_benchmark/), even surpassing the
+  recent state-of-art approaches published in CVPR 2021 and ICCV 2021.
+- 2021.03.01: Update the code for pytorch-based O-CNN, including a ResNet and
+  some important modules.
 - 2021.02.08: Release the code for ShapeNet segmentation with HRNet.
 - 2021.02.03: Release the code for ModelNet40 classification with HRNet.
 - 2020.10.12: Release the initial version of our O-CNN under PyTorch. The code
@@ -52,26 +73,6 @@ If you use our code or models, please [cite](docs/citation.md) our paper.
   benchmarks.
 
 
-
-### Contents
-- [Installation](docs/installation.md)
-- [Data Preparation](docs/data_preparation.md)
-- [Shape Classification](docs/classification.md)
-- [Shape Retrieval](docs/retrieval.md)
-- [Shape Segmentation](docs/segmentation.md)
-- [Shape Autoencoder](docs/autoencoder.md)
-- [Shape Completion](docs/completion.md)
-- [Image2Shape](docs/image2shape.md)
-- [Unsupverised Pretraining](docs/unsupervised.md)
-
-
-
-
-We thank the authors of [ModelNet](http://modelnet.cs.princeton.edu), 
-[ShapeNet](http://shapenet.cs.stanford.edu/shrec16/) and 
-[Region annotation dataset](http://cs.stanford.edu/~ericyi/project_page/part_annotation/index.html) 
-for sharing their 3D model datasets with the public.
-
-Please contact us (Pengshuai Wang [email protected], Yang Liu [email protected] ) 
+Please contact us (Peng-Shuai Wang [email protected], Yang Liu [email protected] ) 
 if you have any problems about our implementation.  
 

docs/classification.md

@@ -155,19 +155,52 @@ the data automatically.
 
 ## O-CNN on PyTorch
 
-1. Data preparation. Change the working directory to `O-CNN/pytorch/projects`
-   and follow the instructions [above](classification.md#o-cnn-on-tensorflow) to
-   the 3rd step. Place the preprocesed points files into the folder
-   `dataset/ModelNet40.points`.
+### Prepare the point cloud for ModelNet40   
+1. Change the working directory to `pytorch/projects`. Run the following command
+   to download [ModelNet40](http://modelnet.cs.princeton.edu/ModelNet40.zip)
+   dataset and unzip it to the folder `dataset/ModelNet40` via the following
+   command:
+    ```
+    python tools/modelnet.py --run download_m40
+    ```
+
+2. Convert triangle meshes (in `off` format) to point clouds (in `points` format)
+   with the [virtual_scanner](https://github.com/wang-ps/O-CNN/tree/master/virtual_scanner).
+   This process can be automatically executed by the following command.
+   Remember to provide actual `<The path of the virtual_scanner>` to run the command.
+    ```shell
+    python tools/modelnet.py --run m40_convert_mesh_to_points \
+                             --scanner <The path of the virtual_scanner>
+    ```
+
+3. The generated point clouds are very dense, and if you would like to save disk
+   spaces, you can optionally run the following command to simplify the point cloud.
+    ```shell
+    python tools/modelnet.py --run m40_simplify_points 
+    ```
+   We also provide the converted `point clouds` for convenience. Run the
+   following command to download the zip file
+   [here](https://www.dropbox.com/s/m233s9eza3acj2a/ModelNet40.points.zip?dl=0).
+   ```shell
+   python tools/modelnet.py --run download_m40_points
+   ```
+
+4. Generate the filelists. 
+   ```shell
+   python tools/modelnet.py --run generate_points_filelist
+   ```
+
+### Train a classification with Pytorch-based O-CNN
+
 
-2. Run the following command to train the network.
+1. Run the following command to train the network.
    ```
-   python classification.py --config configs/m40_cls.yaml
+   python classification.py --config configs/cls_m40.yaml
    ```
 
-3. To train a deep Resnet, run the following command.
+2. To train a deep ResNet, run the following command.
    ```
-   python classification.py --config configs/m40_cls.yaml \
-                            SOLVER.logdir logs/m40/0301_resnet \
+   python classification.py --config configs/cls_m40.yaml \
+                            SOLVER.logdir logs/m40/resnet \
                             MODEL.name resnet
    ```

docs/installation.md

@@ -87,16 +87,26 @@ commands: `pip install -U gast==0.2.2 numpy==1.16.4`.
 
 The code has been tested with Ubuntu 16.04 and PyTorch 1.6.0.
 
-1. Install PyTorch and relevent packages with the following commands:
-   ```shell
-   conda create --name pytorch-1.6.0 python=3.7
-   conda activate pytorch-1.6.0
-   conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch
-   conda install tqdm yacs -c conda-forge
-   ```
+1. Enter the subfolder `pytorch`, and install PyTorch and relevant packages with
+   the following commands:
+    ```shell
+    conda create --name pytorch-1.7.0 python=3.7
+    conda activate pytorch-1.7.0
+    conda install pytorch==1.7.0 torchvision==0.8.0 cudatoolkit=10.2 -c pytorch
+    pip install -r requirements.txt
+    ```
+
+    The code is also tested with the following pytorch version:
+    - `pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1`
+    - `pytorch==1.9.0 torchvision cudatoolkit=11.1 `
+    <!--
+    - `pytorch==1.8.0 torchvision==0.9.0 cudatoolkit=11.1` conda  failed
+    - `pytorch/pytorch:1.7.0-cuda11.0-cudnn8-devel`        docker failed
+    - `pytorch/pytorch:1.6.0-cuda10.1-cudnn7-devel`        docker succeed
+    -->
+
 2. Build O-CNN under PyTorch.
    ```shell
-   cd pytorch
    python setup.py install --build_octree
    ```
 

docs/scannet.md

@@ -0,0 +1,52 @@
+# ScanNet Segmentation
+
+### Data preparation
+
+1. Download the data from the [ScanNet benchmark](http://kaldir.vc.in.tum.de/scannet_benchmark/).
+   Unzip the data and place it to the folder `<scannet_folder>`
+
+2. Change the working directory to `pytorch/projects`, run the following command
+   to prepare the dataset.
+   ```shell
+   python tools/scannet.py --run process_scannet --path_in <scannet_folder>
+   ```
+
+3. Download the training, validation and testing file lists via the following command.
+   ```
+   python tools/scannet.py --run download_filelists
+   ```
+
+
+### Training and testing
+
+1. Run the following command to train the network with 4 GPUs.
+   ```shell
+   python segmentation.py --config configs/seg_scannet.yaml SOLVER.gpu 0,1,2,3
+   ```
+   The mIoU on the validation set is 74.0, the training log and weights can be
+   downloaded from this [link](https://www.dropbox.com/s/3grwj7vwd802yzz/D9_2cm.zip?dl=0). 
+
+2. To achieve the 76.2 mIoU on the testing set, we follow the practice described
+   in the MinkowsiNet, i.e. training the network both on the training and
+   validation set via the following command.
+   ```shell
+   python segmentation.py --config configs/seg_scannet.yaml SOLVER.gpu 0,1,2,3  \
+                          SOLVER.logdir logs/scannet/D9_2cm_all  \
+                          DATA.train.filelist data/scannet/scannetv2_train_val_new.txt
+   ```
+   The training log and weights can be downloaded from this [link](https://www.dropbox.com/s/szhjus6kmknxyya/D9_2cm_all.zip?dl=0).
+
+3. To generate the per-point predictions, run the following command.
+   ```shell
+   python segmentation.py --config configs/seg_scannet_eval.yaml
+   ```
+
+4. Run the following command to convert the predictions to segmentation labels.
+   ```shell
+   python tools/scannet.py --path_in  data/scannet/test  \
+                           --path_pred logs/scannet/D9_2cm_eval  \
+                           --path_out logs/scannet/D9_2cm_eval_seg   \
+                           --filelist data/scannet/scannetv2_test_new.txt  \
+                           --run generate_output_seg
+   ```
+   Then the generated segmentation results exist in the folder `logs/scannet/D9_2cm_eval_seg`

docs/segmentation.md

@@ -107,11 +107,11 @@ which consists of pooling and unpooling layers.
    `pytorch/projects/dataset`:
     ```shell
     python seg_shapenet.py 
-    ln -s `pwd`/../script/dataset/shapenet_segmentation `pwd`/../../pytorch/projects/dataset
+    ln -s `pwd`/../script/dataset/shapenet_segmentation `pwd`/../../pytorch/projects/data
     ```
 
 3. Change the working directory to `pytorch/projects`. Run the following script
    to train the network. 
     ```shell
-    python shapenet_seg.py --config configs/shapenet_seg.yaml
+    python segmentation.py --config configs/seg_shapenet.yaml
     ```

pytorch/projects/configs/seg_scannet.yaml

@@ -2,9 +2,9 @@ SOLVER:
   gpu: 0,
   run: train
   logdir: logs/scannet/D9_2cm
-  max_epoch: 400
+  max_epoch: 600
   test_every_epoch: 10
-  weight_decay: 0.0005
+  weight_decay: 0.0001
 
   # learning rate
   lr: 0.05

pytorch/projects/configs/seg_scannet_d8.yaml

@@ -1,10 +1,10 @@
 SOLVER:
   gpu: 0,
   run: train
-  logdir: logs/scannet/D8_4cm
-  max_epoch: 400
+  logdir: logs/scannet/D9_2cm
+  max_epoch: 600
   test_every_epoch: 10
-  weight_decay: 0.0005
+  weight_decay: 0.0001
 
   # learning rate
   lr: 0.05

pytorch/projects/configs/seg_scannet_eval.yaml

@@ -2,8 +2,8 @@ SOLVER:
   gpu: 0,
   logdir: logs/scannet/D9_2cm_eval
   run: evaluate
-  eval_epoch: 1
-  ckpt: logs/scannet/D9_2cm_poly_b6_ep600_w1e-4_lr0.1_0807/checkpoints/00600.model.pth 
+  eval_epoch: 72
+  ckpt: logs/scannet/D9_2cm_all/checkpoints/00600.model.pth 
 
 DATA:
   test:
@@ -15,7 +15,7 @@ DATA:
     offset: 0.0
 
     # data augmentations
-    distort: False  # no data augmentation
+    distort: True  # no data augmentation
     angle: (0, 0, 180)
     scale: 0.1
     jitter: 0.1

pytorch/projects/tools/scannet.py

@@ -7,7 +7,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--path_in', type=str, default='data/scannet/ScanNet_data')
-parser.add_argument('--path_out', type=str, default='data/scannet/scans')
+parser.add_argument('--path_out', type=str, default='data/scannet')
 parser.add_argument('--path_pred', type=str, default='logs/scannet/D9_2cm_eval')
 parser.add_argument('--filelist', type=str, default='scannetv2_test_new.txt')
 parser.add_argument('--run', type=str, default='generate_output_seg')
@@ -29,6 +29,24 @@
 ilabel_dict = dict(zip(np.arange(0, 21), class_ids))
 
 
+def download_filelists():
+  path_out = args.path_out
+  zip_file = os.path.join(path_out, 'filelist.zip')
+  if not os.path.exists(path_out):
+    os.makedirs(path_out)
+
+  # download
+  url = 'https://www.dropbox.com/s/aeizpy34zhozrcw/scannet_filelist.zip?dl=0'
+  cmd = 'wget %s -O %s' % (url, zip_file)
+  print(cmd)
+  os.system(cmd)
+
+  # unzip
+  cmd = 'unzip %s -d %s' % (zip_file, path_out)
+  print(cmd)
+  os.system(cmd)
+
+
 def read_ply(filename, compute_normal=True):
   with open(filename, 'rb') as fid:
     plydata = PlyData.read(fid)