Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration

Guanglu Dong; Chunlei Li; Chao Ren; Jingliang Hu; Yilei Shi; Xiao Xiang Zhu; Lichao Mou

cs.CVMar 2, 2026

Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration

Authors: Guanglu Dong, Chunlei Li, Chao Ren, Jingliang Hu, Yilei Shi, Xiao Xiang Zhu, Lichao Mou

Organizations: Sichuan University · MedAI Technology (Wuxi) Co. Ltd. · Technical University of Munich

Abstract

Recently, significant breakthroughs have been made in all-in-one image restoration (AiOIR), which can handle multiple restoration tasks with a single model. However, existing methods typically focus on a specific image domain, such as natural scene, medical imaging, or remote sensing. In this work, we aim to extend AiOIR to multiple domains and propose the first multi-domain all-in-one image restoration method, DATPRL-IR, based on our proposed Domain-Aware Task Prompt Representation Learning. Specifically, we first construct a task prompt pool containing multiple task prompts, in which task-related knowledge is implicitly encoded. For each input image, the model adaptively selects the most relevant task prompts and composes them into an instance-level task representation via a prompt composition mechanism (PCM). Furthermore, to endow the model with domain awareness, we introduce another domain prompt pool and distill domain priors from multimodal large language models into the domain prompts. PCM is utilized to combine the adaptively selected domain prompts into a domain representation for each input image. Finally, the two representations are fused to form a domain-aware task prompt representation which can make full use of both specific and shared knowledge across tasks and domains to guide the subsequent restoration process. Extensive experiments demonstrate that our DATPRL-IR significantly outperforms existing SOTA image restoration methods, while exhibiting strong generalization capabilities. Code is available at https://github.com/GuangluDong0728/DATPRL-IR.

Figures & tables

Figure 1: This paper makes a preliminary exploration of multi-domain all-in-one image restoration (MD-AiOIR), aiming at further extending the restoration capability of a single model to a broader range of tasks and image domains, including natural scene, medical imaging, and remote sensing.

Figure 2: Framework of the proposed DATPRL-IR for multi-domain all-in-one image restoration. By introducing domain-aware task representation learning, DATPRL-IR can fully utilize both specific and shared knowledge across tasks and domains, effectively reducing the learning difficulty of the model and improving its performance.

Figure 3: A partial visualization of the word clouds generated from the text descriptions produced by LLAVA, and the t-SNE clustering analysis of the text descriptions corresponding to the 9 datasets from different domains and tasks. It can be observed that images from different domains exhibit their own characteristics while also sharing certain overlapping features.

Image Domain		Natural Image				Medical Image				Remote Sensing Image				Average Performance
Task & Dataset		Super-Resolution		Deraining		MRI SR		CT Denoising		RSI SR		Cloud Removal
Task & Dataset		on DIV2K-VAL		on Rain100L		on IXI MRI		on AAPM-Mayo		on UCMerced		on CUHK CR1
Method	Year	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Single-Task Method
MPRNet	CVPR2021	28.82	0.8115	38.07	0.9817	26.84	0.8891	33.60	0.9259	27.70	0.7730	25.35	0.7389	30.06	0.8534
SwinIR	ICCVW2021	28.61	0.8051	36.07	0.9736	26.06	0.8766	33.51	0.9243	27.29	0.7545	24.36	0.6552	29.32	0.8482

Table 1: Quantitative comparison between our method and other SOTA methods on 3 domains & 6 tasks experimental setting. The best and second-best metrics are highlighted in bold and underline .

Figure 4: Comparison of our DATPRL-IR with other SOTA methods on 6-task and 3-domain setting.

Figure 5: Comparison of our DATPRL-IR with other SOTA methods on 9-task and 3-domain setting.

Task & Dataset	Natural SR		Deraining		CT Denoising		PET Synthesis		RSI SR		RSI Dehazing
Task & Dataset	on DIV2K-Val		on Rain100L		on AAPM		on PolarStar		on UCMerced		on RICE1
Method	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
LLaVA1.5-13B	29.04	0.8180	39.71	0.9868	33.77	0.9272	37.08	0.9500	28.31	0.7918	26.92	0.9347
LLaVA1.5-7B	29.05	0.8181	39.67	0.9867	33.77	0.9273	37.12	0.9502	28.31	0.7913	26.94	0.9347
Qwen3-VL-2B-Instruct	29.07	0.8182	39.69	0.9868	33.77	0.9270	37.11	0.9501	28.31	0.7914	26.82	0.9351

Table 4: Effectiveness on different MLLMs on 3 domains & 9 tasks experimental setting. The best metrics are highlighted in bold .

Task & Dataset	Deraining		Deblurring		MRI SR		PET Synthesis		Cloud Removal		RSI Dehazing
Task & Dataset	on Rain100L		on GoPro		on IXI MRI		on PolarStar		on CUHK CR1		on RICE1
Method	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Explicit Domain Prompts	39.59	0.9865	29.31	0.8839	27.77	0.9035	37.08	0.9500	25.76	0.7573	25.91	0.9296
Simple Domain Encodings	39.52	0.9863	29.21	0.8814	27.70	0.9031	37.02	0.9496	25.81	0.7566	25.81	0.9213
Domain Prompt Pool (Ours)	39.67	0.9867	29.57	0.8881	27.86	0.9045	37.12	0.9502	26.00	0.7592	26.94	0.9347

Table 5: Effectiveness on different prompt designs on 3 domains & 9 tasks experimental setting. The best metrics are highlighted in bold .

Panel of Figure 6 — Figure 6: In-depth analysis of our method. Zoom in for better visualization.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Datasets	DF2K	Rain100L	GoPro	IXI MRI	AAPM-Mayo	PolarStar m660	UCMerced
Train	3450	200	2103	40500	18351	27837	1800
Test	100	100	1111	11400	211	2044	300
Tasks	Natural SR	Natural Deraining	Natural Deblurring	MRI SR	CT Denoising	PET Synthesis	RSI SR
Datasets	CUHK CR1	RICE1
Train	534	400
Test	134	100

Appendix

Table 6: Detailed description of the datasets utilized.

Medical Images

“The image is a black and white photo of a brain, with a close up view of the frontal lobe.”

“The image is a black and white picture of a human body, specifically focusing on the abdominal area.”

“The image is a grayscale medical image with low overall contrast, featuring a bright background and a few small dark foci near the center.”

Remote Sensing Images

“The image is a large, green field with a white snowy ground, and it is a winter scene.”

“The image is a bird’s-eye view of a forest with low brightness, where a river flows through it and a highway is visible nearby.”

Appendix

Table 7: Examples of LLaVA-generated textual descriptions for medical, remote sensing, and natural images.

Image Domain	Natural Image			Medical Image			Remote Sensing Image
Task & Dataset	SR	Deraining	Deblurring	MRI SR	CT Denoising	PET Synthesis	RSI SR	Cloud Removal	RSI Dehazing
Task & Dataset	on DIV2K-Val	on Rain100L	on GoPro	on IXI MRI	on AAPM-Mayo	on PolarStar m660	on UCMerced	on CUHK CR1	on RICE1
Single-Task Method
MPRNet	28.32 / 0.8067	37.55 / 0.9797	28.02 / 0.8570	26.69 / 0.8871	33.54 / 0.9253	36.72 / 0.9475	27.47 /0.7646	25.20 / 0.7334	25.66 / 0.9268
SwinIR	28.80 / 0.8109	37.55 / 0.9802	28.17 / 0.8510	26.49 / 0.8844	33.63 /0.9251	36.78 / 0.9468	27.52 / 0.7670	25.39 / 0.7249	25.42 / 0.9244
Restormer	28.63 / 0.8150	38.45 / 0.9833	29.06 / 0.8805	27.43 / 0.8992	33.70 / 0.9269	37.20 / 0.9509	27.94 / 0.7827	25.66 / 0.7451	26.07 / 0.9286

Appendix

Table 8: Quantitative comparison between our method and other SOTA methods on 3 domains & 9 tasks experimental setting. The best and second-best metrics are highlighted in bold and underline .

Figure 7: Comparison of our DATPRL-IR with other SOTA methods on Natural Images.

Figure 8: Comparison of our DATPRL-IR with other SOTA methods on Remote Sensing Images.

Figure 9: Comparison of our DATPRL-IR with other SOTA methods on Medical Images.

Image Domain & Task		AI-Generated Content (AIGC) SR						Comic Image SR
Dataset		on SDXL-Turbo		on DeepFloyd		on Midjourney		on Manga109
Method	Year	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Single-Task Method
MPRNet	CVPR2021	31.13	0.8730	26.76	0.8254	27.19	0.7656	27.13	0.8612
SwinIR	ICCVW2021	31.20	0.8709	27.41	0.8344	27.19	0.7627	27.76	0.8697
Restormer	CVPR2022	32.95	0.8985	28.81	0.8675	28.24	0.8005	28.13	0.8758

Appendix

Table 9: Quantitative zero-shot super-resolution comparison between our method and other SOTA methods on AIGC and comic images. The best metrics are highlighted in bold .

Task & Dataset		Highlight Removal		Image Desnowing
Task & Dataset		on SHIQ		on Snow100K-S
Method	Year	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Single-Task Method
MPRNet	CVPR2021	25.83	0.9297	26.09	0.8882
Restormer	CVPR2022	25.67	0.9295	26.57	0.8955
NAFNet	ECCV2022	25.89	0.9298	26.27	0.8923

Appendix

Table 10: Quantitative zero-shot comparison between our method and other SOTA methods on unseen image specular highlight removal, and image desnowing. The best metrics are highlighted in bold .

Task & Dataset		Real-World Deraining		RSI Deblurring
Task & Dataset		on SPA-Data		on UCMerced
Method	Year	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Single-Task Method
MPRNet	CVPR2021	31.99	0.9263	26.70	0.7741
Restormer	CVPR2022	32.35	0.9282	27.21	0.7928
NAFNet	ECCV2022	31.99	0.9249	26.92	0.7864

Appendix

Table 11: Quantitative generalization comparison between our method and other SOTA methods on real-world image deraining and RSI deblurring. The best metrics are highlighted in bold .

Figure 10: Zero-shot comparison of our DATPRL-IR with other SOTA methods on AIGC and comic domains. Please zoom in for better visualization.

Figure 11: Zero-shot comparison of our DATPRL-IR with other SOTA methods on specular highlight removal task. Please zoom in for better visualization.

Figure 12: Zero-shot comparison of our DATPRL-IR with other SOTA methods on image desnowing task. Please zoom in for better visualization.

Figure 13: Generalization comparison of our DATPRL-IR with other SOTA methods on RSI deblurring task. Please zoom in for better visualization.

Task & Dataset	Deraining		Deblurring		CT Denoising		RSI SR		Cloud Removal		RSI Dehazing
Task & Dataset	on Rain100L		on GoPro		on AAPM		on UCMerced		on CUHK CR1		on RICE1
Method	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Sequential Fusion (domain $\rightarrow$ image $\rightarrow$ task)	39.60	0.9865	29.30	0.8832	33.76	0.9271	28.27	0.7902	25.82	0.7577	26.76	0.9376
Ours	39.67	0.9867	29.57	0.8881	33.77	0.9273	28.31	0.7913	26.00	0.7592	26.94	0.9347

Appendix

Table 13: Effectiveness on different fusion strategies on 3 domains & 9 tasks experimental setting. The best metrics are highlighted in bold .

Image Restoration

Image restoration uses computational methods to recover useful visual content from noisy, blurred, low-resolution, or otherwise degraded images.

See all 372 papers

Beyond Uniform Restoration: Empowering All-in-One Restoration with Pixel-Level Multimodal Guidance
Aug 10, 2026Chunxiao Liu, Wei Liu, Anbin Xiong +1All-in-One Image Restoration Image Restoration
Xiaomi Corporation Beijing, China
TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback
Sep 25, 2026Yanjie Tu, Qingsen Yan, Axi Niu +4All-in-One Image Restoration Image Restoration
Northwestern Polytechnical University · Shenzhen Research Institute of Northwestern Polytechnical University · Xi’an University of Architecture and Technology
Universal Image Restoration via Internalized Chain-of-Thought Reasoning
Jun 16, 2026Yu Guo, Zhengru Fang, Shengfeng He +4All-in-One Image Restoration Image Restoration
Hong Kong JC Lab of Smart City and the Department of Computer Science, City University of Hong Kong · School of Computing and Information Systems, Singapore Management University · Computer Science and Information Engineering, National Taiwan University

Aug 10, 2026J/K move · Enter open · S save

Beyond Uniform Restoration: Empowering All-in-One Restoration with Pixel-Level Multimodal Guidance

Chunxiao Liu, Wei Liu, Anbin Xiong, Erli Meng

Xiaomi Corporation Beijing, China

All-in-One Image Restoration Image Restoration Visual Prompt Learning

Open paper AskPDF arXiv

All-in-one image restoration is a unified low-level vision task that aims to effectively recover high-quality images from inputs degraded by various types and levels of corruption using a single model. Recent works have achieved remarkable progress by learning degradation-adaptive prompts or network architectures. However, these methods typically apply a uniform restoration strategy across the entire image, neglecting the fact that different regions may suffer from distinct degradation types and varying degrees of severity. In contrast, we propose to perform restoration at the pixel level, thereby enabling more fine-grained and precise control over the restoration process. Specifically, we present MGN-AIR, a novel pixel-level restoration framework for all-in-one image restoration. Our approach first learns to estimate a pixel-level visual prompt. Then, it leverages both textual and visual prompts to provide global and local degradation cues, guiding the model on where to look and how to restore at each pixel. We conduct extensive experiments on multiple all-in-one image restoration benchmarks, covering a wide range of tasks including denoising, deraining, deblurring, dehazing, desnowing, and low-light enhancement. Experimental results demonstrate that our proposed method consistently and significantly outperforms existing approaches.

Abstract

Image Restoration

Explore similar work

Beyond Uniform Restoration: Empowering All-in-One Restoration with Pixel-Level Multimodal Guidance

TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback

Universal Image Restoration via Internalized Chain-of-Thought Reasoning