UniRain: Unified Image Deraining with RAG-based Dataset Distillation and Multi-objective Reweighted Optimization

Qianfeng Yang; Qiyuan Guan; Xiang Chen; Jiyu Jin; Guiyue Jin; Jiangxin Dong

cs.CVMar 4, 2026

UniRain: Unified Image Deraining with RAG-based Dataset Distillation and Multi-objective Reweighted Optimization

Authors: Qianfeng Yang, Qiyuan Guan, Xiang Chen, Jiyu Jin, Guiyue Jin, Jiangxin Dong

Organizations: Dalian Polytechnic University · Nanjing University of Science and Technology

Abstract

Despite significant progress has been made in image deraining, we note that most existing methods are often developed for only specific types of rain degradation and fail to generalize across diverse real-world rainy scenes. How to effectively model different rain degradations within a universal framework is important for real-world image deraining. In this paper, we propose UniRain, an effective unified image deraining framework capable of restoring images degraded by rain streak and raindrop under both daytime and nighttime conditions. To better enhance unified model generalization, we construct an intelligent retrieval augmented generation (RAG)-based dataset distillation pipeline that selects high-quality training samples from all public deraining datasets for better mixed training. Furthermore, we incorporate a simple yet effective multi-objective reweighted optimization strategy into the asymmetric mixture-of-experts (MoE) architecture to facilitate consistent performance and improve robustness across diverse scenes. Extensive experiments show that our framework performs favorably against the state-of-the-art models on our proposed benchmarks and multiple public datasets.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of motivation . (a) Rainy image samples from public datasets, illustrating the noticeable differences in data quality. (b) Directly merging existing synthetic and real datasets enlarges data volume, yet quality disparity hinders performance, as shown by PSNR results. (c) The loss curves of DRS, DRD, NRS, and NRD (denoting daytime/nighttime rain streaks and raindrops) show different convergence rates, leading to imbalance in unified training. (d) The PSNR curves indicate that the model tends to favor simpler degradations but struggles with complex ones.

Figure 2: Overall framework of UniRain . (Left) The RAG-based dataset distillation pipeline retrieves real rainy references consistent with the query image via multi-level similarity search and employs vision language models to evaluate its quality, thereby distilling reliable samples from public datasets. (Right) The asymmetric MoE architecture consists of soft-MoE encoder and hard-MoE decoder, optimized via the multi-objective reweighted strategy to achieve balanced learning and robust performance across multiple rain degradation types.

Datasets		RainRAG Dataset
Datasets		DRS			DRD			NRS			NRD			Average
Methods	Venue	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
LQ	-	23.79	0.7514	0.3207	20.79	0.7293	0.2433	26.89	0.8480	0.3583	21.05	0.5590	0.6130	23.13	0.7220	0.3839
PReNet [ 33 ]	CVPR 2019	26.73	0.8233	0.1885	21.45	0.7358	0.2697	32.85	0.9520	0.1957	24.29	0.8058	0.3751	26.33	0.8292	0.2572
RCDNet [ 39 ]	CVPR 2020	25.27	0.7949	0.2777	21.46	0.7295	0.2778	28.44	0.9070	0.2996	20.62	0.5832	0.5931	23.95	0.7537	0.3620
MPRNet [ 50 ]	CVPR 2021	27.99	0.8242	0.1701	22.50	0.7264	0.3308	33.75	0.9541	0.1834	24.46	0.8118	0.3634	27.17	0.8291	0.2619

Table 1: Quantitative evaluations on the proposed RainRAG dataset, where DRS, DRD, NRS, and NRD denote daytime rain streaks, daytime raindrops, nighttime rain streaks, and nighttime raindrops, respectively. The best and second-best values are bold and underlined .

Panel of Figure 3 — Figure 3: Visual comparison of image deraining results on the RainRAG-NRS and RainDS-real-RD datasets. Zoom in for a better view.

Datasets	Metrics	PReNet [ 33 ]	RCDNet [ 39 ]	MPRNet [ 50 ]	IDT [ 44 ]	Restormer [ 51 ]	DRSformer [ 5 ]	RLP [ 53 ]	MSDT [ 4 ]	NeRD-Rain [ 7 ]	URIR [ 45 ]	UniRain (Ours)
RealRain-1k-L [ 26 ]	PSNR $\uparrow$	28.80	28.58	33.89	33.71	32.44	33.65	32.02	33.56	33.98	33.54	36.51
RealRain-1k-L [ 26 ]	SSIM $\uparrow$	0.9307	0.9202	0.9701	0.9664	0.9644	0.9625	0.9558	0.9640	0.9710	0.9686	0.9778
RealRain-1k-H [ 26 ]	PSNR $\uparrow$	25.49	25.46	30.83	30.37	28.80	30.35	29.01	30.91	31.37	31.10	33.74
RealRain-1k-H [ 26 ]	SSIM $\uparrow$	0.8820	0.8613	0.9542	0.9441	0.9389	0.9360	0.9309	0.9451	0.9583	0.9550	0.9672
RainDS-real-RD [ 31 ]	PSNR $\uparrow$	19.08	19.37	19.48	20.35	20.74	20.15	20.86	20.72	21.01	20.44	22.07
RainDS-real-RD [ 31 ]	SSIM $\uparrow$	0.6020	0.5937	0.5957	0.6237	0.6346	0.6133	0.6136	0.6205	0.6155	0.5945	0.6436

Table 2: Quantitative evaluations on the real-world public benchmarks. The best and second-best values are bold and underlined .

Panel of Figure 4 — Figure 4: Visual comparison of image restoration results on the real-world benchmark ( i.e. , WeatherBench [ 14 ] ). Compared to the state-of-the-art methods, our UniRain restores a high-quality image with clear details, even outperforming the GT by removing residual raindrops.

Methods	Restormer [ 51 ]	DRSformer [ 5 ]	NeRD-Rain [ 7 ]	UniRain (Ours)
FLOPs (G)	140.990	220.378	147.978	126.541
Params (M)	26.097	33.627	22.856	24.388

Table 3: Comparisons of model complexity against state-of-the-art methods. The size of the test image is

256\times 256

pixels.

Panel of Figure 5 — Figure 5: Generalization results across multiple scenarios. Our method achieves the best restoration results in driving scenes (first row), UAV scenes (second row), and maritime scenes (third row).

Panel of Figure 6 — Figure 6: Statistical ablation analysis of different combination strategies and their feature distributions across combinations.

Variants	only VLM	w/o Retrieval stage	w/o Generation stage	UniRain (Ours)
PSNR $\uparrow$	27.73	28.39	28.36	28.93
SSIM $\uparrow$	0.8358	0.8453	0.8425	0.8515

Table 4: Ablation analysis for RAG-based dataset distillation.

Base Models	Optimization Strategy			Metrics
Base Models	Continual learning	Prompt learning	Ours	PSNR $\uparrow$	SSIM $\uparrow$
PIGWM [ 55 ]	✔	✗	✗	18.38	0.5758
PIGWM [ 55 ]	✗	✗	✔	18.68	0.5880
PromptIR [ 29 ]	✗	✔	✗	18.58	0.5639
PromptIR [ 29 ]	✗	✗	✔	19.55	0.5759
UniRain (Ours)	✔	✗	✗	18.92	0.5657

Table 5: Ablation analysis on various models using different optimizers, including our adaptive reweighting optimization strategy.

Panel of Figure 7 — Figure 7: Loss and PSNR of the optimization strategy variants.

Variants	Encoder		Decoder		Metrics
Variants	hard-MoE	soft-MoE	hard-MoE	soft-MoE	PSNR $\uparrow$	SSIM $\uparrow$
Baseline	✗	✗	✗	✗	27.54	0.8442
Variant 1	✔	✗	✔	✗	28.39	0.8471
Variant 2	✔	✗	✗	✔	28.15	0.8485
Variant 3	✗	✔	✗	✔	27.91	0.8465
UniRain (Ours)	✗	✔	✔	✗	28.93	0.8515

Table 6: Ablation analysis of asymmetric MoE architecture. Experiments are conducted on the proposed RainRAG dataset.

Methods	TransWeather [ 38 ]	WGWS-Net [ 56 ]	Histoformer [ 37 ]	UniRain (Ours)
PSNR $\uparrow$	24.70	23.89	24.59	26.01
SSIM $\uparrow$	0.7931	0.7811	0.7976	0.8032

Table 7: Extension to all-in-one weather restoration on the WeatherBench dataset [ 14 ] , which contains real-world multi-weather degradations (rain, snow, and haze). Average results are reported.

Panel of Figure 8 — Figure 8: Visual comparison of all-in-one weather restoration results ( e.g. , hazy input (first row) and snowy input (second row)).

Image Enhancement

Image enhancement restores visibility, color, detail, and dynamic range in degraded images for human viewing and downstream vision tasks.

See all 149 papers

Derain-Agent: A Plug-and-Play Agent Framework for Rainy Image Restoration
Mar 12, 2026Zhaocheng Yu, Xiang Chen, Runzhe Li +4Agent-Based Image Restoration Image Restoration
Unpaired Image Deraining Using Reward-Guided Self-Reinforcement Strategy
May 1, 2026Yinghao Chen, Yeying Jin, Xiang Chen +3Self-Training Image Restoration
1National University of Defense Technology · National University of Singapore · 3Nanjing University of Science and Technology +2
Unified Removal of Raindrops and Reflections: A New Benchmark and A Novel Pipeline
Mar 17, 2026Xingyu Liu, Zewei He, Yu Chen +4Diffusion Models for Image Restoration Image Enhancement
Huanjiang Laboratory, Zhuji, China · School of Aeronautics and Astronautics, Zhejiang University, Hangzhou, China · Hangzhou Institute of Technology, Xidian University, Hangzhou, China +1

Mar 12, 2026J/K move · Enter open · S save

Derain-Agent: A Plug-and-Play Agent Framework for Rainy Image Restoration

Zhaocheng Yu, Xiang Chen, Runzhe Li, Zihan Geng+3

Agent-Based Image Restoration Image Restoration Image Enhancement

Open paper AskPDF arXiv

While deep learning has advanced single-image deraining, existing models suffer from a fundamental limitation: they employ a static inference paradigm that fails to adapt to the complex, coupled degradations (e.g., noise artifacts, blur, and color deviation) of real-world rain. Consequently, restored images often exhibit residual artifacts and inconsistent perceptual quality. In this work, we present Derain-Agent, a plug-and-play refinement framework that transitions deraining from static processing to dynamic, agent-based restoration. Derain-Agent equips a base deraining model with two core capabilities: 1) a Planning Network that intelligently schedules an optimal sequence of restoration tools for each instance, and 2) a Strength Modulation mechanism that applies these tools with spatially adaptive intensity. This design enables precise, region-specific correction of residual errors without the prohibitive cost of iterative search. Our method demonstrates strong generalization, consistently boosting the performance of state-of-the-art deraining models on both synthetic and real-world benchmarks.

Abstract

Image Enhancement

Explore similar work

Derain-Agent: A Plug-and-Play Agent Framework for Rainy Image Restoration

Unpaired Image Deraining Using Reward-Guided Self-Reinforcement Strategy

Unified Removal of Raindrops and Reflections: A New Benchmark and A Novel Pipeline