基于U-Net与EfficientNet的衣物智能分割

基于U-Net与EfficientNet的衣物智能分割

如何让机器像人一样精准地识别和分割图像中的服装？这是许多技术领域和行业正在追寻的目标。
本项目使用深度学习技术，从图像中准确提取人物的服装区域，旨在为图像后期处理和服装相关应用提供更智能、更高效的工具支持。

本项目采用了 U-Net 这一经典的分割架构，并将 EfficientNet 引入其中作为特征提取的核心部分。U-Net 因其跳跃连接的设计，能有效整合低层次的细节特征和高层次的语义信息，非常适合细粒度的分割任务。而 EfficientNet 则以其在计算效率和准确度上的良好平衡，为模型的整体表现进一步加码。两者结合，让模型在处理复杂服装样式时，既能识别准确，又能运行流畅。

整体而言，项目基于 PyTorch 框架构建，通过 segmentation_models_pytorch 库完成模型的实现与优化。在数据预处理方面，我们引入了数据增强和标准化技术，提升模型在不同场景下的适应性。同时，针对分割任务的特点，我们选择了特定的损失函数，使模型能够更快收敛，并更好地处理复杂边缘区域。

1. 加载与展示图像数据

在进行图像分割任务时，数据预处理是至关重要的一步。本部分展示了如何加载和处理服装数据集中的图像及其对应的掩膜（mask）文件，并通过可视化展示部分样本。数据集中包含服装的图像及其对应的掩膜，这些掩膜标注了图像中不同区域（如衣物、背景等）的类别信息。

该代码使用了 Pandas 和 NumPy 来处理数据集的路径信息，利用 glob 模块按特定模式读取图像和掩膜路径，并将其封装到一个 DataFrame 中。此外，使用 Matplotlib 可视化了部分图像及其掩膜，以便于检查数据的正确性和质量。

!pip install opencv-python  -i https://pypi.tuna.tsinghua.edu.cn/simple
!pip install albumentations -i https://pypi.tuna.tsinghua.edu.cn/simple
!pip install segmentation-models-pytorch

# 数据处理
import pandas as pd 
import numpy as np

# 图像处理
from PIL import Image
import matplotlib.pyplot as plt 

# 时间相关的操作
import time

# PyTorch核心工具
import torch
import torch.nn as nn  # 定义神经网络模块
import torch.nn.functional as F  # 提供神经网络中常用的函数，例如激活函数和损失函数
import torch.optim as optim  # 优化器模块
from torch.optim import lr_scheduler  # 学习率调度器
from torchvision.datasets import ImageFolder  # 用于加载文件夹组织的图像数据集

# PyTorch数据集工具
from torch.utils.data.sampler import SubsetRandomSampler  # 用于从数据集中随机抽样
from torch.utils.data import Dataset, DataLoader  # 自定义数据集和数据加载器

# 数据增强工具（PyTorch方式）
import torchvision  # 提供计算机视觉相关的工具包
from torchvision import datasets, models, transforms, utils  # 数据集加载、预训练模型、数据增强等
from torchvision.transforms import v2  # 数据增强的新版工具

# 图像处理工具
import cv2  
import os  
from glob import glob  # 查找文件路径模式
from tqdm import tqdm  # 用于显示循环进度条
import shutil  # 高级文件操作，如复制和移动

# 数据分割与评估工具
from sklearn.model_selection import train_test_split  # 数据集的分割
from sklearn.metrics import confusion_matrix, accuracy_score, classification_report  # 性能评估指标
import seaborn as sns  # 数据可视化工具，尤其适合绘制统计图表

# 数据增强工具（基于albumentations）
import albumentations as A  # 高效的数据增强工具包

# 语义分割模型
import segmentation_models_pytorch as smp  # 提供预训练的语义分割模型

# 读取标签信息
labels_df = pd.read_csv("/home/mw/input/clothing2688/clothing/labels (1).csv")

# 获取所有图片路径
img_paths = sorted(glob('/home/mw/input/clothing2688/clothing/jpeg_images/IMAGES/*.jpeg'))

# 获取所有掩膜路径
mask_paths = sorted(glob('/home/mw/input/clothing2688/clothing/jpeg_masks/MASKS/*.jpeg'))

df = pd.DataFrame({"img": img_paths, "mask": mask_paths})

# 查看前几行，确认数据读取正确
df.head()

show_imgs = 4  # 定义要显示的图片数量
idx = np.random.choice(len(df), show_imgs, replace=False)


fig, axes = plt.subplots(show_imgs*2//4, 4, figsize=(15, 8))
axes = axes.flatten() # 将 axes 数组扁平化，便于迭代

# 循环并显示每一个子图（axes 中的每个 ax）
for i, ax in enumerate(axes):
    new_i = i//2
    if i % 2 ==0 :
        full_path = df.loc[idx[new_i]]['img']
    else:
        full_path = df.loc[idx[new_i]]['mask']
    ax.imshow(plt.imread(full_path))
    basename = os.path.basename(full_path) 
    ax.set_title(basename)
    ax.set_axis_off()

2. 定义自定义数据集并应用数据增强技术

在深度学习中，数据预处理和数据增强是提高模型泛化能力的关键步骤。本部分展示了如何为服装图像分割任务准备训练数据和验证数据，并应用不同的数据增强技术，以增强模型对不同情况的适应性。通过将数据集进行随机裁剪、水平翻转、平移、缩放和旋转等增强操作，模型能够学习到更多的图像变换特征，提高其在实际应用中的表现。

我们首先定义了MyDataset类，继承自PyTorch的Dataset类，实现了自定义的数据集加载与处理。数据增强操作通过albumentations库完成，train_transforms和test_transforms分别应用于训练集和测试集。接着，通过train_test_split函数将数据集分割为训练集和验证集，并为每个数据集创建了PyTorch的DataLoader对象，便于批量加载数据进行训练与验证。

# 定义训练时的图像预处理与数据增强
train_transforms = A.Compose([
    A.Resize(576, 576), # 将图像大小调整为 576x576
    A.RandomCrop(height=512, width=512, always_apply=True), # 从中心进行 512x512 的随机裁剪，保持图像尺寸
    A.HorizontalFlip(p=0.5), # 随机水平翻转图像，翻转概率为 50%
    A.ShiftScaleRotate(shift_limit=0.01, scale_limit=(-0.04,0.04), rotate_limit=(-5,5), p=0.5), # 随机平移、缩放和旋转，保持图像结构的变化
])

# 测试集的图像预处理
test_transforms = A.Compose([
    A.Resize(512, 512),
])

# 自定义数据集类 MyDataset，继承自 PyTorch 的 Dataset 类
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, dataframe, transforms_=None):
        """
        初始化数据集
        :param dataframe: 包含图片路径和掩膜路径的DataFrame
        :param transforms_: 用于数据增强和转换的函数（默认是None）
        """
        self.df = dataframe 
        self.transforms_ = transforms_

         # 定义标准化的预处理，基于ImageNet的均值和标准差
        self.pre_normalize = v2.Compose([
            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        self.resize = [512, 512]
        self.class_size = 59
        
    def __len__(self):
        """
        返回数据集的大小
        :return: 数据集的样本数量
        """

        return len(self.df)
    
    def __getitem__(self, index):
        """
        获取一个样本（图片及其对应的掩膜）
        :param index: 样本索引
        :return: 一个包含图像和目标掩膜的字典
        """
        # 读取图片并转换为RGB格式
        img = cv2.cvtColor(cv2.imread(self.df.iloc[index]['img']), cv2.COLOR_BGR2RGB)
        
        # 读取掩膜，并转为灰度图
        mask = cv2.imread(self.df.iloc[index]['mask'],cv2.IMREAD_GRAYSCALE)
        
        # 对掩膜进行处理：将掩膜中大于最大类别数减去1的值设为0，避免出现无效类别
        mask[mask > self.class_size-1] = 0

        #数据增强
        aug = self.transforms_(image=img, mask=mask)            
        img, mask = aug['image'], aug['mask']

        img = img/255 # 将图像归一化至[0,1]区间
        img = self.pre_normalize(img) # 对图像进行标准化（减去均值，除以标准差）
        
        # 将图像转换为PyTorch的Tensor，并调整通道顺序为(C, H, W)
        img = torch.tensor(img, dtype=torch.float).permute(2, 0, 1)
        # 将掩膜转换为长整型的Tensor（适用于分类任务）
        target = torch.tensor(mask, dtype=torch.long)
        
        sample = {'x': img, 'y': target}
        return sample

# 选择设备：如果有可用的GPU则使用GPU，否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将数据集分割为训练集和验证集，20% 作为验证集
train_df, val_df = train_test_split(df, test_size=0.2)

# 创建训练集和验证集的数据集对象，传入对应的变换
train_dataset = MyDataset(train_df, train_transforms)
val_dataset = MyDataset(val_df, test_transforms)

BATCH_SIZE = 4

# 创建训练集的DataLoader对象
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

# 创建验证集的DataLoader对象
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)

print(f'len train: {len(train_df)}')
print(f'len val: {len(val_df)}')

3. UNet模型构建与权重加载

本部分使用了segmentation_models_pytorch（简称SMP）库构建了基于EfficientNet-B4编码器的UNet模型，并从本地加载了预训练的权重用于初始化模型的编码器部分。这种方式能够充分利用预训练模型的特征提取能力，提升训练效率和精度。

模型输入为RGB图像，输出为具有59个类别的分割掩膜。通过在本地加载EfficientNet-B4的权重并应用到模型的编码器部分，我们确保了编码器能够从训练开始就具有强大的特征表达能力。随后，我们通过随机生成的张量测试模型的前向传播，以确认模型构建正确且能够正常运行。

UNet模型原理与架构

UNet是一种典型的卷积神经网络架构，最初设计用于医学图像分割任务，特别是像素级别的二分类分割。它的核心思想是通过对称的编码器-解码器结构来实现图像的分割。UNet模型的关键特性是：

1. 对称的编码器-解码器结构：编码器通过一系列的卷积层和池化层提取图像的高层特征，解码器通过上采样层逐步恢复图像的空间分辨率，使得最终的输出能够与输入图像的尺寸一致。
2. 跳跃连接（Skip Connections）：为了有效地保留图像的空间信息，UNet将编码器阶段的特征图与解码器阶段相应的特征图进行拼接。这些跳跃连接帮助网络在解码阶段更好地恢复图像细节，避免信息丢失。
3. 细粒度信息的保留：通过跳跃连接，UNet能够利用低层特征与高层特征的组合，在分割任务中捕捉到更多细粒度的空间信息。

UNet架构分析

1. 输入层（Input Layer）：

   • 输入图像的尺寸为(batch_size, 3, H, W)，其中H和W是图像的高度和宽度（RGB图像）。

2. 编码器（Encoder）：

   • 编码器部分利用EfficientNet-B4，通过一系列卷积层和池化层提取图像的深层特征。EfficientNet-B4的特点是通过复合缩放（compound scaling）策略，既提高了网络的深度，也扩大了宽度和分辨率。
   • 在此阶段，图像的空间分辨率逐步下降，而特征图的深度逐渐增加。

3. 跳跃连接（Skip Connections）：

   • 在UNet中，编码器各层的输出会与对应的解码器层的输出进行拼接，这样能够将低层的细节特征与高层的语义特征结合起来，从而提高分割精度。

4. 解码器（Decoder）：

   • 解码器通过上采样操作逐步恢复图像的空间分辨率，同时结合编码器阶段的跳跃连接，确保分割结果能精确恢复细节。
   • 每个解码器层包含一个上采样过程和一个卷积过程，帮助恢复图像的尺寸和细节。

5. 输出层（Output Layer）：

   • 输出图像的尺寸和输入图像一致，输出的是每个像素的分类结果，通常通过softmax或sigmoid激活函数来得到每个类别的预测概率。

EfficientNet编码器

EfficientNet是一个高效的卷积神经网络架构，它通过自动化搜索来优化网络的深度、宽度和分辨率，在参数量和计算量上做出了良好的平衡。EfficientNet-B4是EfficientNet系列中的一个变种，它在精度和效率上表现出色。

在UNet模型中，EfficientNet-B4被用作编码器，负责提取输入图像的深度特征。其优势在于：

• 更高的准确性：相比传统的ResNet或者VGG网络，EfficientNet能够在相同的计算资源下提供更高的准确性。
• 较少的计算量和参数：通过对深度、宽度和输入分辨率进行联合优化，EfficientNet能在保持高准确性的同时减少计算量，适用于实际应用。

class_size = 59 # 类别数

# 定义UNet模型，使用EfficientNet-B4作为编码器
model = smp.Unet(
    encoder_name="efficientnet-b4",        # 选择EfficientNet-B4作为编码器（也可以选择如mobilenet_v2, efficientnet-b7等）
    encoder_weights=None,                 # 使用`None`表示不加载预训练的权重，而是从头开始训练
    in_channels=3,                        # 输入通道数（3代表RGB图像）
    classes=class_size,                   # 输出通道数（表示类别数量，这里是59类）
)

# 加载本地的模型权重
local_weights_path = "/home/mw/input/clothing2688/efficientnet-b4-6ed6700e.pth" 
state_dict = torch.load(local_weights_path) 

# 加载权重到模型的编码器部分
model.encoder.load_state_dict(state_dict)  

# 设置设备为CUDA（GPU）或CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 通过传入一个随机生成的输入张量，来检查模型是否能够正常运行
model(torch.randn((1, 3, 512, 512))).shape  

4. UNet模型训练

本部分实现了一个完整的训练循环，利用UNet架构与EfficientNet编码器对语义分割任务进行优化。代码包含训练与验证过程、学习率调度、模型保存以及早停机制，旨在提高训练效率和模型性能。在训练过程中，采用了Dice Loss作为损失函数，并通过多轮训练动态调整学习率以优化收敛速度。最后，验证损失连续多次未改善时，早停机制会中断训练，防止过拟合并节省资源。以下是完整的实现代码及其细节说明。

def train(dataloader, model, loss_fn, optimizer, lr_scheduler):
    # 获取数据集的大小（样本数量）和每个epoch的批次数量
    size = len(dataloader.dataset)
    num_batches = len(dataloader)

    # 设置模型为训练模式
    model.train()

    epoch_loss = 0
    epoch_iou_score = 0

    for batch_i, batch in enumerate(dataloader):
        
        x, y = batch['x'].to(device), batch['y'].to(device) # 将数据转移到GPU
        optimizer.zero_grad()
        pred = model(x)
        loss = loss_fn(pred, y)

        loss.backward() # 反向传播：计算梯度
        optimizer.step() # 更新模型的参数

        epoch_loss += loss.item() # .item() 用来获取单个标量的值
        y = y.round().long()
        pred = torch.argmax(pred,dim=1)

        # 计算TP、FP、FN、TN等统计信息，并使用micro平均计算IOU
        tp, fp, fn, tn = smp.metrics.get_stats(pred, y, mode='multiclass', num_classes=class_size)
        iou_score = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro").item()
        

        epoch_iou_score += iou_score # 累加IOU分数
        lr_scheduler.step() # 更新学习率
    
    # 返回当前epoch的平均损失和平均IOU分数
    return epoch_loss/num_batches, epoch_iou_score/num_batches



def test(dataloader, model, loss_fn):
    # 获取数据集的大小（样本数量）和每个epoch的批次数量
    size = len(dataloader.dataset)
    num_batches = len(dataloader)

    # 设置模型为评估模式
    model.eval() 
    epoch_loss = 0
    epoch_iou_score = 0

    # 在测试时不计算梯度
    with torch.no_grad():
        for batch_i, batch in enumerate(dataloader):
            x, y = batch['x'].to(device), batch['y'].to(device)

            # 前向传播：通过模型获取预测结果, 并计算损失
            pred = model(x)
            loss = loss_fn(pred, y)

            # 累加损失
            epoch_loss += loss.item()
            
            y = y.round().long()
            pred = torch.argmax(pred,dim=1)

            tp, fp, fn, tn = smp.metrics.get_stats(pred, y, mode='multiclass', num_classes=class_size)
            iou_score = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro").item()
            
            # 累加IOU分数
            epoch_iou_score += iou_score
            
    # 返回当前epoch的平均损失和平均IOU分数
    return epoch_loss/num_batches, epoch_iou_score/num_batches

损失函数与优化器

• Dice Loss：该损失函数通常用于处理不平衡数据的分割任务，它计算的是预测结果与真实标签的相似度，能够更好地处理类间不平衡问题。在此代码中使用的是DiceLoss，并指定了mode='multiclass'，即多类别的Dice损失。
• 优化器（Adam）：使用Adam优化器来更新模型参数，它是深度学习中常用的优化方法，结合了动量和自适应学习率，能在训练过程中动态调整学习率。
• 学习率调度器（StepLR）：在训练过程中使用StepLR调度器来控制学习率的变化，step_size=2000表示每2000步降低一次学习率，gamma=0.1表示学习率将缩小为原来的0.1。

EPOCHS = 30  # 设置训练的总轮数
logs = {
    'train_loss': [], 'val_loss': [],
    'train_iou_score': [], 'val_iou_score': [],
}

# 检查并创建模型保存的文件夹
if os.path.exists('checkpoints') == False:
    os.mkdir("checkpoints")

# 初始化参数
loss_fn = smp.losses.DiceLoss(mode='multiclass')
learning_rate = 0.001
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size = 2000, gamma=0.1)

# 如果验证损失在patience个epoch内没有提升，则停止训练
patience = 5
counter = 0
best_loss = np.inf

# 将模型转移到GPU（或CPU）
model.to(device)

# 开始训练循环，进行多个epoch训练
for epoch in tqdm(range(EPOCHS)):
    # 获取训练损失和训练集的IOU分数，以及验证损失和验证集的IOU分数
    train_loss, train_iou_score = train(train_loader, model, loss_fn, optimizer, step_lr_scheduler)
    val_loss, val_iou_score = test(val_loader, model, loss_fn)
    
    # 将本轮训练和验证的损失和IOU分数记录到日志中
    logs['train_loss'].append(train_loss)
    logs['val_loss'].append(val_loss)
    logs['train_iou_score'].append(train_iou_score)
    logs['val_iou_score'].append(val_iou_score)

    # 打印当前epoch的损失、IOU分数和学习率
    print(f'EPOCH: {str(epoch+1).zfill(3)} \
    train_loss: {train_loss:.4f}, val_loss: {val_loss:.4f} \
    train_iou_score: {train_iou_score:.3f}, val_iou_score: {val_iou_score:.3f} \
    lr: {optimizer.param_groups[0]["lr"]}')

    # 每个epoch结束后，保存当前模型的权重
    torch.save(model.state_dict(), "checkpoints/last.pth")

    # 如果验证损失有改善，保存模型为最佳模型
    if val_loss < best_loss:
        counter = 0
        best_loss = val_loss
        torch.save(model.state_dict(), "checkpoints/best.pth")
    else:
        counter += 1
        
    # 如果验证损失连续多个epoch没有改善，则停止训练
    if counter >= patience:
        print("Earlystop!")
        break

5. 可视化训练结果

在语义分割任务的训练中，监控模型的性能指标变化是评估模型训练效果的重要手段。本部分通过绘制损失函数和IoU分数的曲线，直观展示了训练过程中的表现。

plt.figure(figsize=(15,5))

# 第一个子图：绘制训练损失和验证损失
plt.subplot(1,2,1)
plt.plot(logs['train_loss'],label='Train_Loss')
plt.plot(logs['val_loss'],label='Validation_Loss')
plt.title('Train_Loss & Validation_Loss',fontsize=20)
plt.legend()

# 第二个子图：绘制训练IOU分数和验证IOU分数
plt.subplot(1,2,2)
plt.plot(logs['train_iou_score'],label='Train_Iou_Score')
plt.plot(logs['val_iou_score'],label='Validation_Iou_Score')
plt.title('Train_Iou_score & Validation_Iou_score',fontsize=20)
plt.legend()

6. UNet模型测试

在模型训练完成后，验证其泛化能力和实际应用表现至关重要。本部分展示了对测试集样本的预测结果，包括输入图像、模型预测结果和真实标签掩码。通过可视化这些输出，我们能够直观地评估模型的性能，观察模型是否能够准确地分割图像中的不同类别，进而分析其在实际应用中的效果。

# 定义自定义的数据集类，用于加载和处理测试数据
class TestDataset(torch.utils.data.Dataset):
    def __init__(self, dataframe, transforms_=None):
        """
        初始化数据集
        :param dataframe: 包含图片路径和掩膜路径的DataFrame
        :param transforms_: 用于数据增强和转换的函数（默认是None）
        """
        self.df = dataframe
        self.transforms_ = transforms_
        self.pre_normalize = v2.Compose([
            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        self.resize = [512, 512]
        self.class_size = 59
        
    def __len__(self):
        """
        返回数据集的大小
        :return: 数据集的样本数量
        """
        return len(self.df)
    
    def __getitem__(self, index):
        """
        获取一个样本（图片及其对应的掩膜）
        :param index: 样本索引
        :return: 一个包含图像和目标掩膜的字典
        """
        img = cv2.cvtColor(cv2.imread(self.df.iloc[index]['img']), cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.df.iloc[index]['mask'],cv2.IMREAD_GRAYSCALE)
        mask[mask > self.class_size-1] = 0
        aug = self.transforms_(image=img, mask=mask)            
        img, mask = aug['image'], aug['mask']

        # 复制原始图像和掩码，用于可视化（不会影响后续处理）
        img_view = np.copy(img)
        mask_view = np.copy(mask)

        img = img/255
        img = self.pre_normalize(img)

        # 将图像转换为torch tensor，并调整通道顺序 (从 HWC -> CHW)
        img = torch.tensor(img, dtype=torch.float).permute(2, 0, 1)
        # 将掩码转换为torch tensor，并确保类型是long类型（用于分类任务）
        target = torch.tensor(mask, dtype=torch.long)

        # 返回一个包含图像、掩码和可视化图像/掩码的字典
        sample = {'x': img, 'y': target, 'img_view':img_view, 'mask_view':mask_view}
        return sample


test_dataset = TestDataset(val_df, test_transforms)

# 加载训练好的最佳模型
model.load_state_dict(torch.load("checkpoints/best.pth"))  # 加载最好的模型权重
model.to(device)  # 将模型转移到设备（GPU或CPU）

# 设置展示的图片数量
show_imgs = 4
# 随机选择4个样本进行展示
random_list = np.random.choice(len(test_dataset), show_imgs, replace=False)

# 循环展示每个图像的预测结果
for i in range(show_imgs):
    idx = random_list[i] 
    sample = test_dataset[idx]  

    # 通过模型进行预测，将输入图像送入模型并获得预测结果
    pred = model(sample['x'].to('cuda', dtype=torch.float32).unsqueeze(0)) 
    pred = torch.argmax(pred, dim=1).squeeze(0)  
    pred = pred.data.cpu().numpy()  
    
    # 将预测结果（标签）转换为图像显示（8位灰度图）
    pred_view = Image.fromarray(np.uint8(pred), 'L')

    # 获取原图像和掩码用于展示
    img_view = sample['img_view'] 
    img_view = Image.fromarray(img_view, 'RGB')
    mask_view = sample['mask_view'] 
    mask_view = Image.fromarray(mask_view, 'L') 

    # 创建一个包含3个子图的图形，用于显示输入图像、预测结果和真实掩码
    f, axarr = plt.subplots(1, 3) 
    axarr[0].imshow(img_view)
    axarr[0].set_title('Input') 
    axarr[1].imshow(pred_view) 
    axarr[1].set_title('Pred')
    axarr[2].imshow(mask_view)
    axarr[2].set_title('GT') 
    plt.show()

# 若需要完整数据集以及代码 
https://mbd.pub/o/bread/mbd-aJWVmZlp