Diffusion-based RGB-D Semantic Segmentation with Deformable Attention Transformer

TL;DR

This work confronts RGB-D semantic segmentation under noisy depth measurements by introducing a diffusion-based framework conditioned on fused RGB-D features. A Deformable Attention Transformer encoder robustly processes depth data with many invalid pixels, while a diffusion-based mask decoder learns the data distribution to generate segmentation masks, yielding rapid convergence and state-of-the-art results on NYUv2 and SUN-RGBD. The approach is validated through challenging dataset tests and a real-world drone-based volumetric mapping demonstration, underscoring the practical impact for autonomous systems. Overall, the combination of deformable spatial attention and conditional diffusion provides robust, efficient RGB-D segmentation with strong generalization to diverse sensing conditions.

Abstract

Vision-based perception and reasoning is essential for scene understanding in any autonomous system. RGB and depth images are commonly used to capture both the semantic and geometric features of the environment. Developing methods to reliably interpret this data is critical for real-world applications, where noisy measurements are often unavoidable. In this work, we introduce a diffusion-based framework to address the RGB-D semantic segmentation problem. Additionally, we demonstrate that utilizing a Deformable Attention Transformer as the encoder to extract features from depth images effectively captures the characteristics of invalid regions in depth measurements. Our generative framework shows a greater capacity to model the underlying distribution of RGB-D images, achieving robust performance in challenging scenarios with significantly less training time compared to discriminative methods. Experimental results indicate that our approach achieves State-of-the-Art performance on both the NYUv2 and SUN-RGBD datasets in general and especially in the most challenging of their image data. Our project page will be available at https://diffusionmms.github.io/

Figures (7)

• Figure 1: We formulate the RGB-D semantic segmentation task as a denoising diffusion process conditioned by RGB and depth images.
• Figure 2: The architecture of our RGBD semantic mask generation framework. A Deformable Attention Transformer is used as the hierarchical encoder to extract features from RGB and depth images. Multi-scale features from both branches are then processed using fusion modules followed by a Feature Pyramid Network to create a conditioning signal $x$ that matches the shape of the noisy segmentation ground truth. A deformable attention mask encoder is trained to gradually denoise the concatenated signal to generate the segmentation mask.
• Figure 3: Loss comparison between training the original CMX zhang2023cmx model and when using DAT++-S as the encoder for raw depth images.
• Figure 4: GradCAM (Gradient-weighted Class Activation Map) heatmaps visualizing class activations from the final layer of the backbone used in our method and in CMX zhang2023cmx. For each method, we present (from left to right) the predicted semantic segmentation mask, the binary mask for the selected class, and the corresponding GradCAM heatmap. The top row displays results for the class 'window'. The bottom row shows results for the class 'table'.
• Figure 5: Performance and computational cost of our method with different numbers of sampling steps used for inference on NYUv2 dataset