Restoration Adaptation for Semantic Segmentation on Low Quality Images

1. Problem Statement

Semantic segmentation models, particularly those based on CNNs and Transformers (e.g., Mask2Former, SegFormer), are typically trained on high-quality (HQ) standard images. In real-world scenarios, images often suffer from complex, unknown degradations (blur, noise, low resolution, compression artifacts).

• Direct Application Failure: Applying HQ-trained models directly to Low-Quality (LQ) images leads to significant performance drops because the models cannot extract robust features from degraded inputs.
• Limitations of Fine-tuning: Simply fine-tuning segmentation models on LQ data often fails to recover the lost structural details and semantic cues necessary for accurate pixel-level classification.
• Limitations of Pre-processing Restoration: Using standard image restoration (Real-IR) as a pre-processing step is suboptimal. Conventional restoration models focus on pixel-level fidelity and perceptual realism but often lack semantic awareness. They may hallucinate incorrect textures or structures (e.g., turning a blurred bag into a box), leading to misclassification in downstream segmentation tasks. Furthermore, cascading a restoration network and a segmentation network causes error propagation and increased computational cost.

2. Methodology: RASS Framework

The authors propose Restoration Adaptation for Semantic Segmentation (RASS), a unified framework that integrates image restoration directly into the segmentation process using a Stable Diffusion (SD) backbone and LoRA (Low-Rank Adaptation) modules. The framework operates in two stages:

Stage 1: Semantic-Constrained Restoration (SCR)

The goal is to train a restoration model that preserves semantic structures rather than just generating visually pleasing images.

• Backbone: Uses a pre-trained Text-to-Image (T2I) diffusion model (specifically OSEDiff architecture) with a frozen VAE encoder/decoder.
• Semantic Guidance: To prevent semantic hallucinations, the authors introduce a Semantic-Constrained Loss (SCL).
  • They align the cross-attention maps of the diffusion model with ground-truth semantic masks.
  • The loss function ($L_{SCL}$) consists of a region-level loss (ensuring attention focuses on the correct object area) and a pixel-level loss (binary cross-entropy) to enforce spatial consistency.
  • This forces the restoration model to reconstruct content that is semantically consistent with the input text and the underlying object structure.
• Training: The model is trained to reconstruct HQ images from LQ inputs using a combination of pixel-level loss ($L_{MSE}$), perceptual loss ($L_{LPIPS}$), and the semantic constraint loss.

Stage 2: Restoration Adaptation Segmentation (RAS)

The goal is to transfer the learned restoration knowledge into the segmentation task without retraining from scratch or adding heavy inference overhead.

• LoRA Merging: The trained SCR LoRA weights (which encode degradation-invariant features) are merged into the backbone of the segmentation model.
• Task-Specific Fine-tuning: A new Segmentation LoRA branch ($\Delta\theta_{seg}$) is added and fine-tuned specifically for the segmentation task.
• Feature Extraction: Instead of decoding the latent representation back to an image (which is computationally expensive), the model extracts hierarchical features from the VAE Encoder and the Latent Diffusion Network. These "cleaned" latent features are fed directly into a trainable segmentation head ($S_\phi$).
• Loss: The segmentation head is trained using a compound loss ($L_{seg}$) combining Cross-Entropy, Dice Loss, and classification loss.

3. Key Contributions

1. RASS Framework: The first framework designed specifically to address semantic segmentation on real-world LQ images by unifying restoration and segmentation, rather than treating them as separate sequential steps.
2. Semantic-Constrained Restoration (SCR): A novel restoration model that injects segmentation priors by aligning cross-attention maps with semantic masks. This ensures the restored image is not just visually clear but semantically faithful to the original scene structure.
3. Efficient Knowledge Transfer: Utilizes LoRA-based module merging to transfer restoration knowledge into the segmentation backbone. This allows the model to operate on "denoised" latent features during inference, avoiding the computational cost of full image reconstruction and error propagation.
4. RealLQ Dataset: The construction of a new, carefully annotated real-world LQ segmentation dataset containing 100 degraded images with high-quality annotations, covering 82 categories, to facilitate comprehensive evaluation.

4. Experimental Results

The authors evaluated RASS on both synthetic (ADE20K with simulated degradation) and real-world (RealLQ, RealSR, DrealSR) benchmarks.

• Segmentation Performance:
  • RASS achieved a 47.42 mIoU on the degraded ADE20K validation set (Fine-Tuning setting), outperforming the strongest baseline (Mask2Former fine-tuned) by 3.23 mIoU.
  • Under the "Restoration-to-Segmentation" (R2S) protocol, RASS also led with 46.81 mIoU, significantly outperforming other methods that rely on pre-processing.
  • On the RealLQ dataset, RASS achieved 39.80 mIoU, demonstrating superior robustness to real-world degradations compared to lightweight models which often degrade further after fine-tuning.
• Restoration Performance:
  • The SCR model demonstrated state-of-the-art performance in perceptual quality (MUSIQ, MANIQA, CLIPIQA) on RealSR and DIV2K datasets, often surpassing multi-step diffusion models while using only one inference step.
• Efficiency:
  • Although the backbone is large (1.7B parameters), RASS only requires training 27.9M parameters (via LoRA).
  • Inference speed is 3.4 FPS, which is competitive given the complexity of the task, and significantly faster than methods requiring full VAE decoding.

5. Significance

• Paradigm Shift: RASS challenges the traditional pipeline of "Restore $\to$ Segment." By fusing the tasks, it prevents the error propagation caused by disjoint pipelines and ensures that the restoration process is explicitly guided by the needs of the segmentation task.
• Robustness: It provides a viable solution for deploying vision systems in harsh, real-world environments (e.g., surveillance, autonomous driving in bad weather) where image quality is unpredictable.
• Semantic Fidelity: The introduction of semantic constraints into generative restoration addresses a critical flaw in current Real-IR models: the tendency to generate semantically incorrect details that confuse downstream tasks.
• Resource Efficiency: By leveraging LoRA and latent space processing, the method achieves high performance without the prohibitive computational cost of retraining massive models or running full image reconstruction at inference.

In conclusion, RASS represents a significant advancement in handling low-quality visual data, proving that integrating semantic awareness into the restoration process is crucial for robust downstream vision tasks.