pytorch实现supervisely 发布的人像分割数据集

整个工程文件已放到Github上：

https://github.com/yaoyi30/PyTorch_Image_Segmentation

一、训练图像分割网络主要流程

1. 构建数据集
2. 数据预处理、包括数据增强和数据标准化和归一化
3. 构建网络模型
4. 设置学习率、优化器、损失函数等超参数
5. 训练和验证

二、各个流程简要说明

1. 构建数据集

本文使用supervisely 发布的人像分割数据集，百度网盘地址：https://pan.baidu.com/s/1B8eBqg7XROHOsm5OLw-t9g提取码: 52ss

在工程目录下，新建data文件夹，在文件夹内分别新建last和last_msk文件夹，用来放图片和对应的mask图片,结构如下：

加载数据集代码：

from torch.utils.data import Dataset
import os
import cv2
import numpy as np


class MyDataset(Dataset):
    def __init__(self, train_path, transform=None):
        self.images = os.listdir(train_path + '/last')
        self.labels = os.listdir(train_path + '/last_msk')
        assert len(self.images) == len(self.labels), 'Number does not match'
        self.transform = transform
        self.images_and_labels = []  # 存储图像和标签路径
        for i in range(len(self.images)):
            self.images_and_labels.append((train_path + '/last/' + self.images[i], train_path + '/last_msk/' + self.labels[i]))

    def __getitem__(self, item):
        img_path, lab_path = self.images_and_labels[item]
        img = cv2.imread(img_path)
        img = cv2.resize(img, (224, 224))
        lab = cv2.imread(lab_path, 0)
        lab = cv2.resize(lab, (224, 224))
        lab = lab / 255  # 转换成0和1
        lab = lab.astype('uint8')  # 不为1的全置为0
        lab = np.eye(2)[lab]  # one-hot编码
        lab = np.array(list(map(lambda x: abs(x - 1), lab))).astype('float32')  # 将所有0变为1（1对应255， 白色背景），所有1变为0（黑色，目标）
        lab = lab.transpose(2, 0, 1)  # [224, 224, 2] => [2, 224, 224]
        if self.transform is not None:
            img = self.transform(img)
        return img, lab

    def __len__(self):
        return len(self.images)


if __name__ == '__main__':
    img = cv2.imread('data/train/last_mask/ache-adult-depression-expression-41253.png', 0)
    img = cv2.resize(img, (16, 16))
    img2 = img / 255
    img3 = img2.astype('uint8')
    hot1 = np.eye(2)[img3]
    hot2 = np.array(list(map(lambda x: abs(x - 1), hot1)))
    print(hot2.shape)
    print(hot2.transpose(2, 0, 1))

模型代码：

import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.encode1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.encode2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.encode3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.encode4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.encode5 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.decode1 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=3,
                               stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True)
        )
        self.decode2 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 3, 2, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(True)
        )
        self.decode3 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 3, 2, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.decode4 = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 3, 2, 1, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(True)
        )
        self.decode5 = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 3, 2, 1, 1),
            nn.BatchNorm2d(16),
            nn.ReLU(True)
        )
        self.classifier = nn.Conv2d(16, 2, kernel_size=1)

    def forward(self, x):  # b: batch_size
        out = self.encode1(x)  # [b, 3, 224, 224]  =>  [b, 64, 112, 112]
        out = self.encode2(out)  # [b, 64, 112, 112] =>  [b, 128, 56, 56]
        out = self.encode3(out)  # [b, 128, 56, 56]  =>  [b, 256, 28, 28]
        out = self.encode4(out)  # [b, 256, 28, 28]  =>  [b, 512, 14, 14]
        out = self.encode5(out)  # [b, 512, 14, 14]  =>  [b, 512, 7, 7]
        out = self.decode1(out)  # [b, 512, 7, 7]    =>  [b, 256, 14, 14]
        out = self.decode2(out)  # [b, 256, 14, 14]  =>  [b, 128, 28, 28]
        out = self.decode3(out)  # [b, 128, 28, 28]  =>  [b, 64, 56, 56]
        out = self.decode4(out)  # [b, 64, 56, 56]   =>  [b, 32, 112, 112]
        out = self.decode5(out)  # [b, 32, 112, 112] =>  [b, 16, 224, 224]
        out = self.classifier(out)  # [b, 16, 224, 224] =>  [b, 2, 224, 224]   2表示类别数，目标和非目标两类
        return out


if __name__ == '__main__':
    img = torch.randn(2, 3, 224, 224)
    net = Net()
    sample = net(img)
    print(sample.shape)

训练代码：

# python
# 导入必要的库
import os  # 导入操作系统库，用于文件操作
import model  # 导入自定义的模型模块
import torch  # 导入PyTorch库
import torch.nn as nn  # 导入PyTorch神经网络模块
import torch.optim as optim  # 导入PyTorch优化器模块
import numpy as np  # 导入NumPy库，用于数组操作
from load_img import MyDataset  # 从load_img模块中导入自定义的MyDataset类
from torchvision import transforms  # 导入torchvision的transforms模块，用于图像预处理
from torch.utils.data import DataLoader  # 导入PyTorch的数据加载器模块

# 设置超参数
batchsize = 16  # 设置每个批次的大小为8
epochs = 50  # 设置训练的总轮数为50
train_data_path = 'data/train'  # 设置训练数据的路径

# 定义图像预处理流程：转换为tensor并归一化
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化处理，使用预定义的均值和标准差
])

# 加载数据集，并设置批处理大小和是否打乱数据
bag = MyDataset(train_data_path, transform)  # 初始化MyDataset类，加载数据并应用预处理
dataloader = DataLoader(bag, batch_size=batchsize, shuffle=True)  # 使用DataLoader来加载数据，设置批次大小和打乱数据

# 选择使用GPU进行计算
device = torch.device('cuda')  # 设置设备为CUDA（即GPU）

# 将模型加载到GPU上
net = model.Net().to(device)  # 初始化模型，并将其移动到GPU上

# 设置损失函数和优化器
criterion = nn.BCELoss()  # 使用二分类交叉熵损失函数
optimizer = optim.SGD(net.parameters(), lr=1e-2, momentum=0.7)  # 使用随机梯度下降优化器，设置学习率和动量

# 检查是否存在保存模型的文件夹，如果不存在则创建
if not os.path.exists('checkpoints'):
    os.mkdir('checkpoints')

# 开始训练循环
for epoch in range(1, epochs + 1):  # 对每个epoch进行循环
    for batch_idx, (img, lab) in enumerate(dataloader):  # 对每个batch的数据进行循环
        img, lab = img.to(device), lab.to(device)  # 将数据和标签移动到GPU上
        # 通过模型进行前向传播，并使用sigmoid函数得到输出
        output = torch.sigmoid(net(img))  # sigmoid函数用于将输出转换为概率值
        # 计算损失
        loss = criterion(output, lab)  # 使用损失函数计算损失值

        # 每20个batch打印一次损失值
        if batch_idx % 20 == 0:
            print('Epoch:[{}/{}]\tStep:[{}/{}]\tLoss:{:.6f}'.format(
                epoch, epochs, (batch_idx + 1) * len(img), len(dataloader.dataset), loss.item()
            ))

        # 清空梯度，反向传播，更新权重
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新模型的权重

    # 每10个epoch保存一次模型
    if epoch % 10 == 0:
        torch.save(net, 'checkpoints/model_epoch_{}.pth'.format(epoch))  # 保存模型
        print

测试代码：

# 导入必要的库
import torch  # 导入PyTorch库
import cv2  # 导入OpenCV库，用于图像处理
from torch.utils.data import Dataset, DataLoader  # 从PyTorch库中导入Dataset和DataLoader，用于构建数据集和加载数据
from torchvision import transforms  # 从PyTorch的vision库中导入transforms，用于图像预处理
import numpy as np  # 导入NumPy库，用于数值计算
import os  # 导入os库，用于处理文件路径


# 定义一个名为TestDataset的类，继承自Dataset
class TestDataset(Dataset):
    def __init__(self, test_img_path, transform=None):
        # 初始化方法，接收测试图像的路径和转换方法作为参数
        self.test_img = os.listdir(test_img_path)  # 获取测试图像路径下的所有文件名
        self.transform = transform  # 将转换方法保存为类的属性
        self.images = []  # 初始化一个空列表，用于存储完整的图像路径
        for i in range(len(self.test_img)):  # 遍历所有文件名
            self.images.append(os.path.join(test_img_path, self.test_img[i]))  # 拼接完整的图像路径并添加到列表中

    def __getitem__(self, item):
        # 根据索引获取图像的方法
        img_path = self.images[item]  # 根据索引获取图像的完整路径
        img = cv2.imread(img_path)  # 使用OpenCV读取图像
        img = cv2.resize(img, (224, 224))  # 将图像缩放到224x224的大小
        if self.transform is not None:  # 如果有转换方法
            img = self.transform(img)  # 使用转换方法处理图像
        return img  # 返回处理后的图像

    def __len__(self):
        # 返回数据集的长度的方法
        return len(self.test_img)  # 返回图像文件名的数量


# 定义测试图像的路径和模型检查点的路径
test_img_path = 'data/test/last'
checkpoint_path = 'checkpoints/model_epoch_50.pth'

# 定义图像转换的流水线
transform = transforms.Compose([transforms.ToTensor(),  # 将图像转换为Tensor
                                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])  # 对图像进行归一化处理

# 使用定义的TestDataset类创建数据集对象
bag = TestDataset(test_img_path, transform)

# 使用DataLoader加载数据集
dataloader = DataLoader(bag, batch_size=1, shuffle=None)  # 批处理大小为1，不打乱数据顺序

# 加载模型检查点
net = torch.load(checkpoint_path)

# 将模型移动到GPU上
net = net.cuda()

# 遍历数据加载器中的每一个数据批次
for idx, img in enumerate(dataloader):
    # 将图像数据移动到GPU上
    img = img.cuda()

    # 前向传播，得到模型输出
    output = torch.sigmoid(net(img))  # 对输出应用sigmoid函数，将其转换为概率值

    # 将输出从GPU转移到CPU，并转换为NumPy数组
    output_np = output.cpu().data.numpy().copy()

    # 找到每个样本输出中的最小值索引
    output_np = np.argmin(output_np, axis=1)

    # 挤压数组，移除可能的单维度条目
    img_arr = np.squeeze(output_np)

    # 将索引值转换为0-255的范围（这看起来是不正确的，因为索引通常是整数，而不是概率值）
    img_arr = img_arr * 255

    # 保存结果图像
    cv2.imwrite('result/%03d.png' % idx, img_arr)

    # 打印保存的图像的文件名
    print('result/%03d.png' % idx)