TSMD-Net: A two-stage mixed dehazing network with feature fusion and multi-window self attention on Jetson board

Abstract

Single image dehazing is a pivotal endeavor focused on enhancing the quality of images by removing haze-induced noise. Achieving this goal requires a delicate equilibrium between retaining spatial intricacies and contextual understanding, particularly crucial in real-world scenarios where complexity is amplified. In response to this challenge, we present the Two-Stage Mixed Dehazing Network (TSMD-Net), designed to reconcile these conflicting objectives by leveraging the foundational U-Net architecture. Our approach entails a dual-phase strategy: initially assimilating contextual comprehension through encoder-decoder frameworks, followed by amalgamating it with high-resolution branches to uphold spatial integrity. To streamline the analysis and comparison of this architecture, we opt to substitute or eliminate nonlinear activation functions with multiplication operations, effectively mitigating system complexity. Furthermore, to foster seamless information exchange across hierarchical layers of the encoder-decoder at varying scales, we introduce the Feature Fusion (FF) mechanism, facilitating the integration of insights from higher layers progressively down to the base layer. This amalgamated data significantly enhances the fidelity of the original image transformation by harmonizing the inclusion of contextual feature maps with the preservation of spatial details, resulting in noise-free, high-quality images. Additionally, we propose the “Multi-Window Self Attention” (MWSA) mechanism, characterized by linear time complexity, as the central block of the encoder-decoder. This module mitigates the confined receptive fields inherent in convolutional neural networks (CNNs), enabling the aggregation of more comprehensive feature-map data. Subsequently, our TSMD-Net undergoes rigorous training and testing on both synthetic and real-world datasets, showcasing superior performance compared to current state-of-the-art (SOTA) networks while demanding fewer processing resources. To validate the robustness and practical utility of our methodology, we implement TSMD-Net on a low-end Jetson Nano board, confirming its effectiveness and efficiency in real-world settings.