Repository files navigation

This repo contains a U-Net implementation for satellite segmentation

Due to a severe lack of training data, several pre-processing steps are taken to try andalleviate this. First, the annotations are converted from their json files to image masks. Then, the satellite images are further tileddown, to tiles of size(512x512), so that the images can be fed into a fully convolutional network (FCN) for semantic segmentation.The augmentations are:

• blur: combination of median and bilateral blur
• bright increase: increase brightness artificially
• distort: elastic deformation of image
• gaussian blur: gaussian blurring
• HSV: convert channels to HSV
• medial blur: median blur
• mirror: mirror image
• rotation invariance: apply rotation invariance (seedata_utils.augment.rotation_invariance() for details)
• crop + resize: crop image randomly and resize to size expected by network (crop is applied to both image and mask)

blur	bright increase	distort	gaussian blur

HSV shift	median blur	mirror	gaussian blur + rotation

distort + rot	rotation invariance	crop + resize

U-Net is used as a base model for segmentation. The original intention was to use U-Net to show base results,and then trainPSPNet (Pyramid Scene Parsing Network) using a pretrained satellite segmentation model and show comparisons, but time did not allow for this.

U-Net was trained for50 43 epochs, with a batch size of 4.

The base parameters for training can be seen, and adjusted, inrun_training.py:

defsetup_run_arguments():args=EasyDict()args.epochs=50args.batch=4args.val_percent=0.2args.n_classes=4args.n_channels=3args.num_workers=8args.learning_rate=0.001args.weight_decay=1e-8args.momentum=0.9args.save_cp=Trueargs.loss="CrossEntropy"

The trained model is provided incheckpoints/unet-augment-final.pth

The expected output is a json annotated file containing the vector points corresponding to the classes.A function for generating such file is found inpredict.prediction_to_json(...). A python notebookis provided showing how to generate the json file, as well as how to generate a color mask file for the json file.

The example:example.ipynb shows howto load the trained model and use it to create the annotation file. It also shows how to create a colored image masks directly with theannotation. (The example notebook was created with u-net after 35 epochs)

There are several improvements that can be made. For starters, some augmentation methods used could be replacedor left out entirely, such as HSV. Simple color shift could have been used instead. Another major issue that shouldhave been addressed at the beginning was class imbalance. It would have been better to apply augmentations with respect tothe class frequency, trying to shift the infrequency balance.

Another obvious issue is that U-Net was a network catered to medical image segmentation, but it is often used as a baseline modebecause it is small and easy to implement. A more suitable network would have been PSPNet, as mentioned above. Similarly,there exists several pretrained model that could have been used with transfer learning. This, coupled with meaningful augmentations,would have yielded better model/results.