[Paper Review] SAM-DiffSR: Structure-Modulated Diffusion Model for Image Super-Resolution

Chengcheng Wang, Zhiwei Hao+5 more|arXiv (Cornell University)|Feb 27, 2024

Photoacoustic and Ultrasonic Imaging4 citations

TL;DR

This paper proposes SAM-DiffSR, a structure-modulated diffusion model for image super-resolution that leverages fine-grained segmentation masks from Segment Anything Model (SAM) to guide noise modulation during the forward diffusion process. By encoding structural position information into SAM masks and using them to independently adjust noise means per region, the method enhances texture and structure fidelity without additional inference cost, achieving a 0.74 dB PSNR gain over prior diffusion-based methods on DIV2K.

ABSTRACT

Diffusion-based super-resolution (SR) models have recently garnered significant attention due to their potent restoration capabilities. But conventional diffusion models perform noise sampling from a single distribution, constraining their ability to handle real-world scenes and complex textures across semantic regions. With the success of segment anything model (SAM), generating sufficiently fine-grained region masks can enhance the detail recovery of diffusion-based SR model. However, directly integrating SAM into SR models will result in much higher computational cost. In this paper, we propose the SAM-DiffSR model, which can utilize the fine-grained structure information from SAM in the process of sampling noise to improve the image quality without additional computational cost during inference. In the process of training, we encode structural position information into the segmentation mask from SAM. Then the encoded mask is integrated into the forward diffusion process by modulating it to the sampled noise. This adjustment allows us to independently adapt the noise mean within each corresponding segmentation area. The diffusion model is trained to estimate this modulated noise. Crucially, our proposed framework does NOT change the reverse diffusion process and does NOT require SAM at inference. Experimental results demonstrate the effectiveness of our proposed method, showcasing superior performance in suppressing artifacts, and surpassing existing diffusion-based methods by 0.74 dB at the maximum in terms of PSNR on DIV2K dataset. The code and dataset are available at https://github.com/lose4578/SAM-DiffSR.

Motivation & Objective

• Address the limitation of conventional diffusion models in handling complex textures and structures across semantic regions due to uniform noise sampling.
• Integrate fine-grained structural information from the Segment Anything Model (SAM) into diffusion-based super-resolution to improve restoration quality.
• Develop a method that leverages SAM's segmentation capability without incurring additional computational cost during inference.
• Enable region-specific noise modulation during the forward diffusion process to enhance local structure and texture recovery.
• Maintain compatibility with existing reverse diffusion processes while injecting structural knowledge during training.

Proposed method

• Pre-generate high-quality segmentation masks for each HR image in the training set using the Segment Anything Model (SAM).
• Introduce a Structural Position Encoding (SPE) module that embeds 2D position information into the SAM masks using Rotary Position Embedding (RoPE).
• Modulate the mean of the noise during the forward diffusion process by applying the SPE-encoded mask to adjust noise independently per segmented region.
• Train the U-Net denoising model to predict the modulated noise, effectively learning to restore structure-aware details.
• Reuse pre-generated SAM masks across all training epochs, minimizing additional training cost.
• Ensure the final model operates without requiring SAM at inference, preserving inference efficiency.

Experimental results

Research questions

• RQ1 How can fine-grained structural information from SAM be effectively integrated into the diffusion process to improve image super-resolution?
• RQ2 What impact does the quality of segmentation masks have on the performance of diffusion-based super-resolution models?
• RQ3 Can structure-aware noise modulation during the forward diffusion process enhance texture and structure fidelity without increasing inference cost?
• RQ4 How does the choice of position encoding method affect the model's ability to localize structural details?
• RQ5 Does the proposed method outperform existing diffusion-based super-resolution approaches in terms of PSNR, SSIM, and artifact suppression?

Key findings

• SAM-DiffSR achieves a 0.74 dB improvement in PSNR over existing diffusion-based methods on the DIV2K dataset, demonstrating state-of-the-art performance.
• The model reduces artifacts significantly, as confirmed by both quantitative artifact map scores and qualitative visual comparisons.
• High-quality segmentation masks from the original SAM yield superior results, with PSNR increasing from 25.33 (low quality) to 25.54 (high quality) on Urban100.
• Using RoPE for structural position encoding leads to better performance than cosine or linear position encodings, with a 0.23 dB PSNR gain on Urban100.
• Quantitative results show consistent improvements in SSIM and FID across all datasets, confirming enhanced perceptual quality and diversity.
• The method achieves these gains without modifying the reverse diffusion process or requiring SAM at inference, maintaining zero additional inference cost.