zorzi-s / PolyWorldPretrainedNetwork

PolyWorld: Polygonal Building Extraction with Graph Neural Networks in Satellite Images
Other
180 stars 28 forks source link

Research @ TUGraz & BlackShark.ai (CVPR 2022)

PolyWorld Inference and Evaluation Code

PolyWorld is a research project conducted by the Institute of Computer Graphics and Vision of TUGraz, in collaboration with BlackShark.ai. PolyWorld is a neural network that extracts polygonal objects from an image in an end-to-end fashion. The model detects vertex candidates and predicts the connection strenght between each pair of vertices using a Graph Neural Network. This repo includes inference code and pretrained weights for evaluating PolyWorld on the CrowdAI Mapping Challenge dataset.

Dependencies

Getting started

After cloning the repo, download the _polyworldbackbone pre-trained weights from here, and place the file in the _trainedweights folder.

The CrowdAI Mapping Challenge dataset can be downloaded here.

Run the evaluation on the CrowdAI Mapping Challenge dataset

To run the evaluation, specify batch size, image folder, and annotation file of the CrowdAI dataset in the main function of the prediction.py script. Then simply run:

python prediction.py

The code is tested on an Nvidia RTX 3090 using _batchsize = 6.

During inference, the script converts the predicted polygons to coco format and saves a json file (predictions.json).

If you wish to visualize the results in QGIS, we suggest to convert the predictions from coco json format to shapefile using the _coco_toshp.py script. To run the conversion, specify the json file and the output folder in the main function, and then type:

python coco_to_shp.py

In order to compute AP and AR metrics with the COCO API, or the MTA metric, please use the script provided by the Frame Field Learning repo.

If you want to compute IoU and C-IoU metrics, use the _coco_IoUcIoU.py script. To run the evaluation, specify json file and ground truth annotations in the main function and then run:

python coco_IoU_cIoU.py

Download results

A download link for the PolyWorld predictions corresponding to the val-set of the CrowdAI dataset is also provided:

BibTeX citation

If you use any ideas from the paper or code from this repo, please consider citing:

@inproceedings{zorzi2022polyworld,
  title={PolyWorld: Polygonal Building Extraction with Graph Neural Networks in Satellite Images},
  author={Zorzi, Stefano and Bazrafkan, Shabab and Habenschuss, Stefan and Fraundorfer, Friedrich},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={1848--1857},
  year={2022}
}