基于收敛性感知的任务调度与位置约束语义Mamba模型，实现一体化恶劣天气图像恢复

《Neurocomputing》：Convergence-aware task scheduling with position-constrained semantic mamba for all-in-one adverse weather image restoration

　　

Xianhao Wu|Guili Xu|Xiang Chen|Qianfeng Yang|Qiyuan Guan|Xianglong Zeng|Junwei He

中国江苏省南京市211106，南京航空航天大学自动化工程学院

摘要

雨、雾、雪等恶劣天气条件会严重降低图像的可见度并破坏结构细节，给视觉感知及后续视觉任务带来挑战。现有的综合修复方法通常侧重于统一特征学习，而对任务间的优化失衡以及考虑退化特征的交互作用关注不足。为解决这些问题，我们提出了AWReNet，这是一种用于综合处理恶劣天气图像修复的统一框架。具体而言，我们引入了一种基于收敛性的任务调度策略（CATS），根据各任务的特定收敛状态动态调整任务权重，从而改善多天气联合训练中的优化平衡。此外，我们还设计了一种带位置约束的动态提示Mamba模块（PDPMamba），该模块通过稀疏动态提示来建模细粒度的退化感知修复状态，并通过位置约束进行语义扫描，以增强特征交互的同时保持空间连续性。在WeatherBench和UAV-Weather6k上的大量实验表明，AWReNet在雨、雾、雪天气条件下都能实现更出色的修复效果。进一步的绝缘子检测实验显示，AWReNet生成的修复图像为基于无人机的电力巡检任务提供了更优质的视觉信息。

引言

雨、雾、雪等恶劣天气条件往往会降低图像的可见度并损害结构细节，从而严重影响视觉感知及后续视觉任务。在基于无人机的成像场景中，这一问题更为复杂，因为天气导致的图像退化往往伴随着复杂的背景、较小的目标以及较大的空间差异。因此，开发一种能够有效处理多种天气退化的统一修复框架，对于在现实世界的恶劣环境中实现稳健的视觉感知至关重要。早期的图像修复方法通常针对某种特定的退化类型，如去雨[1]、[2]、[3]、[4]，去雾[5]、[6]、[7]或除雪[8]、[9]。尽管这些针对特定任务的方法在相应场景下取得了不错的性能，但它们往往需要为不同的退化类型单独构建模型，且对混合或未知的天气条件适应性较差。为克服这一局限，人们提出了综合修复方法，试图在统一框架内修复由不同天气类型引起的图像退化，例如AirNet[10]、TransWeather[11]和PromptIR[12]。这些方法提升了修复模型的灵活性，但由于不同天气退化具有不同的视觉模式和收敛行为，多天气联合训练仍存在任务干扰和优化失衡的问题。

由于不同天气退化具有不同的视觉模式和学习难度，其在联合训练过程中的收敛行为往往不一致。如图1(a)所示，去雾、去雨和除雪任务的损失分布各不相同，这表明静态训练策略无法有效适应各任务的特定学习状态。固定的任务权重可能导致优化失衡，使得相对简单的任务主导训练过程，而那些更复杂的退化类型则得不到充分处理。图1(b)展示了在训练过程中为不同天气任务分配动态权重的思路，这样雾、雨和雪就能得到不同的学习重点，而不会被统一对待。图1(c)中的性能提升进一步证明了，与固定任务权重相比，动态任务分配更适用于多天气修复。

除了训练优化之外，有效的特征建模对于多天气修复也至关重要。基于Transformer的方法，如Restormer[13]和Histoformer[14]，通过捕捉长距离依赖关系提升了修复性能，但其注意力机制往往会导致较高的计算复杂度。近年来，状态空间模型为视觉修复提供了一种高效的替代方案，基于Mamba的方法，如MambaIR[15]和VMambaIR[16]，也展现出了良好的性能。然而，传统的空间Mamba扫描通常遵循固定的光栅、行序或列序，其token之间的交互主要依赖于空间邻近性。这对于多天气修复来说并非最佳方案，因为具有相似退化状态的区域可能在空间上相距较远，而空间上相邻的区域却可能对应着不同的退化强度。因此，固定的空间扫描可能会限制退化感知式的交互，无法捕捉到细微的天气相关特征。

为解决上述问题，我们提出了AWReNet，这是一种用于综合处理恶劣天气图像修复的统一框架。AWReNet从两个互补的角度提升了多天气修复的性能。首先，它引入了基于收敛性的任务调度策略（CATS），根据不同天气任务的收敛状态动态调整学习过程，从而缓解联合训练中的优化失衡问题。其次，它设计了一种带位置约束的动态提示Mamba模块（PDPMamba），该模块通过提示引导的语义扫描来建模退化感知式的修复状态，同时通过位置约束保持空间连续性。通过结合外部训练优化与内部特征建模，AWReNet能够更好地处理雨、雾、雪场景中的多样化退化模式。

本文的主要贡献如下：

• •
  我们提出了AWReNet，这是一种用于综合处理恶劣天气图像修复的统一框架，可在单一模型中有效处理雨、雾、雪等天气导致的图像退化。
• •
  我们引入了一种基于收敛性的任务调度策略，通过根据任务收敛状态动态调整任务权重，从而缓解多天气联合训练中的优化失衡问题。
• •
  我们设计了一种带位置约束的动态提示Mamba模块，该模块利用稀疏动态提示来建模细粒度的修复状态，并在位置约束下引导语义扫描。
• •
  在WeatherBench和UAV-Weather6k上的大量实验表明，AWReNet具有更出色的修复性能，而进一步的检测实验则证明了其在基于无人机的视觉巡检任务中的适用性。

章节节选

综合处理恶劣天气图像修复

现实世界中的图像常常会因雨、雪、雾等各种恶劣天气条件而受损。与那些专为去雨[1]、[2]、[3]，去雾[5]、[6]、[7]或除雪[8]、[9]设计的传统单任务方法不同，综合处理恶劣天气图像修复的方法[17]、[18]、[19]旨在在统一的神经网络中学习通用特征和特定于退化的特征，从而提高模型在复杂天气场景下的适应性。现有研究已经对此进行了探索

整体架构

图2展示了我们所提出的AWReNet，这是一种用于综合处理恶劣天气图像修复的统一框架。受Wave-Mamba[50]的启发，AWReNet采用了小波域的编码器-解码器架构，利用离散小波变换（DWT）和逆离散小波变换（IWT）来进行特征分解与重建[50]。AWReNet从训练优化和特征建模两个方面着手解决多天气修复问题。具体而言，CATS能够动态调整

数据集

我们在两个基准数据集上评估了所提出的AWReNet：WeatherBench[57]和UAV-Weather6k[58]。WeatherBench是一个真实的综合恶劣天气修复基准数据集，涵盖了白天和夜间不同场景下的雨、雪、雾退化情况。该数据集包含42,002对图像，分辨率为512$$512，其中41,402对用于训练，600对用于测试。UAV-Weather6k则是为基于无人机的多天气图像修复而设计的，包含雨、雾、雪等场景的图像

结论

在本文中，我们提出了AWReNet，这是一种用于综合处理恶劣天气图像修复的统一框架。AWReNet通过引入基于收敛性的任务调度来缓解不同天气任务之间的优化失衡问题，同时还设计了一种带位置约束的动态提示Mamba模块，以增强退化感知型的特征建模能力，同时保持空间连续性。在WeatherBench和UAV-Weather6k上的实验充分证明了AWReNet在雨、雾、雪天气条件下的有效性

CRediT作者贡献说明

Xianhao Wu：撰写——初稿、验证、软件开发、方法设计。Guili Xu：撰写——审阅与编辑、指导、资金申请、概念构思。Xiang Chen：撰写——审阅与编辑、方法设计。Qianfeng Yang：验证、数据整理。Qiyuan Guan：验证、形式化分析。Xianglong Zeng：可视化、数据整理。Junwei He：验证、研究分析。

利益冲突声明

作者声明不存在任何可能影响本文研究成果的已知财务利益或个人关系。

致谢

本研究部分得到了中国国家自然科学基金（项目编号：62073161）以及南京航空航天大学研究生科研实践创新计划（项目编号：xcxjh20250322）的支持。

Xianhao Wu目前在中国南京的南京航空航天大学自动化工程学院攻读硕士学位。他于2024年获得中国沈阳航空航天大学电子与信息工程学院的工学学士学位。他的当前研究兴趣包括图像修复、计算机视觉以及智能检测技术。