pytorch U²-Net教程

news2024/9/23 6:35:43

U²-Net (U2-Net) 是一个用于图像分割的神经网络模型，特别擅长于边界复杂的物体分割任务，如前景背景分割和抠图。U²-Net 的独特之处在于其 U 形结构和嵌套 U 形块，能够有效捕捉不同尺度的特征，同时保持较小的模型大小。它非常适合在资源受限的环境下使用。

官方文档链接

U²-Net 本身并没有一个独立的 Python 库，但可以通过官方 GitHub 仓库获取源码和模型细节。

一、U²-Net 架构概述

U²-Net 是基于 U-Net 结构的改进模型，由多个嵌套的 U 形编码器-解码器模块组成。其创新点在于 U2 模块，它在不同尺度上提取特征，增强了对边界信息的捕捉能力。

U²-Net 结构包含：

编码器（Encoder）：使用多尺度卷积核提取图像的特征，逐渐压缩特征图尺寸。
解码器（Decoder）：通过逐步上采样，恢复原始分辨率，同时结合编码器的跳跃连接。
U2 模块：嵌套的 U 形块，能够同时处理不同分辨率的特征，从而保留高分辨率的局部细节和低分辨率的全局语义信息。

二、基础功能

在 U²-Net 中，通常的工作流程是加载预训练模型并对输入图像进行分割。U²-Net 最常见的任务是图像前景提取，比如抠图。

1. 加载 U²-Net 模型

从官方 GitHub 下载预训练模型权重，并通过 PyTorch 加载。

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

# 加载预训练的 U²-Net 模型
model = torch.load('u2net.pth')
model.eval()  # 设置为评估模式

# 准备图像输入
def load_image(image_path):
    transform = transforms.Compose([
        transforms.Resize((320, 320)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)
    return image

# 加载图片并转换为张量
input_image = load_image("input_image.jpg")

# 前向传播，生成分割结果
with torch.no_grad():
    result = model(input_image)

2. 处理模型输出

U²-Net 的输出通常为前景掩码 (mask)，可以通过阈值处理生成二值化图像。

def process_output(output):
    # 提取前景掩码
    mask = output[0][0].squeeze().cpu().numpy()
    
    # 归一化到0-1范围
    mask = (mask - np.min(mask)) / (np.max(mask) - np.min(mask))
    
    # 二值化处理
    mask = (mask > 0.5).astype(np.uint8)
    
    return mask

# 处理输出的前景掩码
foreground_mask = process_output(result)

三、进阶功能

1. 前景提取并保存透明 PNG

U²-Net 可以用于精细化的图像前景提取。通过将背景像素设置为透明，生成透明的 PNG 图片。

from PIL import Image

def save_foreground(image_path, mask, save_path):
    image = Image.open(image_path).convert('RGBA')
    width, height = image.size
    mask = Image.fromarray(mask * 255).resize((width, height), Image.BILINEAR)
    
    # 转换为 RGBA 格式，将背景设置为透明
    image_data = np.array(image)
    mask_data = np.array(mask)
    
    # 将背景区域的 alpha 通道设置为 0（完全透明）
    image_data[:, :, 3] = mask_data
    
    # 保存带有透明背景的 PNG 图片
    output_image = Image.fromarray(image_data)
    output_image.save(save_path)

# 使用掩码提取前景并保存
save_foreground("input_image.jpg", foreground_mask, "output_image.png")

2. 使用其他输入尺寸

虽然 U²-Net 默认是使用 320x320 的输入尺寸，但它对不同的输入尺寸有一定的适应性。我们可以根据需要调整输入图像的大小。

# 自定义输入尺寸
def load_image_custom_size(image_path, size=(320, 320)):
    transform = transforms.Compose([
        transforms.Resize(size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)
    return image

# 调整输入图像尺寸
custom_size_image = load_image_custom_size("input_image.jpg", size=(512, 512))

四、高级教程

U²-Net 的高级用法可以结合其他深度学习框架或任务，例如对分割结果进行进一步的图像处理或增强。

1. 与 OpenCV 结合处理分割结果

可以利用 OpenCV 对分割后的图像进行一些后处理，例如边缘检测、轮廓提取等。

import cv2

def process_with_opencv(mask):
    # 使用 OpenCV 检测轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 绘制轮廓
    contour_image = np.zeros_like(mask)
    cv2.drawContours(contour_image, contours, -1, (255), 2)
    
    return contour_image

# 使用 OpenCV 处理分割结果
contour_image = process_with_opencv(foreground_mask)
cv2.imwrite("contour_image.png", contour_image)

2. 自定义损失函数与训练

如果需要训练自己的 U²-Net 模型，可以基于 Binary Cross Entropy 损失函数进行训练。以下是一个自定义损失函数的示例。

import torch.nn as nn

class U2NetLoss(nn.Module):
    def __init__(self):
        super(U2NetLoss, self).__init__()
        self.bce_loss = nn.BCELoss()

    def forward(self, d0, d1, d2, d3, d4, d5, d6, labels):
        # 对不同尺度的预测进行加权损失计算
        loss0 = self.bce_loss(d0, labels)
        loss1 = self.bce_loss(d1, labels)
        loss2 = self.bce_loss(d2, labels)
        loss3 = self.bce_loss(d3, labels)
        loss4 = self.bce_loss(d4, labels)
        loss5 = self.bce_loss(d5, labels)
        loss6 = self.bce_loss(d6, labels)
        return loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6