Skip to content

Repo for training a U-Net Model to segment rooftops from aerial images using a small training set.

Notifications You must be signed in to change notification settings

dlculver/rooftop-segmentation

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Here is a short summary of what is in this repo. The data in the directory 'data'. I split the training set into a training and validation set. I created a torch DataSet class, RoofDataSet, which also performs data augmentation on the training data to make up for the fact that there are only 20 training images. The code can be found in 'data.py'.

The model that I use here is a U-Net model with back-bone from MobileNetV2. The specifics of the architecture can be found in 'model.py'. In particular, I download a pre-trained MobileNet model and extracted the output from blocks 2, 4, 7, 14, and 17 in MobileNet, then run the output from block 17 through up-sampling blocks and also employ skip connections with the corresponding downsampling blocks. This implementation is based on the implementation in the tensorflow tutorial https://www.tensorflow.org/tutorials/images/segmentation. 

Finally, there are some visualization functions written up in 'utils.py'. Training is done in a Jupyter notebook, 'train.ipynb'. I train an instance of the model on the training set of 20 images for 50 epochs. In the notebook, you can see a plot between the training loss and the validation loss. I also train an instance for 100 epochs to see if training improves. 

About

Repo for training a U-Net Model to segment rooftops from aerial images using a small training set.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published