Joint Learning of Semantic Segmentation and Height Estimation for Remote Sensing Image Leveraging Contrastive Learning
Zhi Gao, Wenbo Sun, Yao Lu, Yichen Zhang, Weiwei Song, Yongjun Zhang, Ruifang Zhai
Abstract
Semantic segmentation and height estimation are two critical tasks in remote sensing scene understanding that are highly correlated with each other. To address both tasks simultaneously, it is natural to consider designing a unified deep learning model that aims to improve performance by jointly learning complementary information among the associated tasks. In this paper, we learn the two tasks jointly under a deep multi-task learning framework and propose two novel objective functions, called cross-task contrastive loss and cross-pixel contrastive loss, respectively, to enhance multi-task learning performance through contrastive learning. Specifically, the cross-task contrastive loss is designed to maximize the mutual information of different task features and enforce the model to learn the consistency between semantic segmentation and height estimation. In addition, our method goes beyond previous approaches that only apply contrastive learning at the instance level. Instead, we design a pixel-wise contrastive loss function that pulls together pixel embeddings belonging to the same semantic class, while pushing apart pixel embeddings from different semantic classes. Furthermore, we find that this semantic-guided contrastive loss simultaneously improves the performance of the height estimation task. Our proposed approach is simple and effective and does not introduce any additional overhead to the model during the testing phase. We extensively evaluate our method on the Vaihingen and Potsdam datasets, and the experimental results demonstrate that our approach significantly outperforms the state-of-the-art methods in both height estimation and semantic segmentation.