UniRestore: Unified Perceptual and Task-Oriented Image Restoration Model Using Diffusion Prior

2025-01-22

AI Summary

Existing image restoration methods suffer from an inherent trade-off between perceptual quality and downstream task performance: perceptual-oriented restoration (PIR) yields visually pleasing results but degrades task accuracy, while task-oriented restoration (TIR) improves task metrics at the expense of perceptual fidelity. To resolve this dichotomy, we propose the first unified framework integrating diffusion priors with task-adaptive mechanisms. Our approach introduces a Complementary Feature Reconstruction Module (CFRM) and a Task Feature Adapter (TFA), enabling joint optimization of human visual perception and high-level vision tasks (e.g., detection, segmentation) within a single model. Built upon diffusion model priors and a VAE latent space, the framework supports end-to-end co-optimization. Extensive experiments across diverse degradations—including noise, blur, and adverse weather—demonstrate consistent and significant improvements over state-of-the-art methods in both PSNR/SSIM and task-specific metrics. This work breaks the long-standing paradigm separation between PIR and TIR, achieving unprecedented fidelity and strong generalization across vision tasks.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

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.

Problem

• Image Restoration
• Perceptual Quality
• Task-oriented Performance

Innovation

• Diffusion Prior
• Complementary Feature Restoration Module (CFRM)
• Task Feature Adapter (TFA)