Research map / Image generation / Image restoration
UniRestore: Unified Perceptual and Task-Oriented Image Restoration Model Using Diffusion Prior
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UniRestore, a unified image restoration model built on Stable Diffusion that uses encoder-feature restoration (CFRM) and a prompt-conditioned Task Feature Adapter (TFA) to serve both perceptual and task-oriented restoration.
Problem
Perceptual image restoration improves visual quality but often hurts downstream vision tasks, while task-oriented restoration improves recognition or segmentation but yields less visually pleasing images; a single model is needed that serves both.
Approach
Adapt a pretrained Stable Diffusion prior by restoring degraded VAE encoder features with a Complementary Feature Restoration Module and fusing them into the decoder via a lightweight task-prompt-driven Task Feature Adapter, so new tasks only require a new learnable prompt.
Key results
- UniRestore reaches 26.00 average PSNR across seen and unseen PIR datasets, versus 24.85 for PromptIR and 24.68 for DiffUIR.
- On CUB unseen classification, UniRestore improves ResNet-50 accuracy by 20.01% over the low-quality input (33.69 to 53.70).
- Object detection on RTTS reaches 58.06 mAP with a newly added prompt, versus 54.29 for PromptIR and 45.63 for low-quality inputs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use UniRestore instead of the strongest listed baseline | image restoration methods · DIV2K test | PSNR: 23.90 -> 24.32 | Helped |
| Use UniRestore instead of the strongest listed baseline | image restoration methods · Rain100L | PSNR: 28.17 -> 30.02 | Helped |
| Use UniRestore instead of the strongest listed baseline | image restoration methods · CUB | accuracy (ResNet-50): 49.64 -> 53.70 | Helped |
| Use UniRestore instead of the strongest listed baseline | image restoration methods · Cityscapes test | mIoU (DeepLabv3+) | Helped |
| Add CFRM feature restoration (ablation) | UniRestore · DIV2K test | PSNR: 21.43 -> 24.32 | Helped |
| Add CFRM feature restoration (ablation) | UniRestore · ImageNet test | accuracy (ResNet-50): 63.10 -> 71.65 | Helped |
| Add CFRM feature restoration (ablation) | UniRestore · Cityscapes test | mIoU: 55.48 -> 66.05 | Helped |
| Add TFA to adapt encoder features (ablation) | UniRestore · DIV2K test | PSNR: 22.16 -> 24.32 | Helped |
| Add TFA to adapt encoder features (ablation) | UniRestore · ImageNet test | accuracy (ResNet-50): 64.25 -> 71.65 | Helped |
| Add TFA to adapt encoder features (ablation) | UniRestore · Cityscapes test | mIoU: 58.13 -> 66.05 | Helped |
| Use one task-specific prompt with one TFA rather than a shared single prompt (ablation) | UniRestore · ImageNet test | accuracy (ResNet-50): 70.05 -> 71.65 | Helped |
| Use one shared TFA rather than multiple TFAs (ablation) | UniRestore · DIV2K test | PSNR: 25.48 -> 24.32 | Hurt |
Surprising findings
- Fine-tuning baseline PIR and TIR methods on multiple downstream tasks (the * variants) generally decreased their performance on each task, which the authors attribute to lack of a mechanism for learning multiple objectives.
Limitations
- Performance on segmentation under the multi-task setting for some baselines is limited, and UniRestore's segmentation results on unseen ACDC lag the HQ reference.
- Extension to new tasks is demonstrated only for object detection on a single dataset (RTTS).
Experimental setup
Datasets: DIV2K Flickr2K OST ImageNet Cityscapes COCO Urban100 BSD68 CBSD68 Kodak McMaster Set12
Benchmarks: DIV2K Rain100L RESIDE UHDSnow Noise (Urban100, BSD68, CBSD68, Kodak, McMaster, Set12) GoPro ImageNet CUB Cityscapes FoggyCityscapes ACDC RTTS
Metrics: PSNR SSIM accuracy mIoU mAP
Baselines: DIP URIE NAFNet PromptIR DiffBIR DiffUIR
Builds on: Stable Diffusion ControlNet-style Controller (Controller) SC-Tuner NAFNet prompt tuning
Closest papers by meaning
- Distilling Semantic Priors from SAM to Efficient Image Restoration Models (CVPR 2024 · similarity 0.89)
- UP-Restorer: When Unrolling Meets Prompts for Unified Image Restoration (AAAI 2025 · similarity 0.86)
- Learning Diffusion Texture Priors for Image Restoration (CVPR 2024 · similarity 0.86)
- ReFIR: Grounding Large Restoration Models with Retrieval Augmentation (NeurIPS 2024 · similarity 0.85)
- Navigating Image Restoration with VAR's Distribution Alignment Prior (CVPR 2025 · similarity 0.85)
- URWKV: Unified RWKV Model with Multi-state Perspective for Low-light Image Restoration (CVPR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.