图像分割FCN实战一

数据集是一个包包的图片

数据集的代码

'''
dataset.py
'''
import os
import torch
from torch.utils.data import DataLoader, Dataset, random_split
from torchvision import transforms
import numpy as np
import cv2

# transform是对图像进行预处理、数据增强等。Compose将多个处理步骤整合到一起。
# ToTensor：将原始取值0-255像素值，归一化为0-1
# Normalize：用像素值的均值和标准偏差对像素值进行标准化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])])

"""
这是一个用于将分类标签转换为one-hot编码的函数。
函数的输入参数包含data和n，其中data是原始的分类标签，n是分类的类别数。
输出结果是一个形状为（data.shape[0],data.shape[1],n)的numpy数组，表示每个样本在n个类别中的one-hot编码
"""
def onehot(data, n):
    buf = np.zeros(data.shape + (n,)) #创建一个形状为data.shape+（n，）的全零数组buf
    nmsk = np.arange(data.size) * n + data.ravel() #使用numpy的arange函数生成一个长度为data.size索引数组，并乘以n，加上data.ravel（）
                                                   # 得到每个元素在buf中的位置。
                                                   # ravel（）是一个numpy数组的方法，用于将多维数组展开成一维数组

    buf.ravel()[nmsk - 1] = 1  #将buf数组中nmask数组对应的位置设置为1，以得到每个样本在n个类别中的one—hot编码
    return buf


class BagDataset(Dataset):
    def __init__(self, transform=None):
        self.transform = transform

    def __len__(self):
        return len(os.listdir('./last'))

    def __getitem__(self, idx):
        # 读取原图
        img_name = os.listdir('./last')[idx] # 获取./last目录中的所有文件和子目录的名称列表，然后通过idx索引获取其中的一个文件夹名
        imgA = cv2.imread('./last/' + img_name) # 读取图像文件
        imgA = cv2.resize(imgA, (160, 160)) # 缩放图像大小

        # 读取标签图，即二值图
        imgB = cv2.imread('last_mask/' + img_name, 0) # 读取标签文件，灰度模式
        imgB = cv2.resize(imgB, (160, 160)) # 缩放便签缩放

        imgB = imgB / 255 # 将标签图像素值归一化[0,1]之间
        imgB = imgB.astype('uint8') # 将标签图像素值转为uint8类型
        imgB = onehot(imgB, 2)  # 因为此代码是二分类问题，即分割出手提包和背景两样就行，因此这里参数是2
        imgB = imgB.transpose(2, 0, 1)  # imgB不经过transform处理，所以要手动把(H,W,C)转成(C,H,W)
        imgB = torch.FloatTensor(imgB) # 将标签图转换为PyTorch张量
        if self.transform:
            imgA = self.transform(imgA)  # imgA通道就变成(C,H,W)
        return imgA, imgB


bag = BagDataset(transform)
train_size = int(0.9 * len(bag))  # 整个训练集中，90%为训练集
test_size = len(bag) - train_size

train_dataset, test_dataset = random_split(bag, [train_size, test_size])  # 按照上述比例(9:1)划分训练集和测试集
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=1)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=1)

if __name__ == '__main__':
    for train_batch in train_dataloader:
        print(train_batch)

    for test_batch in test_dataloader:
        print(test_batch)

模型

# -*- coding: utf-8 -*-
"""
model.py
"""
import torch.nn as nn
from torchvision.models.vgg import VGG


# 继承nn.Module，撰写自己的网络层
class FCNs(nn.Module):
    '''
    类FCNs：将最后一个特征图直接上采样32倍(5次步长为2、卷积核为3*3的反卷积操作)得到的最终
    分割结果。
    '''

    def __init__(self, pretrained_net, n_class):
        super().__init__()
        self.n_class = n_class
        self.pretrained_net = pretrained_net

        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2,
                                          padding=1, dilation=1,
                                          output_padding=1)
        self.bn1 = nn.BatchNorm2d(512)

        self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2,
                                          padding=1, dilation=1,
                                          output_padding=1)
        self.bn2 = nn.BatchNorm2d(256)

        self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2,
                                          padding=1, dilation=1,
                                          output_padding=1)
        self.bn3 = nn.BatchNorm2d(128)

        self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2,
                                          padding=1, dilation=1,
                                          output_padding=1)
        self.bn4 = nn.BatchNorm2d(64)

        self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2,
                                          padding=1, dilation=1,
                                          output_padding=1)
        self.bn5 = nn.BatchNorm2d(32)

        # 分类器是1*1大小的卷积，将channel个数从32减小到n_class
        self.classifier = nn.Conv2d(32, n_class, kernel_size=1)

    def forward(self, x):
        output = self.pretrained_net(x)
        x5 = output['x5']
        x4 = output['x4']
        x3 = output['x3']
        x2 = output['x2']
        x1 = output['x1']

        score = self.bn1(self.relu(self.deconv1(x5)))
        score = score + x4
        score = self.bn2(self.relu(self.deconv2(score)))
        score = score + x3
        score = self.bn3(self.relu(self.deconv3(score)))
        score = score + x2
        score = self.bn4(self.relu(self.deconv4(score)))
        score = score + x1
        score = self.bn5(self.relu(self.deconv5(score)))
        score = self.classifier(score)
        return score


class VGGNet(VGG):
    def __init__(self, pretrained=False, model='vgg16', requires_grad=True, remove_fc=True, show_params=False):
        super().__init__(make_layers(cfg[model]))
        self.ranges = ranges[model]

        if pretrained:
            exec("self.load_state_dict(models.%s(pretrained=False).state_dict())" % model)

        if not requires_grad:
            for param in super().parameters():
                param.requires_grad = False

        # 删除多余的全连接层参数，以节省内存。
        # 去掉vgg最后的全连接层(classifier)
        if remove_fc:
            del self.classifier

        if show_params:
            for name, param in self.named_parameters():
                print(name, param.size())

    def forward(self, x):
        output = {}  # 得到每个最大池化层的输出，VGG网络有5个最大池化层。
        for idx, (begin, end) in enumerate(self.ranges):
            # self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
            for layer in range(begin, end):
                x = self.features[layer](x)
            output["x%d" % (idx + 1)] = x
        return output


ranges = {
    'vgg11': ((0, 3), (3, 6), (6, 11), (11, 16), (16, 21)),
    'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
    'vgg16': ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)),
    'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37))
}

# Vgg网络结构配置
cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


# make layers using Vgg-Net config(cfg)
# 由cfg构建vgg-Net
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)


if __name__ == "__main__":
    pass

训练

'''
train.py
'''
from datetime import datetime
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from dataset import test_dataloader, train_dataloader
from model import FCNs, VGGNet

# 我将FCN文件名改成了model，所以是from model import...
# 同样将BagData文件名改成了dataset，所以是from dataset import...


"""
定义一个device变量，用于确定使用CPU还是GPU进行训练。
如果该变量的值为“cuda”，则表示使用GPU进行训练
如果该变量的值为“cpu”，则表示使用CPU进行训练
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


device=torch.device("cpu")      # 使用cpu
device=torch.device("cuda")     # 使用GPU


"""
定义一个名为train（）的函数，用于执行模型的训练过程。
该函数包含epo_num和show_vgg_params，分别表示训练的epoch数量和是否显示VGG16的参数信息。
函数内部，首先创建一个VGG16的实例vgg_model加载到指定设备上（即之前定义的device变量），
并且定义交叉熵损失函数和SGD优化器。
"""
def train(epo_num=50, show_vgg_params=False): # epoch_num表示训练的epoch数目，默认值为50
                                              # show_vgg_params表示是否显示VGG模型的参数
    # vis = visdom.Visdom()   #pytorch中的可视化工具

    # 创建VGGNet和FCNs模型并赋值给vgg_model和fcn_model
    vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params) # requires_grad=True表示模型的参数需要计算梯度用于反向传播
                                                                        # show_params=show_vgg_params表示是否需要打印VGGNet模型的参数信息
                                                                        # 当show_vgg_params=Ture时，会在控制台上打印VGGNet模型的参数信息
    fcn_model = FCNs(pretrained_net=vgg_model, n_class=2)               # pretrained_net=vgg_model表示使用预训练的VGGNet模型作为FCNs模型的预训练网络
                                                                        # n_class=2表示分类的类别数目为2（可能是黑白色二分类）

    # 将模型加载到指定设备上
    fcn_model = fcn_model.to(device)
    criterion = nn.BCELoss().to(device)
    optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7)

    # 定义两个空列表all_train_iter_loss和all_test_iter_loss用于存储训练接和测试集的损失值
    all_train_iter_loss = []
    all_test_iter_loss = []

    # 计算时间
    prev_time = datetime.now()

    """
    使用两个for循环，分别对训练数据和测试数据进行迭代，
    在迭代中，
    首先将数据和标签加载到指定设备上，
    然后将优化器的梯度清零，并对网络输出和标签进行前向传播和损失计算。
    在计算损失的过程中，代码还将网络输出进行了sigmoid操作，从而将输出值缩放到了0到1之间。
    在计算损失的过程之后，代码还将损失值记录到all_train_iter_loss或all_test_loss中
    并将损失值累加到total_loss中。
    最后，通过反向传播和优化器的更新，更新网络参数。
    """
    for epo in range(epo_num): # 循环遍历每个epoch的训练过程，epoch_num表示总的训练epoch数
        train_loss = 0
        fcn_model.train() # 将模式设置为训练模式，这是因为在训练模式下，模型会启用dropout等正则化
        # 循环遍历数据集中的每个batch
        for index, (bag, bag_msk) in enumerate(train_dataloader): # train_dataloader是一个Pytorch数据加载器，
                                                                  # 可以将训练数据集分成多个batch，并自动对每个batch进行shuffle和数据增强等操作
            # bag.shape is torch.Size([4, 3, 160, 160])
            # bag_msk.shape is torch.Size([4, 2, 160, 160])
            bag = bag.to(device)
            bag_msk = bag_msk.to(device)

            optimizer.zero_grad() # 将优化器的梯度缓存清零，以便计算新的梯度
            output = fcn_model(bag) # 使用FCNs模型对输入数据进行前向传播，得到模型的输出。
            output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
                                            # 将模型输出应用sigmoid函数，将输出值缩放到0到1之间，为了适应交叉熵损失函数的要求。
            # print(output)
            # print(bag_msk)
            loss = criterion(output, bag_msk) # 根据模型输出的标签数据计算损失函数值，这里使用的是交叉熵损失函数
            loss.backward() # 对损失函数进行反向传播，计算参数的梯度
            iter_loss = loss.item()
            all_train_iter_loss.append(iter_loss) # 将当前batch的损失函数值添加到all_train_iter_loss列表中
            train_loss += iter_loss # 累加每个batch的损失函数值，以便最后计算每个epoch的平均损失函数值
            optimizer.step() # 使用优化器更新模型参数

            output_np = output.cpu().detach().numpy().copy()  # output_np.shape = (4, 2, 160, 160)；将模型输出转移到CPU上，并转化为numpy数组
            output_np = np.argmin(output_np, axis=1) # 对输出数组沿着第一个轴求最小值的索引，得到二分类预测结果
            bag_msk_np = bag_msk.cpu().detach().numpy().copy()  # bag_msk_np.shape = (4, 2, 160, 160)；将标签数据移动到CPU上，并转换为numpy数组
            bag_msk_np = np.argmin(bag_msk_np, axis=1) # 对标签数组沿着第一个轴求最小值的索引，得到二分类结果

        test_loss = 0
        """
        fcn_model.eval()是将模型设置为评估模式，这个方法通常用于测试集上进行模型推理。
        在评估模式下，模型将会关闭一些具有随机性的操作，例如dropout、batch normalization等，以避免对模型推理的结果产生影响。
        此外，在评估模式下，模型的权重参数也不会被更新，因此可以减少内存的使用，提高模型推理的速度
        """
        fcn_model.eval() # 将模式设置为评估模式
        with torch.no_grad():
            for index, (bag, bag_msk) in enumerate(test_dataloader):
                bag = bag.to(device)
                bag_msk = bag_msk.to(device)
                optimizer.zero_grad()
                output = fcn_model(bag)
                output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
                loss = criterion(output, bag_msk)  # 预测和原标签图的差
                iter_loss = loss.item()  # item得到一个元素张量里面的元素值，一般用于返回loss，acc
                all_test_iter_loss.append(iter_loss)
                test_loss += iter_loss

                output_np = output.cpu().detach().numpy().copy()  # output_np.shape = (4, 2, 160, 160)
                output_np = np.argmin(output_np, axis=1)
                bag_msk_np = bag_msk.cpu().detach().numpy().copy()  # bag_msk_np.shape = (4, 2, 160, 160)
                bag_msk_np = np.argmin(bag_msk_np, axis=1)

        cur_time = datetime.now()
        h, remainder = divmod((cur_time - prev_time).seconds, 3600)
        m, s = divmod(remainder, 60)
        time_str = "Time %02d:%02d:%02d" % (h, m, s)
        prev_time = cur_time

        print('epoch train loss = %f, epoch test loss = %f, %s'
              % (train_loss / len(train_dataloader), test_loss / len(test_dataloader), time_str))

        """
        这段代码是训练过程中的一个保存模型的操作。
        具体来说，代码中使用了python中的取模运算符%（这里使用了np.mod函数）判断epoch是否为5的倍数，
        如果是，则将当前模型保存到文件中。文件名为“fun_model_{}.pt”,其中{}表示epoch的编号，
        因此每5个epoch保存一次模型，文件名的数字会对应增加
        
        备注：模型的保存可以在训练过程中定期执行，以便在训练中断或出现错误时，可以从已保存的模型中恢复训练
        从而不需要从头训练。此外，也可以使用已保存的模型来进行模型的推理或部署到生产环境中。
        保存模型的方法在Pytorch中有多种，包括使用torch.save()函数保存整个模型或仅保存模型的状态字典（state_dict）
        """
        if np.mod(epo+1, 5) == 0:
            torch.save(fcn_model, './fcn_model_{}.pt'.format(epo))
            print('fcn_model_{}.pt'.format(epo))


if __name__ == "__main__":
    train(epo_num=20, show_vgg_params=False)