RoofSeg

This is the PyToch implementation of the following manuscript:

RoofSeg：An Edge-aware transformer-based network for end-to-end roof plane segmentation

Siyuan You, Guozheng Xu, Pengwei Zhou, Qiwen Jin, Jian Yao, Li Li

Dataset

Three roof datasets： RoofNTNU is a small real-world dataset provide by RoofNTNU, which can be downloaded from baiduyun. Roofpc3D and Building3D can be derived from DeepRoofPlane.

Two generic plane segmentation datasets: KSR42 is a plane shape detection dataset bulit by KSR42. We slected 10 models and annonated them with accurate plane lables, which can be direived from baiduyun. ABC can be download from [ABC CAD MODELS] (https://deep-geometry.github.io/abc-dataset/) or baiduyun.

If you want to train and test on your own dataset, please make sure that the input point cloud contain 5 columns [x, y, z, ins_label, edge_label], where the first three columns are the 3D coordinates of each point, ins_label represents the plane instance label, and edge_label represents the edge label (1 for edge points, 0 for non-edge points). If you need to generate edge labels for the dataset, you can run ''' python edge_label.py

## Usage
The environment requires torch1.12.1+cu113 with python=3.9.0 The requirements list is shown as follow:
- numpy
- scipy
- matplotlib
- trimesh
- tensorboardx
- scikit-learn
- pyyaml
- libboost
- plyfile


If the above requirements are incomplete, you can install the missed dependencies according to the compilation errors.

## Train and test
Then, you can run main.py for training and testing the model. Notably, the commands with parameter settings required by the model are stored in the `./scripts/scripts.sh` file. The model checkpoints will be stored in the `./outputs` folder.
An example for training command:
```shell script
CUDA_VISIBLE_DEVICES=0 python main.py --dataset_name building3d --ngpus 1 --nqueries 64 --enc_nlayers 1 --dec_nlayers 1 --encoder_layer_num 4 --decoder_layer_num 4 --decoder_num 2 --max_epoch 180 --matcher_cls_cost 0.5 --loss_no_object_weight 0.2 --loss_sem_cls_weight 1.0 --loss_plane_weight 5.0 --batchsize_per_gpu 16 --save_separate_checkpoint_every_epoch 30 --checkpoint_dir outputs/building3d_q_32_enc_4_dec_4_mask_dec_reverse_shared_decoder --mask_pred --edge_pred --weight_loss --boundary_aware_mask --val_only 

An example for testing command (--test_ckpt $test_model_dir$):

CUDA_VISIBLE_DEVICES=0 python main.py --dataset_name roofpc3d --ngpus 1 --nqueries 16 --enc_nlayers 2 --dec_nlayers 8  --batchsize_per_gpu 1 --checkpoint_dir outputs/roofpc3d_q_16_dec_8 --test_ckpt checkpoints/checkpoint.pth --test_only

If you want clustering-based post-processing, add "--cluster" to testing command.

Results

The model output, such as pred_sem_label, offset, feature, etc., will be saved in ./results. Several examples are provided in ./results.

If you want visualize the output features, we recommend Google's open-source software Embedding Projector. Run txt2tsv.py for format conversion.

python txt2tsv.py

For the quantative evaluation of the experimental results, you can check ./quanti_eval_res.txt.

Citation

If you find our work useful for your research, please consider citing our paper.

RoofSeg：An Edge-aware transformer-based network for end-to-end roof plane segmentation

Siyuan You, Guozheng Xu, Pengwei Zhou, Qiwen Jin, Jian Yao, Li Li

Contact:

Li Li (li.li@whu.edu.cn)