Source-linked AI summary
UniRestore: Unified Perceptual and Task-Oriented Image Restoration Model Using Diffusion Prior
I-Hsiang Chen, Wei-Ting Chen, Yu-Wei Liu, Yuan-Chun Chiang, Sy-Yen Kuo, Ming-Hsuan Yang
TL;DR
PIR can improve visual quality without supporting downstream tasks, while TIR can improve task performance at the expense of visual appeal. UniRestore adapts a diffusion prior with restored encoder features and feature fusion to serve both goals. The paper reports improvements in perceptual and task-oriented restoration and says the approach is extendable.
Problem
PIR may not improve downstream task performance, whereas TIR may produce less visually appealing images, reflecting differing criteria for visual and recognition quality.
Source: §1 · p. 2Method
UniRestore adapts a diffusion prior using CFRM to reconstruct degraded encoder features and TFA to harmonize restored and diffusion features for downstream tasks.
Source: §1 · p. 2Results
UniRestore surpasses existing methods in both visual quality and downstream task performance.
Source: §1 · p. 2Takeaways & Limitations
The paper concludes that UniRestore alleviates trade-offs between PIR and TIR and can be extended to diverse applications.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Image restoration aims to recover content from inputs degraded by various factors, such as adverse weather, blur, and noise. Perceptual Image Restoration (PIR) methods improve visual quality but often do not support downstream tasks effectively. On the other hand, Task-oriented Image Restoration (TIR) methods focus on enhancing image utility for high-level vision tasks, sometimes compromising visual quality. This paper introduces UniRestore, a unified image restoration model that bridges the gap between PIR and TIR by using a diffusion prior. The diffusion prior is designed to generate images that align with human visual quality preferences, but these images are often unsuitable for TIR scenarios. To solve this limitation, UniRestore utilizes encoder features from an autoencoder to adapt the diffusion prior to specific tasks. We propose a Complementary Feature Restoration Module (CFRM) to reconstruct degraded encoder features and a Task Feature Adapter (TFA) module to facilitate adaptive feature fusion in the decoder. This design allows UniRestore to optimize images for both human perception and downstream task requirements, addressing discrepancies between visual quality and functional needs. Integrating these modules also enhances UniRestore's adapability and efficiency across diverse tasks. Extensive expertments demonstrate the superior performance of UniRestore in both PIR and TIR scenarios.
1. Introduction
Perceptual and task-oriented restoration can favor different image qualities, leaving a need for one framework that serves both human perception and downstream tasks. UniRestore adapts a diffusion prior using restored encoder features and reports improvements in both settings.
1. Introduction
- PIR methods can improve visual quality without improving downstream tasks, while TIR methods can improve task performance but produce less visually appealing images.The paper attributes this tension to differences between factors that shape visual quality and recognition quality.Source: §1 · p. 2
- A unified framework is motivated by diverse real-world scenarios and the potential to reduce system redundancy and improve operational efficiency.The paper frames simultaneous perceptual and task performance as the goal.Source: §1 · p. 2
- UniRestore uses a diffusion prior for high-quality image generation and adapts it to meet both perceptual and downstream task needs.The paper notes that diffusion-generated images optimized for human aesthetics may not align with task requirements.Source: §1 · p. 2
- CFRM reconstructs degraded autoencoder encoder features, while TFA harmonizes those restored features with diffusion features in the decoder for downstream tasks.The TFA is described as extendable to accommodate new tasks.Source: §1 · p. 2
- UniRestore surpasses existing methods in both visual quality and downstream task performance.The paper presents it as a unified model for perceptual image restoration and task-oriented image restoration.Source: §1 · p. 2
2. Related Work
Prior work addresses perceptual restoration, task-oriented restoration, and diffusion-based restoration, but the paper identifies using pretrained diffusion models for both PIR and TIR as an open challenge.
2.1. Perceptual Image Restoration
- PIR research ranges from single-degradation methods to unified and holistic approaches, but visual-quality gains do not always improve downstream vision tasks.Examples include MPRNet, NAFNet, SwinIR, Restormer, All-in-One, TransWeather, and PromptIR.Source: §2.1 · p. 2
2.2. Task-oriented Image Restoration
- TIR methods target downstream performance through approaches including end-to-end enhancement and recognition, degradation-adaptive processing, and semantic representation recovery, but may reduce visual appeal.The passage cites URIE [52], DIP [34], and VDR-IR.Source: §2.2 · p. 2
2.3. Diffusion Model
- Diffusion models have been applied to restoration, including super-resolution and degradation restoration, yet using pretrained diffusion models for both PIR and TIR remains open.StableSR, DiffUIR, and DiffBIR are cited as restoration approaches; UniRestore’s overview places CFRMs and TFAs in a pretrained autoencoder.Source: §2.3 · pp. 2–3, Figure 2 · p. 3
3. Proposed Method
UniRestore combines a diffusion prior with restored encoder features to serve perceptual and task-oriented restoration. CFRM reconstructs degraded features, TFA fuses them with decoder features, and a two-stage pipeline trains restoration before task adaptation.
3.1. Architecture of UniRestore
- The degraded input passes through a CFRM-enhanced VAE encoder and controller into the denoising U-Net, then a TFA-enhanced decoder produces task-adapted outputs.The controller uses an SC-Tuner to integrate control signals with the denoising U-Net.Source: §3.1 · p. 3
3.2. Complementary Feature Restoration Module
- CFRM enhances encoder features, models degradation diversity with grouped channel attention, and combines recovered group features with enhanced features through a skip connection.Its output supplies complementary features to the TFA.Source: §3.2 · pp. 3–4
3.3. Task Feature Adapter
- TFA reuses a lightweight learnable prompt per task to update across decoder layers, dynamically fusing CFRM-restored features with diffusion features.This addresses the parameter and scalability burden of separate adapters for many tasks.Source: §3.3 · p. 4
3.4. Training Pipeline
- In Stage 1, PIR data trains CFRM, Controller [74], and SC-Tuner [23]; the CFRM loss restores degraded encoder features to clear states by comparing clear and restored latent features across encoder layers.Source: §3.4 · p. 4
- Stage 2 freezes CFRM, Controller, and SC-Tuner while optimizing TFA with task-specific objectives for perceptual restoration, segmentation, and classification.Afterward, adding a TIR task requires only a new task prompt, its objective, and training data; original-task data and losses are not required.Source: §3.4 · pp. 4–5
4. Implementation Details
Experiments evaluate perceptual restoration, image classification, and semantic segmentation using synthesized degradations. Training combines established restoration datasets with ImageNet and Cityscapes, while evaluation includes both familiar and unseen benchmarks.
4.1. Training Dataset
- Training uses DIV2K, Flickr2K, and OST for perceptual restoration, 80,000 ImageNet training images for classification, and Cityscapes training data for segmentation.Fifteen degradation types, including blur, noise, and adverse weather, are synthesized following [16].Source: §4.1 · p. 5
4.2. Evaluation Dataset
- Perceptual restoration is evaluated on DIV2K and synthetic-degradation benchmarks for denoising, deraining, dehazing, desnowing, and deblurring.The benchmarks include Noise datasets, Rain100L, RESIDE, UHDSnow, and GoPro.Source: §4.2 · p. 5
- Classification evaluation uses ImageNet test images and unseen CUB data, while segmentation evaluation uses Cityscapes, FoggyCityscapes, and unseen ACDC.FoggyCityscapes results average the Fog1, Fog2, and Fog3 subsets.Source: §4.2 · p. 5
4.3. Evaluation Protocol
- Evaluation measures perceptual restoration with PSNR and SSIM, classification with accuracy, and segmentation with mean Intersection over Union.Source: §4.3 · p. 5
5. Experiments
Experiments evaluate UniRestore on perceptual restoration, classification, and segmentation, with comparisons across seen and unseen datasets. Results favor UniRestore across these tasks, while ablations examine its modules, task-specific adapters, and extension to object detection.
5.1. Baselines
- Baselines include task-oriented, perceptual, and diffusion-based methods, evaluated both as originally trained and after PIR training followed by downstream-task fine-tuning.The fine-tuned comparison uses the suffix “*”.Source: §5.1 · p. 5
5.2. Perceptual Image Restoration
- UniRestore achieves the best overall performance on the seen PIR dataset and generalizes to unseen datasets, especially high-resolution UHDSnow and dynamic-scene GoPro.Qualitative comparisons also show more reconstructed detail and better visual quality on unseen datasets.Source: §5.2 · pp. 5–6, Table 1 · p. 6, Figure 4 · p. 6
5.3. Task-oriented Image Restoration
- On unseen CUB images, UniRestore enhances classification accuracy by 20.01% for ResNet-50 and 15.96% for ViT-B.Its classification accuracy is also comparable to high-quality ground-truth inputs.Source: §5.3 · p. 6, Table 2 · p. 7
- UniRestore achieves semantic-segmentation performance on seen and unseen datasets, and its restored images support more accurate object boundaries.The authors attribute the fine-grained detail to TFA integrating diffusion features with restored features.Source: Table 3 · p. 7, §5.3 · p. 7, Figure 6 · p. 7
5.4. Ablation Study
- Ablations show that CFRM and TFA improve performance across PIR and TIR scenarios.The evaluation covers DIV2K for PIR, ImageNet for classification, and Cityscapes for segmentation.Source: §5.4 · pp. 7–8, Table 4 · p. 8
- UniRestore with task-specific prompts performs better in TIR than the single-prompt variant and comparably to Multi-TFA in PIR.Multi-TFA outperforms Multi-Adapter, while UniRestore-SP tunes fewer parameters but performs worse than UniRestore.Source: Table 5 · p. 8, §5.4 · p. 8
- For object detection on RTTS, UniRestore adds a task by optimizing a new prompt with detection data and loss, rather than retraining the full model.The evaluation uses a RetinaNet backbone and COCO training images with synthetic degradation.Source: Table 6 · p. 8, §5.4 · p. 8
6. Conclusion
UniRestore adapts diffusion features for perceptual and task-oriented restoration by restoring encoder features and combining them with diffusion features. Experiments validate its effectiveness and extendability, and the approach alleviates trade-offs associated with existing methods.
6. Conclusion
- UniRestore addresses perceptual and task-oriented image restoration simultaneously by adapting diffusion features for diverse applications.The framework restores encoder features with CFRM and dynamically combines them with diffusion features through TFA.Source: §6 · p. 8
- CFRM restores encoder features, while TFA dynamically combines them with diffusion features for downstream tasks.Source: §6 · p. 8
- Experiments validate UniRestore’s effectiveness and extendability, and show it can alleviate trade-offs associated with existing PIR and TIR methods.Source: §6 · p. 8