基于两阶段蒸馏学习的恶劣天气通用图像复原方法

本发明属于图像处理与计算机视觉，特别涉及一种雨、雪、雾霾等对成像产生干扰的恶劣天气退化图像通用复原方法，可应用于户外影像监控、雨雪天气下智能驾驶系统的高效率识别和决策。

背景技术：

1、近年来，计算机视觉技术迅速发展，在各个领域得到了广泛应用。但是，在实际生活和工作中，由于天气环境对光线折射、散射和遮挡的影响，通过成像设备获取到的可见光图像并不一定是清晰的，这对于人工监测和计算机监测产生了很大的干扰。因此，恶劣图像重建受到了世界各国专家的关注，被国内外研究人员广泛研究，在各领域涌现出了大量的工作。

2、目前，关于恶劣天气条件下的通用图像复原研究不多，各有优势和局限性。li等人在li r,tan r t,cheong l f.all in one bad weather removal usingarchitecturalsearch[c]//proceedings of the ieee/cvf conference on computervision and patternrecognition.2020:3175-3185中设计了包含多个与任务相关的编码器和一个通用解码器的结构。该方法能够利用自动架构搜索技术来优化编码器的特征提取。然而，该方法会显著增加模型规模，需要多个特征提取器来应对不同类型的天气。

3、chen等人在chen w t,huang z k,tsai c c,et al.learning multipleadverseweather removal via two-stage knowledge learning and multi-contrastiveregularization:toward a unified model[c]//proceedings of the ieee/cvf conference oncomputer vision and pattern recognition.2022:17653-17662中提出了一个二阶段的知识学习机制用于恶劣天气通用复原。其在第一阶段，由多个训练有素的教师网络指导学生网络集成和学习多种天气类型的知识，在第二阶段，利用多个对比正则项优化学生网络并改善其对于不同天气类型的识别能力。但是该方法由于在训练时包含多个教师模型，因而训练难度较大且会模型内存占用大。

4、zhu等人在zhu y,wang t,fu x,et al.learning weather-general andweather-specific features for image restoration under multiple adverseweatherconditions[c]//proceedings of the ieee/cvf conference on computervision and patternrecognition.2023:21747-2175中，其设计了两阶段训练策略探索天气类型的一般特征和特定特征。第一阶段将所有天气条件的数据集混合输入到网络，输出粗略恢复的结果，学习天气类型的一般特征；第二阶段自适应地拓展模型中每种天气类型的特定参数，得到多组能分别处理不同天气类别的精调模型参数。但是该方法的图像处理耗时较长，得到模型参数较多，并且在处理前需要预先明确退化图像所对应的天气类别，在实际的部署和智能化上受到限制。

5、综上所述，现有的恶劣天气退化图像的复原方法仍有以下两方面的缺点：

6、其一：均需要大量的计算资源和空间资源，所需资源配置较大且不利于实际的部署。

7、其二：网络结构复杂，对图像的处理效率低下，且大量的计算将导致图像复原的实时性能较差。

技术实现思路

1、本发明的目的在于针对上述现有技术的不足，提出一种基于两阶段蒸馏学习的恶劣天气退化图像通用复原方法，以减小模型占用的计算资源和空间资源，降低网络结构复杂度，提高图像实时复原性能。

2、实现本发明目的的技术方案是：通过使用知识蒸馏减少模型占用的计算资源和空间资源；通过仅使用天气退化感知模块和感知网络复原模块，降低网络结构复杂度；通过引入天气分类器和有监督对比正则，保证在轻型网络结构下也能实现快速的高质量图像复原。

3、根据上述思路，本发明的技术方案包括如下步骤：

4、(1)获取混合的恶劣天气数据，并划分为训练集和测试集；

5、(2)在pytorch框架下构建基于对比学习的两阶段蒸馏学习训练网络：

6、(2a)构建包括四个上下文残差块，三个卷积单元，一个通道维度合并单元，一个池化层和两个线性单元的天气退化感知器模块；

7、(2b)构建包括十八个动态transformer模块，两个卷积层，两个选择融合单元和两个上下采样单元的感知复原网络模块；

8、(2c)将天气退化感知器模块和感知复原网络模块级联，并将天气退化感知器模块与天气分类器级联得到一阶段训练网络；

9、(2d)将两个天气退化感知器模块并联，并将其中的学生天气退化感知器模块与感知复原网络模块级联得到二阶段蒸馏训练网络；

10、(3)构建一阶段训练网络和二阶段蒸馏训练网络损失：

11、(3a)构建由复原图像和真值图像的绝对值损失lr和天气分类损失lc组成的一阶段损失ls1：

12、ls1＝lr+αlc，其中，权重α是实验经验值

13、(3b)构建由复原图像和真值图像的绝对值损失lr，退化映射损失ldm，强-弱增强复原图像的绝对值损失l1和有监督对比损失lsc组成的二阶段损失ls2：

14、ls2＝lr+β1ldm+β2l1+β3lsc

15、其中，权重β1，β2，β3是实验经验值，分别取0.1，0.3和0.05；

16、(4)将数据集输入到一阶段训练网络之中，采用adam优化算法迭代更新网络参数，获得训练好的预训练天气退化感知器模型；

17、(5)将数据集输入到二阶段蒸馏训练网络之中，采用adam优化算法迭代更新网络参数，协同一阶段训练好的预训练天气退化感知器模型，获得训练好的学生退化感知器和感知复原网络模型；

18、(6)将测试集的多种退化的恶劣天气图像输入到训练好的学生退化感知器和感知复原网络模型中，输出复原后的图像。

19、本发明与现有技术相比，具有如下优点：

20、1)本发明由于构建了包含天气退化感知器模块和感知复原网络模块的网络，能够在简单的网络架构中准确获取图像的先验信息和退化特征，得到高质量的复原图像。

21、2)本发明由于引入天气分类损失，让退化特征更具辨识性，并在二阶段增加有监督对比损失，通过该损失增强教师与学生天气退化感知器的蒸馏结果一致性，能够进一步提升模型精度，压缩网络参数，提高模型的实时性。

22、仿真结果表明，本发明方法在1.5m左右的低参数量情况下，能提高合成恶劣天气图像的复原性能，在混合天气数据下其峰值信噪比psnr指标超越了现有先进算法，具有部署实际场景的潜力。