PolyGNN

---

PolyGNN is an implementation of the paper PolyGNN: Polyhedron-based Graph Neural Network for 3D Building Reconstruction from Point Clouds. PolyGNN learns a piecewise planar occupancy function, supported by polyhedral decomposition, for efficient and scalable 3D building reconstruction.

🛠️ Setup

Repository

Clone the repository:

git clone https://github.com/chenzhaiyu/polygnn && cd polygnn

All-in-one installation

Create a conda environment with all dependencies:

conda env create -f environment.yml && conda activate polygnn

Manual installation

Still easy! Create a conda environment and install mamba for faster parsing:

conda create --name polygnn python=3.10 && conda activate polygnn
conda install mamba -c conda-forge

Install the required dependencies:

mamba install pytorch torchvision sage=10.0 pytorch-cuda=11.7 pyg=2.3 pytorch-scatter pytorch-sparse pytorch-cluster torchmetrics rtree -c pyg -c pytorch -c nvidia -c conda-forge
pip install abspy hydra-core hydra-colorlog omegaconf trimesh tqdm wandb plyfile

🚀 Usage

Quick start

Download the mini dataset and pretrained weights:

python download.py dataset=mini

In case you encounter issues (e.g., Google Drive limits), manually download the data and weights here, then extract them into ./checkpoints/mini and ./data/mini, respectively. The mini dataset contains 200 random instances (~0.07% of the full dataset).

Train PolyGNN on the mini dataset:

python train.py dataset=mini

The data will be automatically preprocessed the first time you initiate training.

Evaluate PolyGNN with option to save predictions:

python test.py dataset=mini evaluate.save=true

Generate meshes from predictions:

python reconstruct.py dataset=mini reconstruct.type=mesh

Remap meshes to their original CRS:

python remap.py dataset=mini

Generate reconstruction statistics:

python stats.py dataset=mini

Available configurations

# check available configurations for training
python train.py --cfg job

# check available configurations for evaluation
python test.py --cfg job

Alternatively, review the configuration file: conf/config.yaml.

Custom data

PolyGNN requires polyhedron-based graphs as input. To prepare this from your own point clouds:

1. Extract planar primitives using tools such as Easy3D or GoCoPP, preferably in VertexGroup format.
2. Build CellComplex from the primitives using abspy. Example code:

   from abspy import VertexGroup, CellComplex
   vertex_group = VertexGroup(vertex_group_path, quiet=True)
   cell_complex = CellComplex(vertex_group.planes, vertex_group.aabbs,
                              vertex_group.points_grouped, build_graph=True, quiet=True)
   cell_complex.prioritise_planes(prioritise_verticals=True)
   cell_complex.construct()
   cell_complex.save(complex_path)

   Alternatively, you can modify CityDataset or TestOnlyDataset to accept inputs directly from VertexGroup or reference mesh.
3. Structure your dataset similarly to the provided mini dataset:

   YOUR_DATASET_NAME
   └── raw
       ├── 03_meshes
       │   ├── DEBY_LOD2_104572462.obj
       │   ├── DEBY_LOD2_104575306.obj
       │   └── DEBY_LOD2_104575493.obj
       ├── 04_pts
       │   ├── DEBY_LOD2_104572462.npy
       │   ├── DEBY_LOD2_104575306.npy
       │   └── DEBY_LOD2_104575493.npy
       ├── 05_complexes
       │   ├── DEBY_LOD2_104572462.cc
       │   ├── DEBY_LOD2_104575306.cc
       │   └── DEBY_LOD2_104575493.cc
       ├── testset.txt
       └── trainset.txt

4. To train or evaluate PolyGNN using your dataset, run the following commands:

   # start training
   python train.py dataset=YOUR_DATASET_NAME
   
   # start evaluation
   python test.py dataset=YOUR_DATASET_NAME

   For evaluation only, you can instantiate your dataset as a TestOnlyDataset, as in this line.

👷 TODOs

• Demo with mini data and pretrained weights
• Short tutorial for getting started
• Host the entire dataset (>200GB)

🎓 Citation

If you use PolyGNN in a scientific work, please consider citing the paper:

@article{chen2024polygnn,
title = {PolyGNN: Polyhedron-based graph neural network for 3D building reconstruction from point clouds},
journal = {ISPRS Journal of Photogrammetry and Remote Sensing},
volume = {218},
pages = {693-706},
year = {2024},
issn = {0924-2716},
doi = {https://doi.org/10.1016/j.isprsjprs.2024.09.031},
url = {https://www.sciencedirect.com/science/article/pii/S0924271624003691},
author = {Zhaiyu Chen and Yilei Shi and Liangliang Nan and Zhitong Xiong and Xiao Xiang Zhu},
}