当前位置：首页 > news >正文

人脸美颜磨皮祛痘3：深度学习SUNet神经网络实现图片修复(含训练代码、数据集和GUI交互界面)

news 来源：原创 2025/5/25 6:07:04

第一步：SUNet介绍

SUNet（Spatial U-Net）是一种基于U-Net架构改进的深度学习网络，主要用于图像分割（本文把它用在了图像修复任务上）任务，尤其在医学影像、遥感图像等需要高精度空间信息保留的领域表现突出。它通过结合空间注意力机制和U-Net的编码器-解码器结构，增强了模型对图像中重要区域的定位能力，同时保持了原始U-Net的多尺度特征融合优势。

核心特点

空间注意力机制（Spatial Attention）
- 在U-Net的跳跃连接（Skip Connection）中引入空间注意力模块，动态调整不同空间位置的特征权重，使网络更关注目标区域（如病灶、边缘等），抑制无关背景。
- 典型实现方式：通过卷积层生成注意力图（Attention Map），与原始特征图逐点相乘，突出关键区域。
多尺度特征融合
- 保留U-Net的经典编码器-解码器结构，通过下采样（编码器）捕获全局上下文，上采样（解码器）恢复空间细节，结合跳跃连接传递局部信息。
改进的跳跃连接
- 传统U-Net的跳跃连接直接拼接编码器和解码器特征，可能引入冗余信息。SUNet通过空间注意力筛选重要特征，提升分割精度。
轻量化设计（可选）
- 部分变体采用深度可分离卷积（Depthwise Separable Convolution）减少参数量，适合移动端或实时应用。

第二步：SUNet网络结构

SUnet遵循Unet的编码器-解码器结构，并采用跳跃连接来传达低级语义信息。通道号表示为、、和，定义为 = 64、= 128、= 320、= 512。ESRA Transformer块用作 SUnet 中的特征提取模块，本文将两个通道堆叠在一起每级都有 ESRA Transformer块。 ESRA Transformer块的示意图如图1（b）所示。 ESRA Transformer块内的 Mix-FFN 模块如图1（c）所示，它与传统的前馈网络（FFN）不同，它在两个线性层之间使用深度卷积。在 SUnet 模型中，重叠补丁嵌入使用内核大小为 7 × 7 且步幅为 3 的卷积层进行补丁嵌入。重叠补丁合并采用内核大小为 3 × 3 的卷积层。嵌入补丁的重叠属性有助于减轻传统补丁嵌入造成的信息丢失。图1（d）展示了EFF的架构，主要由三个子模块组成：EAG、ECA和SA。本文使用组数为 32 的分组卷积和残差连接来改进原始 AG。 EAG模块通过上采样获得的高级语义特征来增强通过跳跃连接传输的低级语义特征。 EAG的结构如图1(e)所示。连接后，ECA和SA主要用于突出特征图中任务相关区域的重要通道和空间位置，以提高特征表达能力。应该注意的是，底部ECA和SA模块仅包含单个输入特征，因此本文仅使用ECA和SA来强调特征。

第三步：模型代码展示

import torch.nn as nn
from model.SUNet_detail import SUNetclass SUNet_model(nn.Module):def __init__(self, config):super(SUNet_model, self).__init__()self.config = configself.swin_unet = SUNet(img_size=config['SWINUNET']['IMG_SIZE'],patch_size=config['SWINUNET']['PATCH_SIZE'],in_chans=3,out_chans=3,embed_dim=config['SWINUNET']['EMB_DIM'],depths=config['SWINUNET']['DEPTH_EN'],num_heads=config['SWINUNET']['HEAD_NUM'],window_size=config['SWINUNET']['WIN_SIZE'],mlp_ratio=config['SWINUNET']['MLP_RATIO'],qkv_bias=config['SWINUNET']['QKV_BIAS'],qk_scale=config['SWINUNET']['QK_SCALE'],drop_rate=config['SWINUNET']['DROP_RATE'],drop_path_rate=config['SWINUNET']['DROP_PATH_RATE'],ape=config['SWINUNET']['APE'],patch_norm=config['SWINUNET']['PATCH_NORM'],use_checkpoint=config['SWINUNET']['USE_CHECKPOINTS'])def forward(self, x):if x.size()[1] == 1:x = x.repeat(1, 3, 1, 1)logits = self.swin_unet(x)return logitsif __name__ == '__main__':from utils.model_utils import network_parametersimport torchimport yamlfrom thop import profilefrom utils.model_utils import network_parameters## Load yaml configuration filewith open('../training.yaml', 'r') as config:opt = yaml.safe_load(config)Train = opt['TRAINING']OPT = opt['OPTIM']height = 256width = 256x = torch.randn((1, 156, height, width))  # .cuda()model = SUNet_model(opt)  # .cuda()out = model(x)flops, params = profile(model, (x,))print(out.size())print(flops)print(params)

第四步：运行交互代码