<kaggle_start><data_title>VOC2012<data_description>Image dataset，only use by learning 
Image dataset，only use by learning 
Image dataset，only use by learning 
Image dataset，only use by learning 
Image dataset，only use by learning 
Image dataset，only use by learning 
Image dataset，only use by learning<data_name>voc2012
<code># **导入必要的包**
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import PIL.Image as pil_image
import matplotlib.pyplot as plt

# **加载自定义的数据******
# * https://zhuanlan.zhihu.com/p/32506912
# 分割的类别
classes = [
    "background",
    "aeroplane",
    "bicycle",
    "bird",
    "boat",
    "bottle",
    "bus",
    "car",
    "cat",
    "chair",
    "cow",
    "diningtable",
    "dog",
    "horse",
    "motorbike",
    "person",
    "potted plant",
    "sheep",
    "sofa",
    "train",
    "tv/monitor",
]
# 每个类别对应的RGB值
colormap = [
    [0, 0, 0],
    [128, 0, 0],
    [0, 128, 0],
    [128, 128, 0],
    [0, 0, 128],
    [128, 0, 128],
    [0, 128, 128],
    [128, 128, 128],
    [64, 0, 0],
    [192, 0, 0],
    [64, 128, 0],
    [192, 128, 0],
    [64, 0, 128],
    [192, 0, 128],
    [64, 128, 128],
    [192, 128, 128],
    [0, 64, 0],
    [128, 64, 0],
    [0, 192, 0],
    [128, 192, 0],
    [0, 64, 128],
]
cm2lbl = np.zeros(256**3)  # 每个像素点有 0 ~ 255 的选择，RGB 三个通道
for i, cm in enumerate(colormap):
    cm2lbl[(cm[0] * 256 + cm[1]) * 256 + cm[2]] = i  # 建立索引


# 将图片映射成索引数组
def image_to_label(image):
    """将图片映射成类别索引的数组"""
    data = np.array(image, dtype="int32")
    # 按照上面一样的计算规则，得到对应的值
    index = (data[:, :, 0] * 256 + data[:, :, 1]) * 256 + data[:, :, 2]
    return np.array(cm2lbl[index], dtype="int64")


def random_crop(image, width, height):
    """随机裁剪"""
    pass


# 测试
image = pil_image.open(
    "../input/voc2012/data/VOCdevkit/VOC2012/SegmentationClass/2007_000032.png"
).convert("RGB")


# image = transforms.RandomCrop((224, 224))(image)
# print(image)
# plt.imshow(image)
# label = transforms.FixedCrop(*rect)(label)
# image_array=image_to_label(image)
class VOCDataset(Dataset):
    """自定义数据类加载规则"""

    def __init__(self, file_path=None, transform=None):
        """初始化函数"""
        images_labels = []
        file = open(file_path)
        for name in file.readlines():
            # 移除空格和换行符
            name = name.strip()
            image = (
                "../input/voc2012/data/VOCdevkit/VOC2012/JPEGImages/" + name + ".jpg"
            )
            label = (
                "../input/voc2012/data/VOCdevkit/VOC2012/SegmentationClass/"
                + name
                + ".png"
            )
            images_labels.append((image, label))
        self.images_labels = images_labels
        self.transform = transform

    def __getitem__(self, index):
        """在DataLoader中会调用这个方法读取一个batch的数据"""
        image_path, label_path = self.images_labels[index]
        # 使用image.open加载目标图和特征图
        image = pil_image.open(image_path)
        label = pil_image.open(label_path).convert("RGB")
        #         # 裁剪图片，使其所有的图片输入一致
        #         x,y,width,height=transforms.RandomCrop.get_params(img=image,output_size=(224,224))
        #         image=function.crop(image,x,y,width,height)
        #         label=function.crop(label,x,y,width,height)
        image = transforms.Resize((512, 512))(image)
        label = transforms.Resize((512, 512))(label)
        # 转化特征图
        if self.transform is not None:
            image = self.transform(image)
        # 映射目标图
        label = image_to_label(label)
        # 从numpy数组转化成张量
        label = torch.from_numpy(label)
        # 返回
        return image, label

    def __len__(self):
        """获取整个dataset的数据大小"""
        return len(self.images_labels)


# 数据预处理，增强，归一化
transform_train = transforms.Compose(
    [
        # 将数据转化成张量，并且归一化到[0,1]
        transforms.ToTensor(),
        # 将数据标准化到[-1,1]，image=(image-mean)/std
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
transform_test = transforms.Compose(
    [
        # 将数据转化成张量，并且归一化到[0,1]
        transforms.ToTensor(),
        # 将数据标准化到[-1,1]，image=(image-mean)/std
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
train_datasets = VOCDataset(
    file_path="../input/voc2012/data/VOCdevkit/VOC2012/ImageSets/Segmentation/train.txt",
    transform=transform_train,
)
test_datasets = VOCDataset(
    file_path="../input/voc2012/data/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt",
    transform=transform_test,
)
train_loader = DataLoader(
    dataset=train_datasets, batch_size=8, shuffle=False, sampler=None
)
test_loader = DataLoader(
    dataset=test_datasets, batch_size=8, shuffle=False, sampler=None
)
print(len(train_loader))
# print(next(iter(train_loader)))


# **FCN网络结构如下**
# 前面由VGG16组成，分5个卷积层，后面有3个（1，1）卷积层，输出32s，16s，8s，然后在进行融合，上采样得出结果
# https://github.com/wkentaro/pytorch-fcn/blob/main/torchfcn/models/fcn8s.py
# ![](https://gimg2.baidu.com/image_search/src=http%3A%2F%2Fwww.itdaan.com%2Fi%2Fd91a35b4436eba676dc2ee6cc6e45276.jpg&refer=http%3A%2F%2Fwww.itdaan.com&app=2002&size=f9999,10000&q=a80&n=0&g=0n&fmt=auto?sec=1653238248&t=2c48611d3c312d0356b18fff8f203e7f)
# **构建FCN网络模型**
class FCN8s(nn.Module):
    def __init__(self):
        super(FCN8s, self).__init__()
        # 本项目有20个类别，一个背景，一共21类
        n_class = 21
        # conv1
        # 输入图像为3通道，输出64个特征图，卷积核大小为（3，3），步长为1，padding为100（避免图片不兼容，其实也可以为1的）
        # 卷积输出公式：output=(input+2*padding-kernel_size)/stride+1
        #  512=(512+2*1-3)/1+1
        self.conv1_1 = nn.Conv2d(
            in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1
        )
        self.bn1_1 = nn.BatchNorm2d(num_features=64)
        self.relu1_1 = nn.ReLU(inplace=True)
        self.conv1_2 = nn.Conv2d(
            in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1
        )
        self.bn1_2 = nn.BatchNorm2d(num_features=64)
        self.relu1_2 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        # 采样输出公式：output=(input+2*padding-kernel_size)/stride+1
        # 256=(512+2*0-2)/2+1
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv2
        self.conv2_1 = nn.Conv2d(
            in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1
        )
        self.bn2_1 = nn.BatchNorm2d(num_features=128)
        self.relu2_1 = nn.ReLU(inplace=True)
        self.conv2_2 = nn.Conv2d(
            in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1
        )
        self.bn2_2 = nn.BatchNorm2d(num_features=128)
        self.relu2_2 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv3
        self.conv3_1 = nn.Conv2d(
            in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_1 = nn.BatchNorm2d(num_features=256)
        self.relu3_1 = nn.ReLU(inplace=True)
        self.conv3_2 = nn.Conv2d(
            in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_2 = nn.BatchNorm2d(num_features=256)
        self.relu3_2 = nn.ReLU(inplace=True)
        self.conv3_3 = nn.Conv2d(
            in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_3 = nn.BatchNorm2d(num_features=256)
        self.relu3_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv4
        self.conv4_1 = nn.Conv2d(
            in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_1 = nn.BatchNorm2d(num_features=512)
        self.relu4_1 = nn.ReLU(inplace=True)
        self.conv4_2 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_2 = nn.BatchNorm2d(num_features=512)
        self.relu4_2 = nn.ReLU(inplace=True)
        self.conv4_3 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_3 = nn.BatchNorm2d(num_features=512)
        self.relu4_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv5
        # 输入图像为3通道，输出64个特征图，卷积核大小为（3，3），步长为1，padding为100（避免图片不兼容）
        self.conv5_1 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_1 = nn.BatchNorm2d(num_features=512)
        self.relu5_1 = nn.ReLU(inplace=True)
        self.conv5_2 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_2 = nn.BatchNorm2d(num_features=512)
        self.relu5_2 = nn.ReLU(inplace=True)
        self.conv5_3 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_3 = nn.BatchNorm2d(num_features=512)
        self.relu5_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool5 = nn.MaxPool2d(kernel_size=2, stride=2)
        # cnov6
        self.conv6 = nn.Conv2d(
            in_channels=512, out_channels=4096, kernel_size=7, stride=1, padding=1
        )
        self.bn6 = nn.BatchNorm2d(num_features=4096)
        self.relu6 = nn.ReLU(inplace=True)
        self.drop6 = nn.Dropout2d(p=0.5)
        # cnov7
        self.conv7 = nn.Conv2d(
            in_channels=4096, out_channels=4096, kernel_size=1, stride=1, padding=1
        )
        self.bn7 = nn.BatchNorm2d(num_features=4096)
        self.relu7 = nn.ReLU(inplace=True)
        self.drop7 = nn.Dropout2d(p=0.5)
        # cnov8，本项目有20个类别，一个背景，一共21类
        self.conv8 = nn.Conv2d(
            in_channels=4096, out_channels=n_class, kernel_size=1, stride=1, padding=1
        )
        # 上采样2倍（16，16，21）————>（32，32，21）
        self.up_conv8_2 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=2,
            stride=2,
            bias=False,
        )
        # 反卷积ConvTranspose2d操作输出宽高公式
        # output=((input-1)*stride)+outputpadding-(2*padding)+kernelsize
        # 34=(16-1)*2+0-(2*0)+4
        # 第4层maxpool值做卷积运算
        self.pool4_conv = nn.Conv2d(
            in_channels=512, out_channels=n_class, kernel_size=1, stride=1
        )
        # 利用反卷积上采样2倍
        self.up_pool4_2 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=2,
            stride=2,
            bias=False,
        )
        # 第3层maxpool值做卷积运算
        self.pool3_conv = nn.Conv2d(
            in_channels=256, out_channels=n_class, kernel_size=1, stride=1
        )
        # 利用反卷积上采样8倍
        self.up_pool3_8 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=8,
            stride=8,
            bias=False,
        )

    def forward(self, x):
        """正向传播"""
        # 记录初始图片的大小（32，21，512，512）
        h = x
        # conv1
        x = self.relu1_1(self.bn1_1(self.conv1_1(x)))
        x = self.relu1_2(self.bn1_2(self.conv1_2(x)))
        x = self.maxpool1(x)
        # conv2
        x = self.relu2_1(self.bn2_1(self.conv2_1(x)))
        x = self.relu2_2(self.bn2_2(self.conv2_2(x)))
        x = self.maxpool2(x)
        # conv3
        x = self.relu3_1(self.bn3_1(self.conv3_1(x)))
        x = self.relu3_2(self.bn3_2(self.conv3_2(x)))
        x = self.relu3_3(self.bn3_3(self.conv3_3(x)))
        x = self.maxpool3(x)
        pool3 = x
        # conv4
        x = self.relu4_1(self.bn4_1(self.conv4_1(x)))
        x = self.relu4_2(self.bn4_2(self.conv4_2(x)))
        x = self.relu4_3(self.bn4_3(self.conv4_3(x)))
        x = self.maxpool4(x)
        pool4 = x
        # conv5
        x = self.relu5_1(self.bn5_1(self.conv5_1(x)))
        x = self.relu5_2(self.bn5_2(self.conv5_2(x)))
        x = self.relu5_3(self.bn5_3(self.conv5_3(x)))
        x = self.maxpool5(x)
        # conv6
        #         print(self.conv6(x).shape)
        #         print(self.bn6(self.conv6(x)).shape)
        #         print(self.relu6(self.bn6(self.conv6(x))).shape)
        #         print(self.drop6(self.relu6(self.bn6(self.conv6(x)))).shape)
        x = self.drop6(self.relu6(self.bn6(self.conv6(x))))
        # conv7
        x = self.drop7(self.relu7(self.bn7(self.conv7(x))))
        # conv8
        x = self.up_conv8_2(self.conv8(x))
        up_conv8 = x
        # 计算第4层的值
        x2 = self.pool4_conv(pool4)
        # 相加融合
        x2 = up_conv8 + x2
        # 反卷积上采样8倍
        x2 = self.up_pool4_2(x2)
        up_pool4 = x2
        # 计算第3层的值
        x3 = self.pool3_conv(pool3)
        x3 = up_pool4 + x3
        # 反卷积上采样8倍
        x3 = self.up_pool3_8(x3)
        return x3


model = FCN8s()
# print(model)


# **模型训练**
# 配置训练参数
# 选择设置，优先GPU
device = torch.device("cuda" if torch.cuda.is_available else "cpu")
# 训练次数
epochs = 100
# 评判标准，损失函数
criterion = torch.nn.CrossEntropyLoss()
# 优化方法
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 将模型赋值到GPU中
model = model.to(device)
for epoch in range(epochs):
    num = 1
    loss_num = 0
    for i, (image, label) in enumerate(train_loader):
        # 将数据复制到GPU中
        image = image.to(device)
        label = label.to(device)
        # 梯度清零
        optimizer.zero_grad()
        # 正向传播
        output = model(image)
        # 计算损失
        loss = criterion(output, label)
        # 反向传播
        loss.backward()
        # 更新梯度
        optimizer.step()
        # 获取当前损失
        running_loss = loss.data.item()
        loss_num += running_loss
        num += i
    epoch_loss = loss_num / len(train_loader)
    print("次数", epoch, "平均损失", epoch_loss)
# 保存模型
torch.save(model.state_dict(), "./fcn8s.pt")

# **模型预测**
# 加载模型
# model=torch.load("./fcn8s.pt")
# 进入评估模式
model.eval()
cm = np.array(colormap).astype("uint8")
# 测试集梯度不更新
with torch.no_grad():
    for image, label in test_loader:
        # 将数据复制到GPU中
        image = image.to(device)
        label = label.to(device)
        print("image", image.shape, "label", label.shape)
        # 正向传播
        output = model(image)
        # 把数据从GPU复制到CPU中，plt才能调用
        output = output.max(1)[1].squeeze().cpu().data.numpy()
        pred = cm[output]
        plt.imshow(pred[0])
        #         output=output.cpu().numpy()
        #         print(output.shape)
        #         for i,eval_image in enumerate(output):
        #             plt.imshow(eval_image)
        break
<|endoftext|>
[{"Id": 94772222, "ScriptId": 26481417, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 3411147, "CreationDate": "05/04/2022 18:53:08", "VersionNumber": 7.0, "Title": "[\u56fe\u50cf\u5206\u5272]Pytorch-VGG/FCN-\u8bed\u4e49\u5206\u5272", "EvaluationDate": "05/04/2022", "IsChange": true, "TotalLines": 416.0, "LinesInsertedFromPrevious": 22.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 394.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 0}]
[{"Id": 129113281, "KernelVersionId": 94772222, "SourceDatasetVersionId": 2887798}]
[{"Id": 2887798, "DatasetId": 1768994, "DatasourceVersionId": 2934854, "CreatorUserId": 7951664, "LicenseName": "U.S. Government Works", "CreationDate": "12/05/2021 02:55:48", "VersionNumber": 1.0, "Title": "VOC2012", "Slug": "voc2012", "Subtitle": "Image sets Image sets Image sets Image sets Image sets Image sets", "Description": "Image dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning \nImage dataset\uff0conly use by learning", "VersionNotes": "Initial release", "TotalCompressedBytes": 0.0, "TotalUncompressedBytes": 0.0}]
[{"Id": 1768994, "CreatorUserId": 7951664, "OwnerUserId": 7951664.0, "OwnerOrganizationId": NaN, "CurrentDatasetVersionId": 2887798.0, "CurrentDatasourceVersionId": 2934854.0, "ForumId": 1791237, "Type": 2, "CreationDate": "12/05/2021 02:55:48", "LastActivityDate": "12/05/2021", "TotalViews": 2390, "TotalDownloads": 235, "TotalVotes": 4, "TotalKernels": 1}]
# **导入必要的包**
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import PIL.Image as pil_image
import matplotlib.pyplot as plt

# **加载自定义的数据******
# * https://zhuanlan.zhihu.com/p/32506912
# 分割的类别
classes = [
    "background",
    "aeroplane",
    "bicycle",
    "bird",
    "boat",
    "bottle",
    "bus",
    "car",
    "cat",
    "chair",
    "cow",
    "diningtable",
    "dog",
    "horse",
    "motorbike",
    "person",
    "potted plant",
    "sheep",
    "sofa",
    "train",
    "tv/monitor",
]
# 每个类别对应的RGB值
colormap = [
    [0, 0, 0],
    [128, 0, 0],
    [0, 128, 0],
    [128, 128, 0],
    [0, 0, 128],
    [128, 0, 128],
    [0, 128, 128],
    [128, 128, 128],
    [64, 0, 0],
    [192, 0, 0],
    [64, 128, 0],
    [192, 128, 0],
    [64, 0, 128],
    [192, 0, 128],
    [64, 128, 128],
    [192, 128, 128],
    [0, 64, 0],
    [128, 64, 0],
    [0, 192, 0],
    [128, 192, 0],
    [0, 64, 128],
]
cm2lbl = np.zeros(256**3)  # 每个像素点有 0 ~ 255 的选择，RGB 三个通道
for i, cm in enumerate(colormap):
    cm2lbl[(cm[0] * 256 + cm[1]) * 256 + cm[2]] = i  # 建立索引


# 将图片映射成索引数组
def image_to_label(image):
    """将图片映射成类别索引的数组"""
    data = np.array(image, dtype="int32")
    # 按照上面一样的计算规则，得到对应的值
    index = (data[:, :, 0] * 256 + data[:, :, 1]) * 256 + data[:, :, 2]
    return np.array(cm2lbl[index], dtype="int64")


def random_crop(image, width, height):
    """随机裁剪"""
    pass


# 测试
image = pil_image.open(
    "../input/voc2012/data/VOCdevkit/VOC2012/SegmentationClass/2007_000032.png"
).convert("RGB")


# image = transforms.RandomCrop((224, 224))(image)
# print(image)
# plt.imshow(image)
# label = transforms.FixedCrop(*rect)(label)
# image_array=image_to_label(image)
class VOCDataset(Dataset):
    """自定义数据类加载规则"""

    def __init__(self, file_path=None, transform=None):
        """初始化函数"""
        images_labels = []
        file = open(file_path)
        for name in file.readlines():
            # 移除空格和换行符
            name = name.strip()
            image = (
                "../input/voc2012/data/VOCdevkit/VOC2012/JPEGImages/" + name + ".jpg"
            )
            label = (
                "../input/voc2012/data/VOCdevkit/VOC2012/SegmentationClass/"
                + name
                + ".png"
            )
            images_labels.append((image, label))
        self.images_labels = images_labels
        self.transform = transform

    def __getitem__(self, index):
        """在DataLoader中会调用这个方法读取一个batch的数据"""
        image_path, label_path = self.images_labels[index]
        # 使用image.open加载目标图和特征图
        image = pil_image.open(image_path)
        label = pil_image.open(label_path).convert("RGB")
        #         # 裁剪图片，使其所有的图片输入一致
        #         x,y,width,height=transforms.RandomCrop.get_params(img=image,output_size=(224,224))
        #         image=function.crop(image,x,y,width,height)
        #         label=function.crop(label,x,y,width,height)
        image = transforms.Resize((512, 512))(image)
        label = transforms.Resize((512, 512))(label)
        # 转化特征图
        if self.transform is not None:
            image = self.transform(image)
        # 映射目标图
        label = image_to_label(label)
        # 从numpy数组转化成张量
        label = torch.from_numpy(label)
        # 返回
        return image, label

    def __len__(self):
        """获取整个dataset的数据大小"""
        return len(self.images_labels)


# 数据预处理，增强，归一化
transform_train = transforms.Compose(
    [
        # 将数据转化成张量，并且归一化到[0,1]
        transforms.ToTensor(),
        # 将数据标准化到[-1,1]，image=(image-mean)/std
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
transform_test = transforms.Compose(
    [
        # 将数据转化成张量，并且归一化到[0,1]
        transforms.ToTensor(),
        # 将数据标准化到[-1,1]，image=(image-mean)/std
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
train_datasets = VOCDataset(
    file_path="../input/voc2012/data/VOCdevkit/VOC2012/ImageSets/Segmentation/train.txt",
    transform=transform_train,
)
test_datasets = VOCDataset(
    file_path="../input/voc2012/data/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt",
    transform=transform_test,
)
train_loader = DataLoader(
    dataset=train_datasets, batch_size=8, shuffle=False, sampler=None
)
test_loader = DataLoader(
    dataset=test_datasets, batch_size=8, shuffle=False, sampler=None
)
print(len(train_loader))
# print(next(iter(train_loader)))


# **FCN网络结构如下**
# 前面由VGG16组成，分5个卷积层，后面有3个（1，1）卷积层，输出32s，16s，8s，然后在进行融合，上采样得出结果
# https://github.com/wkentaro/pytorch-fcn/blob/main/torchfcn/models/fcn8s.py
# ![](https://gimg2.baidu.com/image_search/src=http%3A%2F%2Fwww.itdaan.com%2Fi%2Fd91a35b4436eba676dc2ee6cc6e45276.jpg&refer=http%3A%2F%2Fwww.itdaan.com&app=2002&size=f9999,10000&q=a80&n=0&g=0n&fmt=auto?sec=1653238248&t=2c48611d3c312d0356b18fff8f203e7f)
# **构建FCN网络模型**
class FCN8s(nn.Module):
    def __init__(self):
        super(FCN8s, self).__init__()
        # 本项目有20个类别，一个背景，一共21类
        n_class = 21
        # conv1
        # 输入图像为3通道，输出64个特征图，卷积核大小为（3，3），步长为1，padding为100（避免图片不兼容，其实也可以为1的）
        # 卷积输出公式：output=(input+2*padding-kernel_size)/stride+1
        #  512=(512+2*1-3)/1+1
        self.conv1_1 = nn.Conv2d(
            in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1
        )
        self.bn1_1 = nn.BatchNorm2d(num_features=64)
        self.relu1_1 = nn.ReLU(inplace=True)
        self.conv1_2 = nn.Conv2d(
            in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1
        )
        self.bn1_2 = nn.BatchNorm2d(num_features=64)
        self.relu1_2 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        # 采样输出公式：output=(input+2*padding-kernel_size)/stride+1
        # 256=(512+2*0-2)/2+1
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv2
        self.conv2_1 = nn.Conv2d(
            in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1
        )
        self.bn2_1 = nn.BatchNorm2d(num_features=128)
        self.relu2_1 = nn.ReLU(inplace=True)
        self.conv2_2 = nn.Conv2d(
            in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1
        )
        self.bn2_2 = nn.BatchNorm2d(num_features=128)
        self.relu2_2 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv3
        self.conv3_1 = nn.Conv2d(
            in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_1 = nn.BatchNorm2d(num_features=256)
        self.relu3_1 = nn.ReLU(inplace=True)
        self.conv3_2 = nn.Conv2d(
            in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_2 = nn.BatchNorm2d(num_features=256)
        self.relu3_2 = nn.ReLU(inplace=True)
        self.conv3_3 = nn.Conv2d(
            in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1
        )
        self.bn3_3 = nn.BatchNorm2d(num_features=256)
        self.relu3_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv4
        self.conv4_1 = nn.Conv2d(
            in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_1 = nn.BatchNorm2d(num_features=512)
        self.relu4_1 = nn.ReLU(inplace=True)
        self.conv4_2 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_2 = nn.BatchNorm2d(num_features=512)
        self.relu4_2 = nn.ReLU(inplace=True)
        self.conv4_3 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn4_3 = nn.BatchNorm2d(num_features=512)
        self.relu4_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # conv5
        # 输入图像为3通道，输出64个特征图，卷积核大小为（3，3），步长为1，padding为100（避免图片不兼容）
        self.conv5_1 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_1 = nn.BatchNorm2d(num_features=512)
        self.relu5_1 = nn.ReLU(inplace=True)
        self.conv5_2 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_2 = nn.BatchNorm2d(num_features=512)
        self.relu5_2 = nn.ReLU(inplace=True)
        self.conv5_3 = nn.Conv2d(
            in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1
        )
        self.bn5_3 = nn.BatchNorm2d(num_features=512)
        self.relu5_3 = nn.ReLU(inplace=True)
        # 最大池化层进行下采样
        self.maxpool5 = nn.MaxPool2d(kernel_size=2, stride=2)
        # cnov6
        self.conv6 = nn.Conv2d(
            in_channels=512, out_channels=4096, kernel_size=7, stride=1, padding=1
        )
        self.bn6 = nn.BatchNorm2d(num_features=4096)
        self.relu6 = nn.ReLU(inplace=True)
        self.drop6 = nn.Dropout2d(p=0.5)
        # cnov7
        self.conv7 = nn.Conv2d(
            in_channels=4096, out_channels=4096, kernel_size=1, stride=1, padding=1
        )
        self.bn7 = nn.BatchNorm2d(num_features=4096)
        self.relu7 = nn.ReLU(inplace=True)
        self.drop7 = nn.Dropout2d(p=0.5)
        # cnov8，本项目有20个类别，一个背景，一共21类
        self.conv8 = nn.Conv2d(
            in_channels=4096, out_channels=n_class, kernel_size=1, stride=1, padding=1
        )
        # 上采样2倍（16，16，21）————>（32，32，21）
        self.up_conv8_2 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=2,
            stride=2,
            bias=False,
        )
        # 反卷积ConvTranspose2d操作输出宽高公式
        # output=((input-1)*stride)+outputpadding-(2*padding)+kernelsize
        # 34=(16-1)*2+0-(2*0)+4
        # 第4层maxpool值做卷积运算
        self.pool4_conv = nn.Conv2d(
            in_channels=512, out_channels=n_class, kernel_size=1, stride=1
        )
        # 利用反卷积上采样2倍
        self.up_pool4_2 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=2,
            stride=2,
            bias=False,
        )
        # 第3层maxpool值做卷积运算
        self.pool3_conv = nn.Conv2d(
            in_channels=256, out_channels=n_class, kernel_size=1, stride=1
        )
        # 利用反卷积上采样8倍
        self.up_pool3_8 = nn.ConvTranspose2d(
            in_channels=n_class,
            out_channels=n_class,
            kernel_size=8,
            stride=8,
            bias=False,
        )

    def forward(self, x):
        """正向传播"""
        # 记录初始图片的大小（32，21，512，512）
        h = x
        # conv1
        x = self.relu1_1(self.bn1_1(self.conv1_1(x)))
        x = self.relu1_2(self.bn1_2(self.conv1_2(x)))
        x = self.maxpool1(x)
        # conv2
        x = self.relu2_1(self.bn2_1(self.conv2_1(x)))
        x = self.relu2_2(self.bn2_2(self.conv2_2(x)))
        x = self.maxpool2(x)
        # conv3
        x = self.relu3_1(self.bn3_1(self.conv3_1(x)))
        x = self.relu3_2(self.bn3_2(self.conv3_2(x)))
        x = self.relu3_3(self.bn3_3(self.conv3_3(x)))
        x = self.maxpool3(x)
        pool3 = x
        # conv4
        x = self.relu4_1(self.bn4_1(self.conv4_1(x)))
        x = self.relu4_2(self.bn4_2(self.conv4_2(x)))
        x = self.relu4_3(self.bn4_3(self.conv4_3(x)))
        x = self.maxpool4(x)
        pool4 = x
        # conv5
        x = self.relu5_1(self.bn5_1(self.conv5_1(x)))
        x = self.relu5_2(self.bn5_2(self.conv5_2(x)))
        x = self.relu5_3(self.bn5_3(self.conv5_3(x)))
        x = self.maxpool5(x)
        # conv6
        #         print(self.conv6(x).shape)
        #         print(self.bn6(self.conv6(x)).shape)
        #         print(self.relu6(self.bn6(self.conv6(x))).shape)
        #         print(self.drop6(self.relu6(self.bn6(self.conv6(x)))).shape)
        x = self.drop6(self.relu6(self.bn6(self.conv6(x))))
        # conv7
        x = self.drop7(self.relu7(self.bn7(self.conv7(x))))
        # conv8
        x = self.up_conv8_2(self.conv8(x))
        up_conv8 = x
        # 计算第4层的值
        x2 = self.pool4_conv(pool4)
        # 相加融合
        x2 = up_conv8 + x2
        # 反卷积上采样8倍
        x2 = self.up_pool4_2(x2)
        up_pool4 = x2
        # 计算第3层的值
        x3 = self.pool3_conv(pool3)
        x3 = up_pool4 + x3
        # 反卷积上采样8倍
        x3 = self.up_pool3_8(x3)
        return x3


model = FCN8s()
# print(model)


# **模型训练**
# 配置训练参数
# 选择设置，优先GPU
device = torch.device("cuda" if torch.cuda.is_available else "cpu")
# 训练次数
epochs = 100
# 评判标准，损失函数
criterion = torch.nn.CrossEntropyLoss()
# 优化方法
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 将模型赋值到GPU中
model = model.to(device)
for epoch in range(epochs):
    num = 1
    loss_num = 0
    for i, (image, label) in enumerate(train_loader):
        # 将数据复制到GPU中
        image = image.to(device)
        label = label.to(device)
        # 梯度清零
        optimizer.zero_grad()
        # 正向传播
        output = model(image)
        # 计算损失
        loss = criterion(output, label)
        # 反向传播
        loss.backward()
        # 更新梯度
        optimizer.step()
        # 获取当前损失
        running_loss = loss.data.item()
        loss_num += running_loss
        num += i
    epoch_loss = loss_num / len(train_loader)
    print("次数", epoch, "平均损失", epoch_loss)
# 保存模型
torch.save(model.state_dict(), "./fcn8s.pt")

# **模型预测**
# 加载模型
# model=torch.load("./fcn8s.pt")
# 进入评估模式
model.eval()
cm = np.array(colormap).astype("uint8")
# 测试集梯度不更新
with torch.no_grad():
    for image, label in test_loader:
        # 将数据复制到GPU中
        image = image.to(device)
        label = label.to(device)
        print("image", image.shape, "label", label.shape)
        # 正向传播
        output = model(image)
        # 把数据从GPU复制到CPU中，plt才能调用
        output = output.max(1)[1].squeeze().cpu().data.numpy()
        pred = cm[output]
        plt.imshow(pred[0])
        #         output=output.cpu().numpy()
        #         print(output.shape)
        #         for i,eval_image in enumerate(output):
        #             plt.imshow(eval_image)
        break
