Task-Guided Prompting for Unified Remote Sensing Image Restoration

1. Problem Statement

Remote Sensing Image Restoration (RSIR) is critical for recovering high-fidelity imagery from degraded observations to support downstream applications like object detection and land-cover segmentation. However, existing approaches face two major limitations:

• Hyper-Specialization: Most current models are "expert" systems designed for a single degradation type (e.g., only cloud removal or only SAR despeckling) or a single sensor modality (optical vs. SAR). This creates a significant operational bottleneck, as real-world scenarios often involve multiple, co-occurring degradations across diverse sensor types (RGB, Multispectral, SAR, Thermal Infrared).
• Domain Gap: Unified restoration frameworks developed for natural images (e.g., denoising, deblurring) fail to generalize to remote sensing due to fundamental differences in imaging physics (e.g., additive noise in optical vs. multiplicative speckle in SAR) and spectral characteristics.

The core challenge is designing a single, unified framework capable of handling heterogeneous degradations (denoising, cloud removal, shadow removal, deblurring, SAR despeckling) across distinct sensor modalities without requiring task-specific retraining or separate model suites.

2. Methodology: TGPNet

The authors propose TGPNet, a unified framework built on an encoder-decoder backbone (adapted from Restormer) enhanced by a novel Task-Guided Prompting (TGP) strategy.

A. Architecture Overview

• Backbone: A U-shaped Transformer-based architecture with multi-head transposed self-attention for long-range dependency modeling and gated depth-wise convolutions for local refinement.
• Task-Guided Prompting (TGP): Instead of hard-coding physical equations (which conflict across modalities), TGPNet uses data-driven implicit priors. The TGP module is inserted at six key points in the decoder (four stages, one refinement module, and one terminal convolution).

B. Core Components of TGP

The TGP module consists of two synergistic parts:

1. Learnable Task-Specific Embedding (LTSE):
   • Takes a task prompt ($F_{TP}$) as input (a learnable vector representing the specific degradation type).
   • Processes this through a hierarchical convolutional architecture (7x7 and 3x3 kernels with downsampling) to generate a compact, degradation-aware embedding vector ($F_E$).
   • This vector acts as a "latent physical descriptor" encoding the specific characteristics of the target degradation.
2. Hierarchical Feature Modulation (HFM):
   • Uses the embedding $F_E$ to generate stage-wise affine parameters: a scaling factor ($\gamma$) and a shifting offset ($\beta$).
   • These parameters are applied via an affine transformation ($F_{M} = \gamma \cdot F_{D} + \beta$) to the decoder features at each stage.
   • Mechanism: This allows the network to dynamically reconfigure feature statistics to match the specific signal distribution of the current task (e.g., suppressing multiplicative noise for SAR or removing additive noise for optical) while sharing the same backbone weights.

C. Training Strategy

• Joint Training: The model is trained on a unified pool of datasets covering five tasks and four modalities (RGB, MS, SAR, TIR).
• Prompt Retrieval: During each iteration, the specific task prompt corresponding to the degradation type of the input sample is retrieved to modulate the features, enabling simultaneous optimization for all tasks.
• Loss Function: Primary optimization uses L1 loss (Mean Absolute Error) to preserve edge details.

3. Key Contributions

1. First Unified Cross-Domain Framework: TGPNet is the first model to unify restoration across fundamentally distinct sensor modalities (Optical, SAR, and Thermal Infrared) and five diverse degradation types within a single network.
2. Task-Guided Prompting (TGP) Strategy: Introduces a novel mechanism where learnable task embeddings dynamically modulate hierarchical features via affine transformations, enabling a single model to act as multiple "experts" without weight duplication.
3. Comprehensive Benchmark (URSIR): The authors constructed and will release a unified Remote Sensing Image Restoration benchmark covering:
   • Datasets: UCMLUD (Denoising), RICE1/2 & SEN12MS-CR (Declouding), SRD & UAV-TSS (Deshadowing), NRD (SAR Despeckling), HIT-UAV (Deblurring).
   • Scope: Covers RGB, Multispectral, SAR, and Thermal Infrared modalities.

4. Experimental Results

The paper validates TGPNet through extensive experiments on the new benchmark and comparisons with state-of-the-art (SOTA) unified and specialized models.

• Unified Performance: TGPNet achieves State-of-the-Art (SOTA) performance across 2-task, 3-task, 4-task, and 5-task scenarios.
  • It outperforms existing unified frameworks (e.g., PromptIR, AdaIR, MOCE-IR) and specialized models (e.g., CR-former for clouds) in average metrics.
  • Notably, it surpasses specialized cloud removal models by 0.26–0.58 dB in PSNR on the RICE datasets, proving that unification does not sacrifice specialization.
• Cross-Domain Generalization: The model successfully handles the transition from Optical to Microwave (SAR) and Thermal (TIR) domains, maintaining high performance even when adding new modalities to the training set.
• Composite Degradations (Out-of-Distribution):
  • For unseen composite degradations (e.g., Noise + Cloud), TGPNet employs a Sequential Processing strategy. Instead of averaging prompts (which fails), it decomposes the task into ordered subtasks (e.g., Denoise $\to$ Decloud).
  • This approach yields massive gains (e.g., +8 dB over baselines) on composite tasks, demonstrating robust generalization to real-world scenarios.
• Efficiency: TGPNet is highly efficient, with 21.27M parameters and 71.42G FLOPs. It achieves the fastest inference speed (60.84 ms) among high-performing competitors, making it suitable for operational pipelines.
• Ablation Studies:
  • Feature clustering analysis (t-SNE, Silhouette, ARI) confirms that TGP effectively disentangles task-specific features, creating distinct clusters for each degradation type.
  • The linear affine modulation is proven sufficient, offering a better trade-off between performance and complexity compared to non-linear variants.

5. Significance

This work represents a paradigm shift in Remote Sensing Image Restoration:

• Operational Scalability: It eliminates the need to deploy and maintain dozens of specialized models for different sensors and degradation types, offering a practical "all-in-one" solution for satellite and UAV data processing.
• Handling Complexity: By successfully unifying optical and SAR modalities, it addresses the "domain gap" that has previously prevented unified models from being effective in remote sensing.
• Future-Proofing: The ability to handle unseen composite degradations via sequential prompting suggests a path toward truly robust, blind restoration systems for complex, real-world environments.

The authors conclude that while the current model relies on explicit task prompts (non-blind), the framework lays the groundwork for future self-supervised, blind unified restoration systems.