OTM-UNet: Optimized Semantic Segmentation of Remote Sensing Imagery with Learned Optimal Transport Maps

Abstract

Deep learning techniques have recently shown remarkable effectiveness in the semantic segmentation of natural and remote sensing (RS) images. However, despite advances in conventional networks, accurate detection and segmentation of small objects in complex scenes remains a major challenge. The detection and segmentation of small features, such as vehicles and pedestrians, is complex due to the occlusion and density of contextual information. In this paper, we propose an enhanced UNet architecture, called Optimal Transport Maps (OTM-UNet), which uses optimal transport layers to compute learned transport maps that align feature maps from both the encoder and decoder. This alignment is critical for
preserving spatial orientation and improving semantic consistency during segmentation. Optimized transport layers are strategically placed deep in the decoder and perform exhaustive transformations on feature maps sliced from encoders before concatenation. The resulting transport maps bridge the gap between encoder and decoder feature distributions, facilitating effective information transfer and preserving spatial detail throughout the architecture.
The performance of OTM-UNet was evaluated using two publicly available Remote Sensing Imagery datasets, and a comprehensive quantitative and qualitative comparison was made with other models. Results from the evaluation of the Vaihingen dataset showed that the proposed model achieved an impressive average F1 score of 90.90% and an accuracy of 93.17%. In addition, the visual qualitative results showed a significant reduction in object class confusion, improved ability to segment different scales of object, and improved object integrity, highlighting the model`s effectiveness in addressing small objects in remote sensing segmentation challenges.