dlculver / rooftop-segmentation Goto Github PK

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.