💻 computer science

TDiR: Transformer based Diffusion for Image Restoration Tasks

This paper introduces TDiR, a transformer-based diffusion model that outperforms 18 state-of-the-art techniques across five benchmarks in underwater enhancement, denoising, and deraining, effectively restoring degraded image quality for downstream tasks.

Original authors: Abbas Anwar, Ibrahim Radwan, Ali Arshad Nasir, Mudassir Masood, Saeed Anwar

Published 2026-07-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Abbas Anwar, Ibrahim Radwan, Ali Arshad Nasir, Mudassir Masood, Saeed Anwar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to look at a beautiful painting, but someone has sprayed it with fog, splattered it with mud, or covered it in a thick layer of rain. In the world of computers, this is called "image degradation." When a camera captures a photo in a dark cave, underwater, or during a storm, the resulting picture is often blurry, noisy, or the wrong color. This isn't just an annoyance for photographers; it's a huge problem for robots, self-driving cars, and doctors who need to see clearly to make decisions. For years, computers have tried to fix these pictures using "discriminative" models. Think of these like a strict art teacher who looks at a muddy painting and tries to guess exactly what the original art must have been, following a single, rigid path to the answer. But real-world messiness is rarely that simple. Sometimes, a blurry spot could be fixed in many different ways, and a single rigid guess often misses the mark.

Enter a newer, more creative approach called "diffusion models." If the old way is a strict teacher, a diffusion model is like a sculptor working with clay. Instead of guessing the final shape immediately, the sculptor starts with a lump of random noise (like a chaotic pile of clay dust) and slowly, step-by-step, chips away the noise to reveal the hidden image underneath. This paper introduces a new tool called TDiR (Transformer-based Diffusion for Image Restoration). It combines the step-by-step "sculpting" power of diffusion models with a "Transformer" architecture—a type of computer brain that is really good at understanding how different parts of an image relate to each other, much like how we understand a sentence by looking at how words connect. The authors wanted to see if this combination could fix three very messy types of photos: underwater shots (which are often green and hazy), photos covered in rain, and photos full of static noise.

The researchers built TDiR by taking an existing smart network called PromptIR and giving it a "noise-conditioned" upgrade. They added a special decoder pathway that tells the computer, "Hey, we are at this specific step of the cleaning process, and the noise level is this high." This allows the model to adjust its strategy as it works. They tested this new system on five different standard benchmarks, comparing it against 18 other top-tier techniques. The results were impressive. In the underwater category, TDiR achieved a score of 22.90 dB in PSNR (a measure of how close the restored image is to the original) and 0.8724 in SSIM (a measure of structural similarity), exceeding the closest competitor, MetaUE, by 0.89dB and 0.011 SSIM points. For removing rain, it scored 37.43 dB on the Rain100L dataset, narrowly beating a model that was specifically retrained just for that task (PromptIR*) by 0.40dB.

However, the paper is careful not to call this a perfect, solved problem. The authors found that while TDiR is a "generalist" that can handle noise, rain, and underwater issues all at once, it does not always beat "specialist" models that only do one job. In fact, for denoising at lower noise levels, a model trained solely on that task still performed better. They also noted a specific quirk: because the model processes images in small patches (like a mosaic), it sometimes leaves faint "blocking" artifacts at the edges where the patches meet, which can slightly lower the technical scores even if the image looks great to the human eye. Furthermore, the process is slower and more computationally expensive than older methods because it has to take many steps to "sculpt" the image back to life.

Ultimately, the study suggests that mixing diffusion models with transformers is a powerful way to restore degraded images, often producing results that look more natural and handle complex, multi-layered messes better than older methods. It suggests that for tasks like underwater exploration or autonomous driving, where the "right" answer isn't always obvious, this flexible, step-by-step approach offers a significant upgrade in visual quality, even if it comes with a higher cost in computing power and a few minor technical trade-offs.

Technical Summary: TDiR (Transformer based Diffusion for Image Restoration)

Problem Statement

Image restoration is a critical challenge in computer vision, aiming to recover clean, original images from degraded versions affected by noise, blur, color casts, and light scattering. These degradations significantly impair the utility of images in downstream tasks such as object detection, mapping, and classification. While non-generative methods have shown success in high-level vision tasks, they often struggle with the complex, multimodal distributions found in low-level vision tasks like underwater reconstruction, denoising, and deraining. Specifically, underwater imagery presents unique challenges due to light scattering, absorption, and color shifts that vary with environmental conditions, often requiring models capable of navigating multiple potential degradation states rather than mapping a single degraded input to a single clean output.

Methodology

The paper proposes TDiR, a transformer-based conditional diffusion model designed to address multiple image restoration tasks within a unified framework.

Architecture: The core of TDiR is a Transformer-based U-Net architecture, adapted from the PromptIR backbone. PromptIR utilizes prompt modules to adapt to varying levels of degradation without requiring distinct models for each task.
Diffusion Process: The model operates within a diffusion framework consisting of a forward diffusion process (progressively adding Gaussian noise) and a backward denoising stage.
- Forward Process: Gaussian noise is added to the input image $y_0$ over $T$ iterations until it resembles pure noise.
- Backward Process: A denoising network learns to reverse this process, moving from the Gaussian distribution back to the empirical distribution of the dataset.
Key Architectural Modification: The authors introduce a lightweight modification where a dedicated decoder pathway concatenates the degraded input with an embedding of the diffusion timestep. This allows the network to condition its restoration process based on the specific noise level introduced during the forward pass.
Training Strategy: To address the data-hungry nature of transformers, the authors utilize a pre-trained encoder (frozen) as the backbone. Only the decoder parameters are optimized during training. The model is trained on a combined dataset covering three tasks: underwater enhancement, denoising, and deraining. The objective function minimizes the $L_1$ loss between the predicted noise and the actual noise added to the image.
Inference: The restoration process involves an iterative denoising procedure, requiring multiple forward passes per image.

Key Contributions

Unified Conditional Diffusion Model: The proposal of TDiR, which converts a pretrained PromptIR backbone into a noise-conditioned denoiser, enabling application across diverse restoration tasks (underwater, denoising, deraining) without task-specific model switching.
Lightweight Timestep Conditioning: The introduction of a decoder pathway that concatenates the degraded input with diffusion timestep embeddings, allowing the model to dynamically adapt its restoration strategy based on the noise level.
Comprehensive Evaluation: A rigorous assessment of TDiR across three distinct tasks using five benchmark datasets (UIEB, Test-60, BSD68, Rain100L, etc.) and four evaluation metrics (PSNR, SSIM, UCIQE, UIQM). The study compares TDiR against 18 state-of-the-art techniques.
Qualitative Analysis: An examination of the diffusion-based restoration process compared to purely discriminative networks, highlighting advantages in handling multimodal degradations while acknowledging constraints like residual patch-based artifacts.

Experimental Results

The paper reports that TDiR outperforms current state-of-the-art methods across the evaluated tasks:

Underwater Image Enhancement: On the UIEB dataset, TDiR achieved the highest full-reference scores with 22.90 dB PSNR and 0.8724 SSIM, surpassing the closest competitor (MetaUE) by 0.89 dB and 0.011 SSIM, respectively. On the no-reference Test-60 subset, it achieved a UIQM score of 2.73, significantly outperforming other methods. Qualitative results show superior color fidelity and suppression of green color casts compared to PromptIR and other baselines.
Image Denoising: On the BSD68 dataset, TDiR outperformed all benchmark methods (including specialized single-task models) across noise levels $\sigma=15, 25, 50$ . Notably, at $\sigma=50$ , TDiR surpassed the specialist PromptIR model. While SSIM slightly decreased at high noise levels due to patch artifacts, the model produced perceptually cleaner results with fewer visible grain artifacts in uniform regions.
Image Deraining: On the Rain100L dataset, TDiR achieved a PSNR of 37.43, outperforming the retrained PromptIR* baseline by 0.40 dB and significantly exceeding other methods like AirNet and MPRNet. Qualitative analysis confirmed the complete removal of rain streaks in uniform regions (sky, sand) and better preservation of fine textures compared to competitors.

Significance and Claims

The paper claims that the combination of diffusion models and transformers offers a robust solution for image restoration, particularly in scenarios involving complex, multimodal degradations. The authors emphasize that diffusion models excel at navigating the diverse distributions of degraded images, enabling a more nuanced restoration that reflects real-world variability, which is particularly beneficial for underwater imaging where lighting and clarity vary significantly.

However, the paper maintains a modest tone regarding its limitations. It acknowledges that:

Computational Cost: The iterative nature of diffusion makes TDiR significantly more expensive to infer than single-pass discriminative models, potentially limiting real-time applications.
Artifacts: The patch-wise processing of the transformer-based denoiser can lead to blocking artifacts at stitch boundaries, which may marginally reduce PSNR/SSIM scores relative to perceptual quality.
Generalization Trade-offs: While the unified model is versatile, it consistently underperforms compared to specialized single-task models in specific scenarios (e.g., low-depth underwater scenes with flat backgrounds), suggesting a trade-off between versatility and peak accuracy.
Depth Dependency: The model relies on depth-related cues for effective underwater restoration, which can be absent in certain challenging scenes.

The authors conclude that while TDiR demonstrates significant improvements in image quality across multiple tasks, future work could address its limitations through techniques such as overlapping-patch inference, distillation to fewer diffusion steps, or task-conditioned fine-tuning.