Boosting All-in-One UAV Image Restoration via Degradation-Aware Frequency Prompt State Space Model

Highlights

What are the main findings?

• A degradation-aware Frequency Prompt State Space Model is proposed, integrating Prompt-Guided Mamba Blocks and Adaptive Frequency Prompt Blocks to jointly achieve spatial–frequency collaborative restoration for UAV images.
• The proposed method consistently outperforms existing all-in-one restoration approaches on multiple benchmarks, demonstrating strong robustness and generalization capability across blur, noise, haze, raindrop, and low-light degradations.

What are the implications of the main findings?

• Explicitly modeling frequency-domain degradation characteristics and embedding them as prompts effectively mitigates cross-degradation interference and enhances detail preservation and structural fidelity in UAV image restoration.
• The proposed unified framework provides a practical and efficient solution for UAV perception systems, improving visual reliability under complex and coupled degradations and supporting downstream UAV vision tasks.

Abstract

Unmanned Aerial Vehicle (UAV)-captured images are easily affected by various degradations such as motion blur, noise, low illumination, haze, and raindrops in complex environments, and these degradations exhibit significant differences in the frequency domain. Existing all-in-one models typically operate in the spatial domain, making it difficult to effectively distinguish different degradation types, which leads to degradation interference and suboptimal restoration quality. To address these challenges, this paper proposes a degradation-aware prompt state space model for unified UAV image restoration. Specifically, we design a Prompt-Guided Mamba Block (PGMB) that injects dynamic degradation prompts into state space modeling, achieving differentiated global structure modeling. Meanwhile, we introduce an Adaptive Frequency Prompt Block (AFPB) that explicitly perceives frequency characteristics of diverse degradations and dynamically guides the reconstruction process via frequency-domain prompts, enabling collaborative restoration in both spatial and frequency domains. Extensive experiments demonstrate that the proposed method outperforms existing all-in-one approaches across multiple typical degradation tasks, significantly enhancing visual quality and robustness of UAV imagery.

1. Introduction

Unmanned Aerial Vehicles (UAV) offer strong mobility, flexible deployment, and diverse imaging perspectives, and they have been widely applied in security inspection [1,2], remote sensing mapping [3,4,5], agricultural monitoring [6,7], and disaster response [8]. High-quality aerial images are essential for ensuring the performance of downstream visual [9] tasks such as object detection, semantic segmentation, and target tracking. However, due to unstable flight attitudes, complex illumination conditions, and various environmental disturbances, UAV-captured images often suffer from different levels of visual degradation, which severely compromises the reliability and practicality of UAV-based vision systems.

In real-world applications, UAV-captured images are often affected by multiple complex degradations such as motion blur, sensor noise, low-light conditions, haze, and raindrop occlusion. These degradations have different causes and destructive behaviors: motion blur reduces edge sharpness and structural clarity, noise corrupts high-frequency texture details, low illumination decreases the overall signal-to-noise ratio, and haze and raindrops lead to contrast attenuation and local occlusions. The coexistence of these factors poses more complex and stringent challenges for UAV image restoration.

With the rapid development of image processing technology, extensive studies have been conducted on image restoration. Traditional methods are typically designed for a single degradation type, such as denoising [10,11,12], deraining [13,14,15], deblurring [16,17], dehazing [18,19,20], or low-light enhancement [21,22]. In recent years, deep learning has significantly improved the performance of single-task restoration. However, these approaches generally lack degradation-awareness and cannot dynamically adjust their inference strategies when dealing with multiple degradations. As a result, their generalization capability remains limited, making them insufficient for real-world UAV applications where diverse degradations frequently coexist.

To improve model generalization, all-in-one image restoration methods [23,24,25,26,27,28] have gradually become a research focus. Convolution-based [29] models are limited by their restricted receptive fields and struggle to capture long-range contextual dependencies. Transformer-based [30,31] approaches can model global features but suffer from the quadratic computational complexity of self-attention. Recent methods, such as PromptIR [28], achieve multi-degradation restoration through prompt modulation; however, their prompts are highly implicit and insufficient to suppress interference among different degradations. InstructIR [32] introduces textual prompts but still fails to handle the performance imbalance caused by heterogeneous degradation distributions. AMIR [33] enhances adaptiveness via dynamic expert routing, yet it ignores the inherent frequency characteristics of degradations, leading to insufficient detail recovery.

Moreover, although existing all-in-one methods are capable of handling multiple degradations within a single framework, they primarily rely on spatial-domain feature interaction and fail to explicitly leverage the distinct frequency-domain distributions of different degradations. We observe that degradations such as blur, noise, low-light conditions, haze, and raindrop occlusion exhibit significant differences in their frequency characteristics. As shown in Figure 1, motion blur causes clear anisotropic streak patterns, noise introduces uniformly distributed high-frequency components forming a cross-shaped spectrum, low-light and haze lead to over-concentrated low-frequency energy and reduced global contrast, while raindrops result in approximately isotropic attenuation in the mid-to-high frequency bands. Ignoring these distinctions leads to feature interference among different degradations, compromising the stability and overall performance of unified restoration frameworks. Therefore, explicitly analyzing degradation-specific frequency cues and transforming them into learnable prompts to guide targeted reconstruction is crucial for mitigating cross-degradation interference and improving all-in-one UAV image restoration.

In recent years, state space models (SSMs) [34,35,36] have shown great potential in the vision domain due to their efficient long-range dependency modeling and linear computational complexity. Although some all-in-one restoration studies have adopted SSM-based architectures, they typically rely on fixed-direction state scanning and static modeling parameters, making it difficult to flexibly adapt to heterogeneous degradation distributions. Moreover, such static state propagation often leads to ineffective state accumulation, resulting in insufficient detail recovery under diverse degradation conditions.

To address the above issues, we propose a degradation-aware state space model guided by frequency prompts for unified UAV image restoration. Specifically, we integrate state space modeling with a prompt-guided strategy and introduce a Prompt-Guided Mamba Block (PGMB), where learnable degradation prompts adaptively regulate the state update process to achieve differentiated modeling of multiple degradations without requiring explicit degradation labels or prior knowledge. In addition, we design an Adaptive Frequency Prompt Block (AFPB) that adaptively identifies degradation-specific frequency characteristics and generates frequency prompts to accurately distinguish diverse degradation patterns. This enables collaborative optimization between global structural consistency and high-frequency detail enhancement. Extensive experiments demonstrate that the proposed method significantly outperforms existing state-of-the-art all-in-one approaches under complex and coupled degradation conditions such as motion blur, noise, low-light, haze, and raindrop occlusion, effectively improving the visual quality of UAV images and the usability of downstream perception tasks.

The main contributions of this work can be summarized as follows:

• We propose a degradation-aware state space model guided by frequency prompts for unified UAV image restoration. By introducing learnable degradation prompts into state space modeling, our method achieves differentiated global representation without requiring explicit degradation labels or task priors.
• We design an Adaptive Frequency Prompt Block that exploits frequency-domain degradation priors to dynamically perceive degradation characteristics and guide selective restoration of texture and structural details.
• Extensive experiments on multiple public benchmarks demonstrate that our method consistently outperforms existing state-of-the-art all-in-one restoration approaches, validating its effectiveness and robustness in complex UAV imaging conditions.

In the following sections, we first review the relevant research on multi-degradation image restoration, including recent advances in single-degradation restoration, multi-degradation restoration, and all-in-one image restoration approaches. Section 3 then presents the proposed degradation-aware state space model guided by frequency prompts in detail, including the overall network architecture and each of its core components. Section 4 provides comprehensive experimental evaluations, covering datasets and implementation details, quantitative and qualitative comparisons against state-of-the-art methods, as well as ablation studies. Finally, Section 5 concludes this work and discusses potential directions for future research.

2. Related Work

2.1. Image Restoration for Single Degradation

For single-degradation image restoration tasks, extensive research has been conducted, with remarkable progress in typical scenarios such as deblurring, dehazing, deraining, denoising, and low-light enhancement. For motion deblurring, Lee et al. [17] proposed a UAV image deblurring network that enhances structural detail visibility through deep feature modeling, while Li et al. [37] employed a Transformer-based architecture tailored to UAV imaging characteristics to effectively restore sharp images in complex aerial scenes. In the field of dehazing, the HazyDet dataset [9] provides a realistic benchmark with depth cues for UAV detection under hazy weather, and ClearSight [2] demonstrated efficient dehazing through deep reconstruction. Regarding rain-induced degradation, DRSformer [38] introduces a sparse Transformer framework that improves rain-streak separation, and MSDT [39] leverages multi-scale feature modeling to enhance deraining performance. Chen et al. [40] established a large-scale 4K deraining benchmark and designed an efficient deraining network achieving low complexity and high-quality restoration in ultra-high-resolution settings. Zhou et al. [41] further developed a multi-patch raindrop removal network for UAV imagery, where frequency-aware attention effectively suppresses raindrop artifacts. For image denoising, LAN [13] adaptively adjusts its denoising strategy by learning noise distributions, achieving robust performance in complex environments, while IDR introduces an iterative self-supervised data refinement framework that approaches supervised performance without clean ground truth. Low-light enhancement has also advanced rapidly. RNet [42] enhances UAV night-view images through multi-feature fusion, RetinexFormer [43] integrates Retinex decomposition with Transformer-based modeling for illumination–detail co-optimization, and Wave-Mamba [23] combines wavelet transform with state space modeling to yield excellent performance in ultra-high-resolution low-light settings.

Although these single-task methods achieve superior results for specific degradations, they often struggle to generalize when facing multiple co-existing degradations, leading to performance degradation in real-world UAV scenarios.

2.2. Image Restoration for Multiple Degradations

In the field of multi-degradation image restoration, IPT [44] introduced large-scale pre-training into low-level vision and proposed a multi-head multi-tail Transformer framework that supports unified pre-training across multiple restoration tasks, enabling cross-task feature sharing and transfer. However, its computational cost is extremely high, and its generalization ability remains limited under real-world degradations. Restormer [45] designed an efficient Transformer architecture for high-resolution restoration and achieved strong performance across multiple degradation tasks, making it a widely used backbone network. MambaIR [35] incorporated selective state space modeling to capture long-range dependencies with linear complexity, improving both efficiency and accuracy, while MambaIRv2 [36] further enhanced performance through attentive state space modeling. IR-SDE [46] proposed a mean-reversion stochastic differential equation framework for universal image restoration, which requires no task-specific prior knowledge and can handle different tasks simply by changing the dataset. DGUNet [47] introduced an interpretable deep generalized unfolding framework that improves the performance of various restoration tasks. NAFNet [48] simplified traditional residual blocks by removing nonlinear activations, thereby reducing computational and inference overhead while maintaining or even surpassing the performance of activation-based models. Uformer [49] adopted a universal U-shaped Transformer that integrates local window self-attention with multi-scale representation learning to achieve enhanced restoration ability under diverse degradation settings.

Although these methods demonstrate promising performance across multiple degradation tasks, they typically require separate models for each degradation type, resulting in increased parameter and storage costs, and they still struggle to handle scenarios where multiple degradations coexist, as commonly encountered in UAV imagery.

2.3. All-in-One Image Restoration

All-in-one image restoration aims to handle multiple degradation types within a single model, without training separate networks for each specific task, thus providing stronger generalization capability. As a result, a variety of all-in-one restoration approaches have emerged in recent years. AirNet [24] constructs a degradation-aware network via contrastive learning and introduces degradation guidance to achieve unified modeling across multiple tasks. TransWeather [50] incorporates learnable weather-type query tokens into a single Transformer framework, enabling simultaneous handling of various adverse weather conditions. PromptIR [28] designs learnable prompt modules based on Restormer, where implicit prompts adaptively modulate feature distributions for blind multi-degradation restoration. DiffUIR [25] leverages diffusion models to map different restoration tasks into a shared latent space, achieving a balance in performance across multiple degradations. AutoDIR [26] combines task-agnostic blind image quality assessment with structure-constrained latent diffusion to automatically identify unknown degradations and generate textual prompts for unified restoration. InstructIR [32] uses natural language instructions as conditions to address denoising, deblurring, deraining, dehazing, and low-light enhancement within a single model. Li et al. [27] introduced a UAV-oriented all-in-one dataset, UAVIR-5D, and proposed a prompt-guided wavelet state space model capable of efficiently handling diverse degradations in a unified framework.

Although these all-in-one methods can process multiple degradations under a single architecture and support blind degradation recognition to some extent, they mainly rely on spatial-domain modeling and implicit prompting. Specifically, approaches like PromptIR and InstructIR primarily modulate features in the spatial domain, which limits their ability to effectively decouple mixed degradations that are entangled spatially. Consequently, they insufficiently exploit the explicit frequency-domain differences among degradations. Moreover, most existing approaches are developed for general natural scenes rather than UAV imagery, where degradation is more diverse and severe due to dynamic aerial conditions. These limitations motivate us to explore a UAV-specific all-in-one framework. To clarify the position of our work within the field, we provide a systematic comparison of prompt mechanisms in representative methods see

Table 1. Comparison of prompt mechanisms in representative all-in-one image restoration methods. To clarify the distinction, PG-WSSM relies on fixed wavelet bases for local spatial-frequency analysis, whereas our method leverages the Fourier transform for adaptive global spectral filtering.

Method	Prompt Domain	Mechanism	Key Characteristic
PromptIR [28]	Spatial	Implicit Modulation	Local pixel-wise correction
InstructIR [32]	Textual	Instruction-Conditional	Human-instruction-driven
PG-WSSM [27]	Wavelet Space	Fixed Basis Decomposition	Local Spatial-Frequency analysis
Ours	Fourier Space	Adaptive Spectral Filtering	Global Frequency disentanglement

.

It is worth noting that the recently proposed PG-WSSM [27] also incorporates state space models for UAV restoration. However, a fundamental difference lies in the prompt domain. PG-WSSM employs wavelet transforms, which rely on fixed basis functions to decompose features into spatial-frequency sub-bands. While effective for local analysis, this fixed decomposition limits the flexibility to capture holistic degradation patterns. In contrast, our method leverages the Fourier transform to access the global Fourier space. Unlike spatial, textual, or wavelet-based prompts, our frequency prompts are dynamically learned to adaptively modulate components across the full Fourier spectrum. As shown in Table 1, this adaptive global strategy allows for precise degradation disentanglement, distinguishing our approach from existing spatial or wavelet-based baselines.

3. Materials and Methods

In this section, we first present the overall architecture of our proposed network, followed by detailed descriptions of its major components in the subsequent subsections.

3.1. Overall Architecture

The proposed Degradation-Aware Frequency Prompt State Space Model aims to achieve unified and efficient restoration of multi-degradation UAV imagery. The overall architecture is illustrated in Figure 2. The network adopts a multi-stage U-shaped encoder–decoder structure, consisting of three main components: shallow feature extraction, state space–based deep feature modeling, and high-quality image reconstruction. Specifically, given a degraded input image $I_D\in {\mathbb{R}}^{H\times W\times 3}$, we first employ a $3\times 3$ convolution to extract basic structural and texture information, obtaining the initial feature representation $F_0\in {\mathbb{R}}^{H\times W\times C}$, where H, W denote the spatial resolution and C is the channel dimension. Then, $F_0$ is progressively downsampled and fed into a four-stage Prompt-Guided Mamba Block encoder, which models long-range dependencies through state space modeling while enhancing feature expressiveness. The channel capacity is increased with decreasing spatial resolution, yielding deep feature maps denoted as $F_s\in {\mathbb{R}}^{{\displaystyle \frac{H}{8}}\times {\displaystyle \frac{W}{8}}\times 8C}$, which effectively capture global contextual information and mitigate degradation interference. In the decoding path, features are progressively upsampled and fused with encoder features at the corresponding scales via skip connections to recover high-resolution feature representations $F_r\in {\mathbb{R}}^{H\times W\times 2C}$. Finally, the reconstructed feature $F_r$ is processed by a $3\times 3$ convolution and combined with the input image through residual addition to generate the restored output $I_r$. Moreover, to further reinforce frequency-aware structural restoration, the proposed Adaptive Frequency Prompt Block (AFPB) is inserted into the skip connections between encoder and decoder, enabling frequency-domain guidance for high-frequency details and contrast enhancement. To effectively supervise the parameter learning process and enhance the quality of image restoration, we employ a reconstruction loss that constrains the difference between the restored output and the ground-truth image. The overall loss function $\mathcal{L}$ is defined as follows:

$$\mathcal{L}=\parallel I_r-I_{gt}{\parallel }_1$$

(1)

where $I_{gt}$ denotes the ground truth image, and ${\parallel ·\parallel }_1$ represents the $\mathcal{L}1$ norm.

3.2. Prompt-Guided Mamba Block

The Prompt-Guided Mamba Block (PGMB) is the core component of our network for modeling long-range dependencies and achieving adaptive degradation-aware restoration. The key idea is to inject dynamic degradation prompts into the visual state space modeling process to enhance the capability of identifying and restoring different degradation patterns. The overall procedure consists of two stages. Specifically, given the input feature $F_{in}$, we first apply Layer Normalization to obtain the normalized feature representation, which is then fed into the Vision State Space Model (VSSM) to capture long-range contextual dependencies. Subsequently, a residual connection is applied between the VSSM output and the original input feature to form the output $F_1$ of the first stage:

$$F_1=\mathrm{VSSM}\left(\mathrm{LN}\left(F_{in}\right)\right)+F_{in}$$

(2)

where LN denotes the Layer Normalization operation.

Next, the output feature $F_1$ from the first stage is fed into the second stage, where it is sequentially processed by LayerNorm and a Feed-Forward Network (FFN). The resulting features are further combined with $F_1$ through a residual connection to obtain the final output representation $F_{out}$. This process can be formulated as:

$$F_{out}=\mathrm{FFN}\left(\mathrm{LN}\left(F_1\right)\right)+F_1$$

(3)

Through the two-stage processing, the Prompt-Guided Mamba Block achieves both stable training and a balance between global structure modeling and local detail enhancement. Unlike conventional modules that rely solely on convolution or self-attention, and inspired by the design philosophy of Mamba-based [36] architectures, our PGMB incorporates degradation-aware prompts into the state space framework, enabling stronger adaptability and generalization to diverse UAV degradation conditions. Next, we provide a detailed introduction to the Vision State Space Model, which serves as the core component of dynamic context aggregation within the PGMB.

To enable the state space model to dynamically adapt to and accurately restore various UAV degradations, we incorporate a Dynamic Prompt Modulation mechanism into the Prompt-Guided Mamba Block (PGMB). This mechanism enhances feature discrimination and selective representation across different degradation types. Specifically, given an input feature map $F_{in}\in {\mathbb{R}}^{H\times W\times C}$, positional encoding is first applied to preserve spatial structural cues. Then, we adopt a semantics-guided neighborhood unfolding (SGN-unfold) strategy to reorganize the 2D feature map into a 1D sequence along the channel dimension, represented as $X\in {\mathbb{R}}^{L\times C}$, where $L=H\times W$.

Unlike conventional row or column unfolding strategies which organize tokens strictly based on spatial raster scanning, SGN-unfold reorganizes tokens according to their learned semantic attributes. Specifically, semantic similarity is measured by the affinity of pixels to the learnable degradation prompts. We generate a routing map via a linear projection of input features, assigning each pixel to a specific latent prompt category. Pixels that share the same prompt assignment are grouped together in the unfolded sequence. This transformation allows the subsequent state space modeling to efficiently exploit long-range dependencies among semantically related regions while maintaining computational efficiency. Based on the unfolded sequence representation, we introduce a learnable Prompt Pool $\mathcal{P}\in {\mathbb{R}}^{N\times d}$ into the Vision State Space Modeling stage, where N denotes the number of prompt components and d corresponds to the latent dimension consistent with the state representation.

The prompt pool acts as a trainable degradation knowledge base, whose internal embeddings are optimized during training and gradually encode the latent priors of diverse UAV degradations (e.g., blur, noise, low-light, haze, and raindrops). To enhance representation sharing across layers while maintaining the expressiveness of degradation-specific differences, we decouple the prompt pool into two parts:

$$\mathcal{P}=\{{\mathcal{P}}^{shared},{\mathcal{P}}^{specific}\}$$

(4)

where ${\mathcal{P}}^{shared}$ stores global degradation-agnostic priors shared across different PGMB layers, ensuring consistent cross-layer guidance. ${\mathcal{P}}^{specific}$ contains local degradation-dependent patterns, enabling fine-grained adaptation to degradation variations at different processing stages. Together, these prompts provide hierarchical conditional cues to guide state transition refinement, improving degradation-aware feature selection and long-range dependency restoration.

To achieve pixel-level degradation-aware prompt selection, we project the unfolded sequence features into a prompt attention space using a lightweight linear transformation. Different from standard attention mechanisms that use Softmax to produce a soft mixture of features, we aim for a sharper selection to distinctively identify specific degradation patterns. Specifically, the projection results are then activated by a LogSoftmax function to ensure numerical stability, optionally coupled with a Gumbel sampling operation to facilitate differentiable sharper selection, to obtain the prompt attention weights:

$$W=\mathrm{Proj}\left(X\right)\in {\mathbb{R}}^{L\times N}$$

(5)

$$\alpha =\mathrm{LogSoftmax}\left(W\right)\in {\mathbb{R}}^{L\times N}$$

(6)

where $\mathrm{Proj}(·)$ represents the projection operator, and $\alpha$ represents the log-probability of each token’s preference toward different prompt components.

Based on these learned attention weights, the degradation-aware prompt embedding is dynamically aggregated from the prompt pool, enabling token-wise prompt variation:

$$\mathbf{P}=\alpha ·\mathcal{P}\in {\mathbb{R}}^{L\times d}$$

(7)

where L is the number of prompt tokens in the prompt pool and d denotes the feature dimensionality of each token.

Thus, different spatial regions within the same UAV image are allowed to select distinct prompt vectors according to their underlying degradation properties, forming spatially heterogeneous restoration constraints that adaptively modulate the reconstruction behavior.

Finally, the dynamically generated prompt P is injected into the state space output transformation, forming a prompt-guided state transition formulation that tightly integrates with VSSM:

$$h_t=\overline{\mathbf{A}}{h}_{t-1}+\overline{\mathbf{B}}{x}_t,y_t=(\mathbf{C}+\mathbf{P})h_t+\mathbf{D}{x}_t$$

(8)

where $\overline{\mathbf{A}}$, $\overline{\mathbf{B}}$, $\mathbf{C}$, and $\mathbf{D}$ denote learnable state transition and projection matrices.

By modulating the output projection matrix, the degradation-aware prompt directly influences the hidden state decoding, allowing the model to adaptively adjust feature propagation with respect to different degradation types—without relying on explicit degradation labels or predefined classifiers. The updated sequence representation is then projected back into the 2D feature space through SGN-fold, serving as the output of the first stage in the Prompt-Guided Mamba Block. Owing to this design, the proposed dynamic prompt modulation enables single-direction scanning while preserving robust degradation adaptivity, thereby avoiding the extra computational cost and redundancy that exist in multi-direction scanning Mamba variants.

3.3. Adaptive Frequency Prompt Block

UAV images often suffer from multiple coexisting degradations such as rain streaks, Gaussian noise, haze, and low-light conditions. Existing spatial-domain restoration methods struggle to effectively disentangle such hybrid degradations due to the entanglement of structural and texture corruptions. However, we observe that these degradations exhibit distinct spectral behaviors in the frequency domain: rain streaks and noise mainly introduce sharp fluctuations in the high-frequency components, whereas haze and low-light conditions predominantly weaken low-frequency responses by reducing illumination and global contrast.

Motivated by this observation, we propose an Adaptive Frequency Prompt Block (AFPB) to explicitly capture degradation discrepancies in the frequency domain. The module learns adaptive degradation-aware prompts that selectively suppress specific corrupted frequency patterns while preserving valid image structures. Specifically, given an input feature map $F_{in}\in {\mathbb{R}}^{H\times W\times C}$, we first apply the Fast Fourier Transform (FFT) [51] to convert spatial-domain representations into frequency-domain spectra, enabling global spectral modeling:

$${\mathcal{F}}_{freq}=\mathrm{FFT}\left(F_{in}\right)$$

(9)

Next, we divide the frequency-domain representations into low-frequency and high-frequency components to separately characterize global structural information and local texture variations. To make the restoration strategy degradation-adaptive, we establish parallel prompt generation paths for both frequency branches. For each frequency component (either low-frequency or high-frequency), we first apply global average pooling (GAP) followed by a linear projection and Softmax activation to derive a normalized frequency-wise prompt weight W. Specifically, the Softmax function is employed to ensure that the generated attention weights form a valid probability distribution (i.e., $\sum W_j=1$). This design introduces a competition mechanism among the prompt components, encouraging the model to selectively emphasize the most relevant frequency patterns for the current input while suppressing irrelevant ones. The calculation is formulated as:

$$W=\mathrm{Softmax}\left(\mathrm{Linear}\left(\mathrm{GAP}\left({\mathcal{F}}_{freq}\right)\right)\right)$$

(10)

Based on the learned frequency weights, we linearly combine a set of learnable frequency prompt components $P_{comp}$ to produce degradation-aware frequency prompts $P_{freq}$:

$$P_{freq}=\sum _{j=1}^N{W}_j·P_{comp}^j,P_{freq}=\{P_{low},P_{high}\}$$

(11)

where $P_{comp}^j$ denotes the j-th prompt component and $W_j$ corresponds to its dynamically predicted importance. $P_{freq}$ includes both high-frequency prompt $P_{high}$ and low-frequency prompt $P_{low}$ embeddings.

This weighted aggregation adaptively adjusts prompt embeddings according to input-dependent degradation characteristics, enabling accurate degradation pattern encoding in the frequency domain. Finally, the low-frequency prompt $P_{low}$ and high-frequency prompt $P_{high}$ are concatenated and injected into the frequency features to simultaneously strengthen valid structures and suppress corruption artifacts. After the frequency-domain modulation, we apply the inverse Fast Fourier Transform (IFFT) to project the enhanced frequency representation back into the spatial domain:

$$F_{spatial}=\mathrm{IFFT}\left(Concat(P_{low},P_{high})\right)$$

(12)

where $Concat(·)$ denotes the concatenation operation along the channel dimension.

Although frequency-domain operations excel at handling global degradations, they may introduce slight artifacts or inconsistencies in local textures. To alleviate this issue and further fuse the features, we apply a $3\times 3$ convolution after the IFFT operation for local spatial refinement. The refined feature maps are then concatenated with the original spatial features to preserve both global structural cues and local details. Finally, the fused representation is fed into the Prompt-Guided Mamba Block (PGMB), where the degradation-aware prompts dynamically regulate feature propagation and selection, ensuring adaptive restoration according to the input degradation type. A concluding $1\times 1$ convolution is employed to adjust the channel dimension, yielding the final output feature $F_{out}$:

$$F_{out}=Con{v}_{1\times 1}\left(PGMB(Con{v}_{3\times 3}\left(F_{spatial}\right)+F_{in})\right)$$

(13)

where $Con{v}_{1\times 1}$ and $Con{v}_{3\times 3}$ denote the convolution operations with kernel sizes of $1\times 1$ and $3\times 3$, respectively. Through this collaborative spatial–frequency design, the AFPB significantly enhances model robustness against complex and mixed degradations commonly observed in UAV imagery.

4. Result

In this section, we present the datasets, evaluation metrics, and implementation details used in our experiments, followed by comprehensive quantitative and qualitative comparisons. In addition, we conduct ablation studies to verify the effectiveness of the proposed components.

4.1. Datasets and Evaluation Metrics

Datasets. In UAV image restoration experiments, we adopt the UAVIR-5D [27] dataset to verify the performance of our method. UAVIR-5D is a multi-degradation benchmark tailored for UAV viewpoints, containing 4520 paired training images and 500 paired test images, with five typical degradation types including blur, noise, haze, low-light, and raindrops. UAVIR-5D covers various common UAV aerial scenarios such as urban roads, building complexes, and campus environments, providing a unified evaluation benchmark for multi-degradation UAV image restoration.

To further evaluate the generalization capability of our model, we also conduct experiments on the MDRS-Landsat [52] remote sensing dataset proposed in Ada4DIR. MDRS-Landsat is constructed using 5500 high-quality Landsat-8 images from the RSHaze dataset, and synthetic degradations including blur, noise, haze, and low-light are systematically applied. This dataset contains 5130 training samples, 100 validation samples, and 270 test samples, covering diverse large-scale land surfaces and man-made structures, providing a more realistic benchmark for cross-domain generalization in multi-degradation restoration.

Evaluation Metrics. For evaluation, consistent with most existing image restoration works, we report Peak Signal-to-Noise Ratio (PSNR) [53] and Structural Similarity Index Measure (SSIM) [54] for each degradation task as well as overall performance. PSNR measures pixel-wise reconstruction fidelity between restored and reference images, where a higher value indicates lower reconstruction error. SSIM assesses perceptual similarity in terms of luminance, contrast, and structure, focusing more on structural information and human visual consistency.

4.2. Implementation Details

For implementation, our model is developed based on the PyTorch 2.0.1 framework and trained using an environment configured with Anaconda3, and CUDA 11.8. All experiments are carried out on a single NVIDIA RTX 4090D GPU. To ensure reproducibility, we specify the key architectural parameters: the initial channel dimension C is set to 48, and the encoder follows a four-stage hierarchical design with channel dimensions of $\{48,96,192,384\}$. The prompt pool size is set to $N=128$, with components initialized using a standard Gaussian distribution with a standard deviation of $\sigma =0.02$. In the AFPB, we employ a fixed frequency separation strategy, where the low-frequency component is defined as the central rectangular region occupying 50% of the spectrum area, while the remaining area is treated as high-frequency. Additionally, the number of Prompt Components is set to 5. We adopt the Adam [55] optimizer with ${\beta }_1=0.9$ and ${\beta }_2=0.999$, setting an initial learning rate of $3\times {10}^{-4}$ and employ a cosine annealing strategy to gradually decay the learning rate to $1\times {10}^{-6}$ for more stable convergence. The batch size is set to 4, and input images are randomly cropped into $128\times 128$ patches during training. The total number of training iterations is fixed to 300 K. To improve the robustness and generalization of our model, standard data augmentation strategies such as random horizontal and vertical flips are applied, while strictly maintaining alignment between degraded images and their corresponding ground truth. During inference, the best checkpoint is selected to generate the final results.

4.3. Comparisons with the State of the Art

In this section, we evaluate the overall performance of different models across various restoration tasks on the UAVIR-5D dataset, including deblurring, denoising, dehazing, de-raindropping, and low-light enhancement. Both quantitative and qualitative results are provided to comprehensively assess the effectiveness of each method. Furthermore, to verify the generalization capability of our approach, we additionally conduct experiments on the publicly available MDRS-Landsat dataset for multi-degradation remote sensing image restoration.

Evaluations on UAVIR-5D datasets for all-in-one UAV image restoration. To validate the effectiveness of the proposed method in handling multiple UAV degradation tasks, we conduct comprehensive comparisons on the UAVIR-5D dataset with a series of representative all-in-one image restoration approaches, including AirNet [24], TransWeather [50], DGUNet [47], IDR [11], IR-SDE [46], PromptIR [28], MambaIR [35], AutoDIR [26], DiffUIR [25], InstructIR [32], PG-WSSM [27], and FoundIR [56].

As shown in

Table 2. Quantitative comparisons (PSNR/SSIM) on the UAVIR-5D dataset under five degradation types for different restoration methods. The best results are highlighted in bold.

Methods	Venue	Blur		Noise		Haze		Raindrop		Lowlight		Average
		PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
AirNet	CVPR 22	21.54	0.6312	27.92	0.7749	13.76	0.7858	14.70	0.6064	8.19	0.2933	17.22	0.6183
TransWeather	CVPR 22	23.30	0.6386	27.90	0.7250	11.97	0.7139	15.06	0.5862	8.73	0.3535	17.39	0.6034
DGUNet	CVPR 22	25.90	0.7804	31.73	0.8602	23.96	0.9374	20.90	0.7960	18.16	0.8984	24.13	0.8545
IDR	CVPR 23	19.06	0.4820	16.77	0.1911	11.12	0.3419	13.36	0.2559	18.17	0.4522	15.70	0.3446
IR-SDE	ICML 23	23.38	0.7248	29.27	0.8731	12.64	0.7774	16.35	0.7615	8.87	0.4000	18.10	0.7074
PromptIR	NeurIPS 23	29.99	0.9013	33.99	0.9256	19.62	0.9108	20.47	0.8155	33.60	0.9871	27.53	0.9081
MambaIR	ECCV 24	23.51	0.6553	29.77	0.8724	20.86	0.8797	21.33	0.7860	28.33	0.9320	24.76	0.8247
AutoDIR	ECCV 24	21.44	0.6082	27.67	0.8086	16.89	0.7364	14.99	0.4959	20.07	0.7236	20.21	0.6745
DiffUIR	CVPR 24	29.83	0.9017	33.24	0.9141	24.49	0.9451	21.27	0.8306	25.22	0.9698	26.81	0.9122
InstructIR	ECCV 24	25.89	0.7790	33.80	0.9165	15.93	0.8586	15.19	0.6411	22.29	0.9016	22.62	0.8194
PG-WSSM	JARS 25	30.04	0.9039	33.40	0.9156	28.49	0.9650	24.59	0.8782	39.52	0.9879	31.21	0.9301
FoundIR	ICCV 25	29.86	0.9006	33.32	0.9128	24.80	0.9444	22.31	0.8611	29.91	0.9786	28.04	0.9195
Ours	-	30.07	0.9044	33.48	0.9163	29.18	0.9706	25.40	0.8867	40.05	0.9908	31.64	0.9338

, the proposed method achieves superior performance on most degradation types, except for a few cases where the leading results are not obtained. In terms of overall restoration quality, our approach reaches an average PSNR of 31.64 dB across the five degradation tasks, outperforming the second-best method PG-WSSM by 0.43 dB, and thus achieving the best overall results. The notable improvement clearly demonstrates the effectiveness of our frequency prompts mechanism in mitigating degradation interference and enhancing robust multi-degradation restoration performance.

In terms of individual restoration tasks, our method achieves different levels of improvement under various degradations. For the deblurring task, the performance gain over the second-best approach is relatively small, which may be attributed to the fact that motion blur typically causes more severe structural information loss, thereby increasing the difficulty of accurate recovery. Regarding the denoising task, we observe a slight performance gap compared with the best-performing method PromptIR. This trade-off is deeply rooted in the distinct mechanisms of frequency-domain versus spatial-domain prompting. As illustrated in the frequency spectrum visualization in Figure 1, Gaussian noise manifests as a uniform distribution across high-frequency bands, lacking the specific directional or structural patterns observed in motion blur or rain streaks. Existing spatial-based methods like PromptIR excel at pixel-wise correction, enabling them to effectively suppress such stochastic, non-structural noise through local feature modulation.

In contrast, our proposed frequency-domain prompts are designed to capture and disentangle global degradation patterns based on their spectral signatures. Since random noise presents as a global disturbance without a compact spectral representation, our frequency-aware filters are slightly less precise in distinguishing noise from high-frequency image textures compared to localized spatial attention. However, this marginal trade-off in pure denoising allows our model to gain significant robustness in handling complex, spatially coupled degradations. For instance, in the dehazing task, our approach surpasses the second-best method by 0.69 dB, benefiting from the proposed frequency prompt that compensates for low-frequency attenuation and effectively improves the overall contrast and visibility of the restored images. For the de-raindrop task and low-light enhancement, our method achieves remarkable improvements of 0.81 dB and 0.53 dB, respectively. This demonstrates that while slightly less sensitive to random noise, the frequency prompts provide superior guidance in distinguishing occlusion artifacts (rain) from background textures and ensuring robust luminance recovery.

To further assess the restoration capability of the proposed method in UAV imaging scenarios, we provide visual comparisons across various degradation tasks. Figure 3, Figure 4, Figure 5 and Figure 6 illustrate the results of our approach and several representative methods under different degradation conditions. Overall, compared with existing advanced approaches, our method preserves better global visual consistency, enhances contrast and dynamic range, and substantially improves texture reconstruction quality, achieving superior structural fidelity and color restoration.

Figure 3 presents the deblurring results. As observed, our method effectively suppresses streak-like motion artifacts and significantly improves object boundaries and overall clarity. In contrast, TransWeather and AutoDIR still retain noticeable blur, while MambaIR and DiffUIR generate overly smoothed results, leading to the loss of fine details.

Figure 4 shows dehazing performance under challenging non-uniform haze conditions, where strong local scattering causes spatially varying degradation. Our method better captures low-frequency haze components and reduces scattering halos, restoring natural color appearances for buildings and vegetation. In comparison, PromptIR and InstructIR fail to fully remove haze and produce residual artifacts, while AutoDIR and DiffUIR partially suppress haze but suffer from luminance imbalance and color shifts.

Figure 5 illustrates the de-raindrop results. AutoDIR and InstructIR still exhibit prominent residual rain droplets or blur artifacts, whereas DGUNet and PromptIR mistakenly treat large rain streaks as background textures, leaving streak remnants. By contrast, our method effectively removes raindrop occlusions and streak artifacts without compromising background details.

Figure 6 presents the low-light enhancement results. Our approach achieves a more natural illumination enhancement by simultaneously improving brightness, suppressing noise, and maintaining chromatic consistency. DGUNet, although increasing brightness, fails to recover fine-grained textures, while IR-SDE suffers from insufficient contrast enhancement or abnormal color saturation, degrading the overall perceptual quality.

Evaluations on MDRS-Landsat for Remote Sensing Image Restoration. To further validate the cross-scene generalization capability of the proposed method, we conduct experiments on the MDRS-Landsat dataset and compare our approach with several representative image restoration methods, including NAFNet [48], Restormer [45], DGUNet [47], TransWeather [50], AirNet [24], PromptIR [28], IDR [11], SrResNet-AP [57], Restormer-AP [57], Uformer-AP [57], and Ada4DIR-s [52]. All compared methods follow the official train/test split and unified experimental configuration of MDRS-Landsat. Specifically, we randomly select 2048 samples from the dataset as the training set, with 512 images for each degradation type, while the entire validation and test sets are used for evaluation under their respective tasks.

As shown in

Table 3. Quantitative comparisons on the MDRS-Landsat dataset for different methods. The best results are highlighted in bold.

Model	Venue	Blur		Noise		Haze		Lowlight		Average
		PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
NAFNet	ECCV 22	33.10	0.8120	33.08	0.8263	31.56	0.9642	30.40	0.9516	32.04	0.8885
Restormer	CVPR 22	35.23	0.8559	34.53	0.8589	36.18	0.9867	37.86	0.9872	35.95	0.9222
DGUNet	CVPR 22	29.64	0.7822	30.31	0.7314	27.45	0.9338	27.15	0.9010	28.64	0.8371
TransWeather	CVPR 22	33.45	0.8159	33.69	0.8428	35.02	0.9689	36.33	0.9705	34.62	0.8995
AirNet	CVPR 22	28.27	0.7887	30.30	0.7446	24.39	0.9331	28.38	0.9472	27.84	0.8534
PromptIR	NIPS 23	36.41	0.8861	34.99	0.8729	37.61	0.9897	39.09	0.9900	37.02	0.9347
IDR	CVPR 23	36.57	0.8902	34.88	0.8681	36.99	0.9892	35.19	0.9865	35.91	0.9335
SrResNet-AP	CVPR 24	34.63	0.8479	34.70	0.8620	34.78	0.9825	33.87	0.9823	34.50	0.9187
Restormer-AP	CVPR 24	35.75	0.8732	34.96	0.8697	37.36	0.9888	37.27	0.9885	36.34	0.9301
Uformer-AP	CVPR 24	34.64	0.8488	34.39	0.8533	36.06	0.9877	36.58	0.9899	35.41	0.9199
Ada4DIR-s	IF 25	35.99	0.8792	34.79	0.8642	37.80	0.9898	41.32	0.9918	37.47	0.9313
Ours	-	38.46	0.9063	36.85	0.8936	37.87	0.9927	40.39	0.9908	38.39	0.9461

, our method achieves an average PSNR of 38.39 dB over four degradation types, outperforming the second-best approach, Ada4DIR-s, by 0.92 dB. This significant improvement verifies the strong effectiveness and generalization capability of combining frequency prompts with state space modeling for unified restoration in multi-degradation remote sensing scenarios.

Regarding specific restoration tasks, our method delivers consistent and notable performance gains across most degradation types. For the deblurring task, our approach surpasses IDR by 1.89 dB, demonstrating its superior ability to recover structural information. For the dehazing task, we obtain an improvement of 0.07 dB, indicating robustness in compensating low-frequency contrast attenuation. In the denoising task, our method achieves 36.85 dB, outperforming PromptIR by 1.86 dB, highlighting a better balance between noise suppression and detail preservation. For the low-light enhancement task, although our model achieves the second-best PSNR, it still shows clear advantages over most competing methods. Overall, even when trained with only 2048 samples, our method maintains excellent performance across all degradation categories in the remote sensing domain, demonstrating strong generalization ability under limited data conditions.

Evaluation of Model Complexity. To further assess the computational efficiency of our model, we compare its complexity with recent state-of-the-art approaches in terms of floating-point operations (FLOPs) and the number of learnable parameters. As reported in

Table 4. Comparison of model complexity, in terms of the number of parameters (#Param) and FLOPs, for different methods. FLOPs are computed with an input resolution of $256\times 256$.

Method	TransWeather	DGUNet	IR-SDE	PromptIR	MambaIR	AutoDIR	DiffUIR	InstructIR	Ours
#Param (M)	37.93	17.33	34.20	35.59	30.90	115.00	36.60	15.80	24.90
FLOPs (G)	6.14	199.87	98.30	158.14	126.29	63.76	9.88	16.20	55.10

, our method achieves an effective balance between computational cost and parameter capacity, demonstrating competitive efficiency while maintaining strong restoration performance.

4.4. Ablation Study

In this section, we validate the effectiveness of the key components integrated into our network through comprehensive ablation experiments conducted on the UAVIR-5D dataset. We report the average PSNR results in Table 3. To ensure fairness and reliability, all experiments follow the same training settings described in Section 4.2.

Effectiveness of Key Components. We first analyze the contribution of the two major components proposed in our architecture, namely the PGMB and the AFPB. In the experiments, SS2D denotes the Mamba block equipped with a 2D selective scanning strategy, and APB refers to an adaptive prompt block operating purely in the spatial domain. As shown in

Table 5. Quantitative ablation results of key components. SS2D denotes the selective 2D scanning strategy, and APB denotes the general prompt module. The checkmark (✓) indicates that the corresponding component or strategy is employed in the experiment.

Method	PGMB		AFPB	APB	PSNR	SSIM
	VSSM	SS2D
(a)		✓			28.05	0.9049
(b)	✓				28.72	0.9115
(c)		✓		✓	29.18	0.9191
(d)		✓	✓		29.99	0.9256
(e)	✓			✓	30.85	0.9276
Ours	✓		✓		31.64	0.9338

, each component contributes noticeably to the overall performance, indicating that both modules play indispensable roles in multi-degradation restoration. Moreover, the visual comparisons in Figure 7 further demonstrate the qualitative benefits introduced by each component.

Effectiveness from Structural and Parametric Perspectives. Following standard ablation study protocols, we extended our analysis to investigate the optimal configuration of the prompt modules from both structural and parametric perspectives. First, regarding the structural placement,

Table 6. Comparison of AFPB module at different positions. The best results are highlighted in bold.

Method	PSNR	SSIM
AFPB in Encoder	30.49	0.9226
AFPB in Decoder	31.18	0.9265
AFPB in Skip Connection	31.64	0.9338

demonstrates that integrating the AFPB into the skip connections is crucial for performance enhancement. This highlights that a rational structural design is as critical as the presence of the module itself. Second, to provide a detailed parameter comparison, we investigated the impact of the number of Prompt Components (N) on restoration performance. As reported in

Table 7. Ablation study on the number of Prompt Components. We investigate the impact of the component quantity (N) on restoration performance.

Model	Number (N)	PSNR (dB)	SSIM
(a)	3	30.31	0.9275
(b)	5	31.64	0.9338
(c)	7	31.60	0.9332
(d)	9	31.48	0.9335

, the results exhibit a clear trend: the performance peaks at N = 5 with a PSNR of 31.64 dB. Specifically, at N = 3, the restricted representational capacity limits the restoration quality. Conversely, increasing N beyond 5 to 7 and 9 results in performance saturation or slight degradation rather than improvement. This non-monotonic trend implies that excessive components may introduce semantic redundancy or feature interference. Consequently, N = 5 is adopted as the optimal balance between representation capability and model compactness.

Beyond these quantitative analyses, we further examine the internal working mechanism of the proposed modules. The AFPB functions effectively because different degradation types exhibit distinct spectral signatures—for instance, haze concentrates in low-frequency bands, whereas noise dominates high-frequency regions. By explicitly separating these components in the Fourier domain, the AFPB enables the model to apply targeted suppression or enhancement strategies, thereby effectively addressing the challenge of decoupling distinct degradation components, which remains difficult for spatial-domain methods (such as standard convolution). This rationale explains why our method outperforms existing approaches in complex coupled tasks such as dehazing and raindrop removal, as verified by the superior PSNR results.

4.5. Limitations and Discussion

Although the proposed Degradation-Aware Frequency Prompt State Space Model achieves state-of-the-art performance across multiple tasks, we have identified several limitations and potential failure cases that merit further discussion.

First, regarding extreme coupled degradation scenarios, the model’s reliance on implicit frequency-domain encoding faces challenges when multiple severe degradations co-occur. As raised in our experimental analysis, restoration performance depends heavily on the accuracy of self-generated frequency prompts. Under extremely challenging conditions—such as the simultaneous presence of dense haze and low illumination—the signal-to-noise ratio drops critically. In these cases, the spectral characteristics of haze and low-light become highly entangled in the low-frequency band. The implicit degradation encoding may fail to correctly disentangle these dominant degradation patterns from the inherent image structure, potentially leading to failure cases characterized by suboptimal brightness recovery or residual color casts.

Second, there is an inherent performance trade-off regarding denoising. As indicated in Table 1, our method performs slightly worse (−0.51 dB) than the spatial-domain method PromptIR in pure denoising tasks. This is a consequence of our architectural design choice: our frequency prompts prioritize capturing global degradation patterns (like streaks or haze distribution) to maximize all-in-one generalization. However, Gaussian noise manifests as a uniform distribution across high frequencies without specific structural patterns. Consequently, our frequency-aware filters are less sensitive to such stochastic pixel-level noise compared to spatially localized modulation methods. We acknowledge this as a necessary trade-off to achieve superior robustness in handling complex, spatially variant degradations.

Third, the domain gap between synthetic and real-world data remains a challenge. Our model is trained on the UAVIR-5D dataset, which is constructed using physics-based synthesis. However, real-world UAV imagery often suffers from more complex, non-uniform artifacts that are difficult to simulate perfectly, such as the coupling of video compression blocking, sensor jitter caused by high-speed flight, and atmospheric turbulence. When applied to real-world wild data, the model may exhibit minor artifacts or reduced restoration quality due to these unseen degradation combinations, limiting its flexibility compared to multi-expert systems trained on diverse real data.

Finally, we must consider computational constraints for practical deployment. While the Mamba-based architecture offers linear complexity and is more efficient than standard Transformers, the proposed AFPB introduces Fast Fourier Transform (FFT) operations. FFT involves complex number calculations that can introduce additional inference latency, especially on edge computing devices with limited hardware acceleration for spectral operations. This could potentially act as a bottleneck for time-sensitive UAV navigation missions requiring high-FPS processing.

In future work, we plan to address these limitations by exploring controllable prompt mechanisms with user intervention to handle extreme cases, and by investigating lightweight spectral approximations to reduce the computational overhead on UAV edge devices.

5. Conclusions

This paper proposes a degradation-aware prompt-guided state space model for all-in-one UAV image restoration. By introducing a Prompt-Guided Mamba Block, the method enables global context modeling with degradation-aware feature learning. In addition, an Adaptive Frequency Prompt Block is designed to capture spectral differences among various degradation types and dynamically guide restoration in the frequency domain. The experimental results demonstrate that our approach achieves outstanding performance across multiple representative degradation tasks, effectively improving the visual quality and reliability of UAV imagery. In the future, we plan to explore finer-grained degradation semantic prompting and cross-domain adaptation strategies to enhance robustness under extreme degradations and real-world scenarios, providing more practical image quality support for intelligent UAV perception applications.