Skip to content

PolyGNN: Polyhedron-based graph neural network for 3D building reconstruction from point clouds [ISPRS 2024]

License

Notifications You must be signed in to change notification settings

chenzhaiyu/polygnn

Repository files navigation

PolyGNN


Paper: HTML Paper: PDF License: MIT

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},
}

About

PolyGNN: Polyhedron-based graph neural network for 3D building reconstruction from point clouds [ISPRS 2024]

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages