A PyTorch implementation of UNet semantic segmentation network, applied to the SpaceNet v1 Building Detection Challenge. View the project report here.
- Liam Coulter
- Teague Hall
- Luis Guzman
- Isaac Kasahara
Download the training dataset images (~3.71 GB), and image masks (~35.7 MB)
Optionally download the geojson vector label data (~201 MB) to test preprocessing scripts
Place the training images in SN1_buildings_train_AOI_1_Rio_3band/3band/*.tif and the masks in building_mask/*.tif
To test a pre-trained model, download the network weights and place it in the trained_models folder.
Official data for SpaceNet v1 Building Challenge is also accessible through AWS.
To train the UNet, run train.py using the command python train.py --batch_size 16
A batch size of 16 takes approximately 12GB of video memory, so adjust the batch size according to what you have available.
To view the results, run test.py
The pre-trained model can be downloaded here and should be placed in a directory called trained_models.
Network evaluation can be calculated with evaluate.py
Note: We provide already-processed image masks, so the following step is not necessary to test our implementation.
To run image_json_to_mask.py, you need to first install the SpaceNet Utilities from here. Follow the instructions and download necessary files. It worked for Liam on MacOS to make a Conda virtual environment, and load everything in there. The only package I had trouble with was rasterio, since I assume it is older. Note that the file geoTools.py at the base repository level is a modified version of geoTools.py that comes with the SpaceNet utilities install, and includes 1 extra function (latlonToPixel).
Once you have the SpaceNet utilities, edit directory variables in image_json_to_mask.py to reflect your setup, and you should be able to run.