A two-stage-UNet network based on group normalization for single image deraining

Abstract

Rain streaks can seriously damage the optical quality of image and affect image processing in many scenes. Deep learning methods achieve state-of-the-art performance in the task of single-image rain removal. However, most deraining models based on deep learning only deal with local relationships, they didn’t sufficiently consider the contextual information over long distances in the task of rain removal. This drawback can lead to residual rain streaks and insufficient recovery of texture details. Therefore, a Two-Stage-UNet Network based on Group Normalization named TSUGN is created to solve these problems. It decomposes the image deraining task into easier and smaller subtasks for capturing more contextual information. And in order to balance spatial details and high-level contextual information, group normalization is also added to our Group Normalization Feature Residual Block (GNFRB). By fully taking into account of multi-scale features information, a Scale-Feature Fusion Module(SFFM)is proposed to learn features with different scales. In addition, a new feature compensation method is proposed to deal with the model bias issue by combining a parameter-free \(3-D\) attention module SimAM with GNFRB. Comprehensive experiments demonstrate the superiority of the proposed network in terms of computational efficiency, end-to-end trainability and easy implementation. It has great potential in image recovery tasks.