Add Mini-Unet for segmentation

[iiroseppa]

I propose to add two versions of Mini-Unet and associated data sampler.

• I have working prototypes, which need further developing. However, I want to submit them now because the project deadline for new algorithms is arriving. The core of the algorithm is finished, however there is a lot things that need implementation before this is ready to be published.

Mini U-net with learned uncertainty estimation

• Semantic segmentation for (small) images, with included learned uncertainty estimation
  • There is separate script for semantic segmentation without uncertainty estimation
• The idea is based on
  • DeVries, T., and G. W. Taylor. 2018b. “Leveraging Uncertainty Estimates for Predicting Segmentation Quality. https://arxiv.org/abs/1807.00502
• This implementation takes strong inspiration from
  • Wu, S., Heitzler, M. & Hurni, L. (2022) ‘Leveraging uncertainty estimation and spatial pyramid pooling for extracting hydrological features from scanned historical topographic maps’, GIScience & Remote Sensing, 59(1), pp. 200–214. Available at: https://doi.org/10.1080/15481603.2021.2023840.
• Loss based on (Sørensen -)Dice coefficient is utilized due to favorable properties related to unbalanced pixel counts.
  • Milletari, F., Navab, N. & Ahmadi, S.-A. (2016) ‘V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation’, in 2016 Fourth International Conference on 3D Vision (3DV). 2016 Fourth International Conference on 3D Vision (3DV), pp. 565–571. Available at: https://doi.org/10.1109/3DV.2016.79.

To summarise the main idea

• In addition to segmentation, the model is asked to create certainty estimation
• When overall certainty is high (close to 1) the loss function operates as if no certainty estimate was used
• When certainty values dip below 1, the model is punished less for making mistakes at those pixels.
  • When certainty is close to zero at certain pixel, the model segmentation output does not matter at all for the loss at that pixel
• To prevent model predicting high uncertainty everywhere, it is punished for being uncertain at large areas
  • This is done by averaging logarithm of the uncertainties at each pixel * -1 together
  • This regularization loss is multiplied with a hyperparameter called uncertainty coefficient
  • The value of the uncertainty coefficient has a large effect on the results, and finding suitable value should be automatised (not yet implemented)
    • Too small values result in poor segmentation results
    • Too large values result in abolute certainty everywhere
    • Suitable value creates resulting values between full range of zero and one for both segmentation prediction and certainty estimate
    • Thus closer (max cert - min cert) + (max pred - min pred) is to to 2, the better the uncertainty coefficient. This could be used as a basis for optimisation

Notes

• Current implementation is experimental and needs further optimization, testing and functionality
  • IMPORTANT The possibility to input a large geotiff should be implemented. In that case the trained model should be used to predict the whole area in a sliding window type of setup. Then it should be parsed back together to create output prediction and uncertainty estimate geotiffs for the whole area. This tool should not be included in the toolkit before that can be done.
  • hyperparameter tuning should be automatized
  • Dice coefficient and Intersection over Union should be implemented as metric to measure the goodness of model
    • custom metrics are needed due to the unusual shape of the output
  • Data suffling should maybe moved to data sampling side
    • Leave one out cross validation should also be implemented and integrated with data sampling, so that user has reliable tools to estimate possible (likely) overfitting
    • current implementation with data sampler does not create reliable test data, because one occurence features multiple times in data with different croppings.
    • The other option is to take out all the times one of the features occurs before suffling and keep data sampling as is
• The current implementation only works with binary classes (0 and 1)
• Image size should be at least 32 x 32 pixels or larger power of 2 (or at least divisible three times with 2, after which size should be at least 3 x 3)
• Using either data from automatic sampling function or user's own data should be given as an option
• The model needs to be trained long enough so that it produces both prediction and uncertainty, even if the uncertainty punishment coefficient is correct. The balancing is quite intricate and highlights the need for automatic hyperparametrisation

Mini U-net

• Semantic segmentation for (small) images, with included learned uncertainty estimation
  • There is separate script for semantic segmentation with learned uncertainty estimation
• Loss based on (Sørensen -)Dice coefficient is utilized due to favorable properties related to unbalanced pixel counts.
  • Milletari, F., Navab, N. & Ahmadi, S.-A. (2016) ‘V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation’, in 2016 Fourth International Conference on 3D Vision (3DV). 2016 Fourth International Conference on 3D Vision (3DV), pp. 565–571. Available at: https://doi.org/10.1109/3DV.2016.79.
• This implementation also takes inspiration from
  • Wu, S., Heitzler, M. & Hurni, L. (2022) ‘Leveraging uncertainty estimation and spatial pyramid pooling for extracting hydrological features from scanned historical topographic maps’, GIScience & Remote Sensing, 59(1), pp. 200–214. Available at: https://doi.org/10.1080/15481603.2021.2023840.

The main idea

• Learns to segment areas of interest with fewer parameters than normal U-net, which reduces the need for training data and also the risk for overfitting

Notes

• Current implementation is experimental and needs further optimization, testing and functionality

  • IMPORTANT The possibility to input a large geotiff in addition to training and testing data should be implemented. In that case the trained model wouldbe used to predict the whole area in a sliding window type of setup. Then it should be parsed back together to create output prediction and uncertainty estimate geotiffs for the whole area. This tool should not be included in the toolkit before that can be done.
  • For example, hyperparameter tuning should be automatized
  • Dice coefficient and Intersection over Union should be implemented as metric to measure the goodness of model
    • custom metrics are needed due to the unusual shape of the output
  • Data suffling should maybe moved to data sampling side
    • Leave one out cross validation should also be implemented and integrated with data sampling, so that user has reliable tools to estimate possible (likely) overfitting
    • current implementation with data sampler does not create reliable test data, because one occurence features multiple times in data with different croppings.
    • The other option is to take out all the times one of the features occurs before suffling and keep data sampling as is

• The current implementation only works with binary classes (0 and 1)

• Image size should be at least 32 x 32 pixels or larger power of 2 (or at least divisible three times with 2, after which size should be at least 3 x 3)

• Using either data from automatic sampling function or user's own data should be given as an option

Data sampler

Creates a set of images that can be used for training a convolutional neural network.

Input

• Geotiff that will be sampled
• Bounds of the geotiff, geopackage or other geospatial vector format
  • TODO: Just extract the extent from geotiff nodata mask
• locations of the wanted class, geopackage or other geospatial vector format

Output

• Two folders with geotiff images inside, one containing images and the other containing labels

Internal workings

• The known occurrences are oversampled and negative samples are selected randomly.
  • This ensures that both sample types are represented equally in the created dataset.
• The sampler works well with sparse datasets, where it is relatively rare that the negative samples would contain an unkown occurence
  • Thus there shouldn't be too much misclassifications in the data that would cause problem
• User has the option to set buffer size for the known occurences, because one pixel locations are harder to locate for model than .

TODO

• Current implementation creates images in order, known occurences first and non-occurences last. This makes suffling necessary at the model script, which complicates things. Suffling should be moved to this script
• This sampling script should be integrated to the cnn scripts so that leave one out cross validation would be easy to do
  • The leave one out method should be implented in a way that one feature is completely removed from the occurence data before any sampling is done, and then test data is creaed by sampling the left out feature alone
  • without this it is hard to assess level of overfitting
  • users should be given a clear warning about htis, if the script is published without these changes. Overfitting is nearly unavoidable with low amounts of data so user should be given tools to asess the level of overfitting