SDNET: MULTI-BRANCH FOR SINGLE IMAGE DERAINING
USING SWIN
前言
最近，流行的transformer具有全局計算特性，可以進一步促進圖像去雨任務的發(fā)展。本文首次將Swim-transformer引入圖像去雨領域，研究了Swim-transformer在圖像去雨領域的性能和潛力。具體來說，我們對Swim-transformer的基本模塊進行了改進，設計了一個三分支模型來實現(xiàn)單幅圖像的去雨。前者實現(xiàn)了基本的雨型特征提取，而后者融合不同的特征進一步提取和處理圖像特征。此外，我們還采用jump connection來融合深層特征和淺層特征。實驗表明，現(xiàn)有的公共數(shù)據(jù)集存在圖像重復和背景相對均勻的問題。因此，我們提出了一個新的數(shù)據(jù)集Rain3000來驗證我們的模型。
Transformer[28]最初是自然語言處理（NLP）領域的一個模型，用于并行處理單詞向量，以加速模型推理。它的全局計算特性適用于遠距離傳遞特征。這正是計算機視覺領域中卷積運算所不擅長的。Dosovitskiy等人[29]將圖像分割成16x16個圖像塊，將不同的圖像塊作為不同的詞輸入到transformer中，提高了圖像分類的精度。近年來，人們從深度[30]、多尺度[31]等角度應用transformer來完成相關任務。然而，Transformer也有不可忽視的缺點，例如計算量與圖像大小之間存在二次關系，這限制了它的應用環(huán)境。Liu等人[32]提出的Swin-transformer使用滑動窗口使模型具有線性計算復雜度，通過跨窗口連接改善了窗口間的信息交換，最終提高了模型在圖像分類、目標檢測和實例分割等方面的性能。
本文提出了一種新的圖像去雨網(wǎng)絡SDNet，它是利用Swim-transformer強大的特征表示能力構(gòu)建的端到端去雨網(wǎng)絡。具體地說，我們改進了Swim-transformer的基本模塊，重新設計了一個雙分支模型，實現(xiàn)了單圖像去雨。前者實現(xiàn)了基本的雨型特征提取，后者融合了不同分支的特征。此外，我們采用jump connection來融合深度特征和淺層特征，以提高網(wǎng)絡模型的性能。
本文貢獻如下：
- 改進了Swim-transformer模塊，使基本塊的輸出大小與特征提取的輸入一致。
- 為了提高特征融合的效率，提出了一種三分支網(wǎng)絡模型，并利用額外的jump connection來融合深特征和淺特征，以提高網(wǎng)絡的學習能力。
- 提出了一個全新的數(shù)據(jù)集。并利用該數(shù)據(jù)集對本文提出的模型和其他主流降雨模型進行了訓練和檢驗。
- 實驗表明，與主流去雨模型相比，該模型具有一定的性能和推理速度優(yōu)勢
最近有大量的研究工作將transformer引入CV域，并取得了良好的效果。具體來說，Dosovitskiy等人[29]將圖像分成16X16個圖像塊，然后將其拉伸成一維向量，然后送入網(wǎng)絡中完成圖像分類任務。Chen等人[38]提出了一種基于卷積運算的transformer與Unet相結(jié)合的TransUnet方法，實現(xiàn)醫(yī)學圖像的分割。蔣等[39]設計了與對抗生成網(wǎng)絡結(jié)構(gòu)相同的圖像生成transformer。transformer中的self-attention導致模型計算直線增長，導致transformer不能在低計算能力的硬件上運行。Liu[32]提出了一種利用滑動窗口方法使網(wǎng)絡計算線性增長并加速網(wǎng)絡推理的方法。我們的方法是基于這種方法來實現(xiàn)一個單一的圖像去雨任務的融合特征。
本文的方法是基于這種方法[32]來實現(xiàn)一個單一的圖像去雨任務，融合不同分支的特征、深度特征和淺層特征。
方法
Swin-transformer
Transformer是一個功能強大的網(wǎng)絡模塊，可以取代CNN操作。但其中的Muti-Head Attention導致模型的計算量迅速增加，導致transformer模型無法在許多底層硬件中測試和使用，注意力的數(shù)學表達式如下：
其中
Network structure
本文使用一個簡單而強大的前饋網(wǎng)絡作為主干，如圖2所示。SDnet網(wǎng)絡基本上由三個多分支融合模塊組成，稱為MSwt，一個多分支模塊MSwt-m和兩個基本block模塊。此外，還增加了跳轉(zhuǎn)連接，目的是融合深特征和淺特征，以提高網(wǎng)絡去雨的性能。
受Swin-transformer的啟發(fā)，本文的Basic-block模塊設計如圖3所示。輸入圖像或特征圖的大小為231× 231× 3，并且在PatchPartition中patch block的大小設置為3× 3，因此矩陣的size of the position encoding 是
為了更靈活地構(gòu)建網(wǎng)絡，提出了Basic-block的概念，并設計了兩個三分支特征融合塊。如圖4和圖5所示，與后者相比，前者有一個用于融合特征的附加基本塊。數(shù)學表達式如下：
其中F（·）表示基本塊的操作。x表示模塊Mswt的輸入。這種設計的思想來源于自我注意中的多頭注意機制。通過學習F1、F2、F3，可以自適應地學習不同的特征。將輸入映射到不同的子空間，分別提取不同的特征。與自我注意不同的是，我們對提取的特征求和，而不是級聯(lián)操作。通過F4融合增加的特征，實現(xiàn)進一步的特征提取。由于設計思想來源于多頭注意機制，多分支具有與該機制相同的特點，即在一定范圍內(nèi)，分支越多，模型性能越好。為了平衡模型的規(guī)模和模型的性能，我們選擇了三個分支進行特征提取。
雖然transformer可以保持特征在長距離傳播，但是仍然需要在網(wǎng)絡中結(jié)合深特征和淺特征，為此我們設計了一個沒有特征融合的Mswt模塊，我們稱之為Mswt-m，如圖5所示，其數(shù)學表達式如下：
F1、F2、F3將輸入映射到三個不同的空間進行特征提取，對提取的特征求和，然后與第二個Mswt模塊的輸出求和，再經(jīng)過一個基本塊，實現(xiàn)深度特征和淺層特征的融合，如圖2中的小跳躍連接所示，而圖2中相對較長的跳躍連接則考慮了主要特征中包含的豐富的空間和紋理信息，有助于完成深度特征中缺失的紋理信息。
Loss function
其中，O為雨圖像，B為對應標簽。是絕對差（SAD）之和，用于計算相似預測圖像和標簽之間的像素損失，如等式6所示。SSIM（結(jié)構(gòu)相似性）是結(jié)構(gòu)相似性，最初用作評估兩個圖像內(nèi)容的結(jié)構(gòu)相似性的度量。Ren等人[41]證明了SSIM作為損失函數(shù)在圖像降額任務中的有效性的負面作用，其數(shù)學表達式如等式7所示。盡管使用該損失函數(shù)可以獲得高SSIM度量，但圖像仍然存在失真和低峰值信噪比（PSNR）。identity loss（等式8）由CycleGAN[42]導出，CycleGAN[42]用于約束生成圖像的顏色丟失，這里我們使用它來約束圖像去雨后的圖像樣式，這減少了圖像失真，提高了網(wǎng)絡性能。α , β , λ 是SAD損失、SSIM損失和identity loss的系數(shù)。在本文中，分別設置為0.2、4和1。
Experimental environment
實驗使用Tesla V100 16G GPU進行訓練，使用Pytorch框架1.7.0和（Adam）[43]，初始學習率為5× 10?4，減少到5× 10?5和5× 10?6當訓練迭代次數(shù)分別為總迭代次數(shù)的3/5和4/5時。輸入模型的圖像大小設置為231×231. batch size為5。
我們提出了一個全新的數(shù)據(jù)集用于網(wǎng)絡訓練和消融實驗。該數(shù)據(jù)集是從ImageNet中隨機抽取的10萬幅圖像，保證了圖像的多樣性。從Efficientderain[12]降雨模式數(shù)據(jù)集中隨機選擇一到四種降雨模式，并添加到選定的圖像中。我們最終選擇了3000張合成圖像作為訓練集，400張作為測試集。我們把這個數(shù)據(jù)集命名為Rain3000。此外，我們還使用公開的數(shù)據(jù)集Rain100L和Rain100H[44]來驗證SDnet模型。兩個公開的數(shù)據(jù)集都包含1800個訓練圖像和200個測試圖像。
使用SSIM和PSNR作為評價指標，這兩種指標已被廣泛用于評價預測圖像的質(zhì)量。PSNR是根據(jù)兩幅圖像之間的像素誤差來計算的，誤差越小，值越大，圖像越相似，除雨效果越好。相反，圖像去雨的效果越差
結(jié)論
首先，本文提出了一種基于Swin-transformer的三分支端到端除雨網(wǎng)絡，它充分利用了Swin-transformer強大的學習能力，用一種改進的Swin-transformer代替卷積運算，并設計了一個多分支模塊來融合不同空間域的信息，使用跳轉(zhuǎn)連接來融合深特征和淺特征。此外，我們提出了一個新的數(shù)據(jù)集，由3000個訓練對和400個測試對組成。該數(shù)據(jù)集是基于ImageNet生成的，具有豐富的背景和雨型組合，便于模型的推廣。我們提出的模型在數(shù)據(jù)集Rain3000和公共數(shù)據(jù)集Rain100L、Rain100H上都達到了最佳性能。我們的工作還有些不足。例如，在參數(shù)數(shù)目相同的情況下，哪種方法更適合于并行或串行的圖像去噪任務還沒有詳細探討。以及是否可以使用多個不同大小的滑動窗口來實現(xiàn)窗口間的進一步信息交換，以提高網(wǎng)絡降容的性能。此外，我們正在使用更簡單的前饋網(wǎng)絡，更復雜的網(wǎng)絡仍然值得研究
