S23DR

Training code for the paper Edge Prediction for Roof Wireframe Reconstruction with Transformers.

This repository contains the code for training a transformer-based network for reconstructing house roof wireframes from sparse SfM point clouds and segmentation images.

Note: This is not a full implementation of the solution described in the paper. For example: the Gestalt autoencoder, the test-time augmentations and the model ensembling are missing. However, this simplified method still achieves competitive results.

virtualenv venv
source venv/bin/activate
pip install -r requirements.txt

Dump training data:

python preprocess.py --split train --output_dir OUTPUT_DIR
python preprocess.py --split validation --output_dir OUTPUT_DIR

This should result in a set of 19603 and 170 point cloud .ply files for the training and validation sets, respectively.

Train the edge predictor:

python train.py --name NAME --ply_dir PLY_DIR --fp16 --compile

The training takes around 12 hours and uses 10 GB of memory. It should reach a validation HSS of approximately 0.47.

Tip: Use --wandb_project <PROJECT> to log results to Weights & Biases.

Evaluate a trained model on the validation set:

python eval.py --checkpoint CHECKPOINT

A pre-trained checkpoint can be downloaded here.