【语义分割专栏】3：Segnet实战篇(附上完整可运行的代码pytorch)

目录

• 前言
  
• Segnet全流程代码
  
  
  
  • 模型搭建(model)
    
  • 数据处理(dataloader)
    
  • 评价指标(metric)
    
  • 训练流程(train)
    
  • 模型测试(test)
    
  
  
• 效果图
  
• 结语
  

前言

Segnet原理篇讲解：【语义分割专栏】3：Segnet原理篇 - carpell - 博客园
代码地址，下载可复现：fouen6/Segnet_semantic-segmentation: 用于学习理解segnet原理

本篇文章收录于语义分割专栏，如果对语义分割领域感兴趣的，可以去看看专栏，会对经典的模型以及代码进行详细的讲解哦！其中会包含可复现的代码！(数据集文中提供了下载地址，下载不到可在评论区要取)
上篇文章已经带大家学习过了Segnet的原理，相信大家对于原理应该有了比较深的了解。本文将会带大家去手动复现属于自己的一个语义分割模型。将会深入代码进行讲解，如果有讲错的地方欢迎大家批评指正！
其实所有的深度学习模型的搭建我认为可以总结成五部分：模型的构建，数据集的处理，评价指标的设定，训练流程，测试。其实感觉有点深度学习代码八股文的那种意思。本篇同样的也会按照这样的方式进行讲解，希望大家能够深入代码去进行了解学习。
请记住：只懂原理不懂代码，你就算有了很好的想法创新点，你也难以去实现，所以希望大家能够深入去了解，最好能够参考着本文自己复现一下。

Segnet全流程代码

模型搭建(model)

首先是我们的crop函数，为什么需要用到这个，因为在测试的时候，我们不会对图像进行resize操作的，所以其就不一定是32的倍数，在下采样的过程中可能会出现从45->22的情况，但是上采样过程中就会变成22->44，这样就会造成shape的不匹配，所以需要对齐两者的shape大小。

1. def crop(upsampled, bypass):
3.     h1, w1 = upsampled.shape[2], upsampled.shape[3]
4.     h2, w2 = bypass.shape[2], bypass.shape[3]
6.     # 计算差值
7.     deltah = h2 - h1
8.     deltaw = w2 - w1
10.     # 计算填充的起始和结束位置
11.     # 对于高度
12.     pad_top = deltah // 2
13.     pad_bottom = deltah - pad_top
14.     # 对于宽度
15.     pad_left = deltaw // 2
16.     pad_right = deltaw - pad_left
18.     # 对 upsampled 进行中心填充
19.     upsampled_padded = F.pad(upsampled, (pad_left, pad_right, pad_top, pad_bottom), "constant", 0)
21.     return upsampled_padded

复制代码

然后就是我们的Segnet模型代码了。其实还是非常好理解的，其编码器的结构就是VGG的结构，只不过其在maxpooling的时候需要保存索引，然后就是解码器的结构，其实就是对编码器做个对称就行了。写好模型参数之后，非常重要的，记得要进行参数的初始化哈，这样能够利于之后的训练过程。

1. class SegNet(nn.Module):
2.     def __init__(self,num_classes=12):
3.         super(SegNet, self).__init__()
4.         self.encoder1 = nn.Sequential(
5.             nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1),
6.             nn.BatchNorm2d(64),
7.             nn.ReLU(),
8.             nn.Conv2d(64,64,kernel_size=3,stride=1,padding=1),
9.             nn.BatchNorm2d(64),
10.             nn.ReLU(),
11.         )
12.         self.encoder2 = nn.Sequential(
13.             nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),
14.             nn.BatchNorm2d(128),
15.             nn.ReLU(),
16.             nn.Conv2d(128,128,kernel_size=3,stride=1,padding=1),
17.             nn.BatchNorm2d(128),
18.             nn.ReLU(),
19.         )
20.         self.encoder3 = nn.Sequential(
21.             nn.Conv2d(128,256,kernel_size=3,stride=1,padding=1),
22.             nn.BatchNorm2d(256),
23.             nn.ReLU(),
24.             nn.Conv2d(256,256,kernel_size=3,stride=1,padding=1),
25.             nn.BatchNorm2d(256),
26.             nn.ReLU(),
27.             nn.Conv2d(256,256,kernel_size=3,stride=1,padding=1),
28.             nn.BatchNorm2d(256),
29.             nn.ReLU(),
30.         )
31.         self.encoder4 = nn.Sequential(
32.             nn.Conv2d(256,512,kernel_size=3,stride=1,padding=1),
33.             nn.BatchNorm2d(512),
34.             nn.ReLU(),
35.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
36.             nn.BatchNorm2d(512),
37.             nn.ReLU(),
38.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
39.             nn.BatchNorm2d(512),
40.             nn.ReLU(),
41.         )
42.         self.encoder5 = nn.Sequential(
43.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
44.             nn.BatchNorm2d(512),
45.             nn.ReLU(),
46.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
47.             nn.BatchNorm2d(512),
48.             nn.ReLU(),
49.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
50.             nn.BatchNorm2d(512),
51.             nn.ReLU(),
52.         )
54.         self.decoder1 = nn.Sequential(
55.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
56.             nn.BatchNorm2d(512),
57.             nn.ReLU(),
58.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
59.             nn.BatchNorm2d(512),
60.             nn.ReLU(),
61.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
62.             nn.BatchNorm2d(512),
63.             nn.ReLU(),
64.         )
65.         self.decoder2 = nn.Sequential(
66.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
67.             nn.BatchNorm2d(512),
68.             nn.ReLU(),
69.             nn.Conv2d(512,512,kernel_size=3,stride=1,padding=1),
70.             nn.BatchNorm2d(512),
71.             nn.ReLU(),
72.             nn.Conv2d(512,256,kernel_size=3,stride=1,padding=1),
73.             nn.BatchNorm2d(256),
74.             nn.ReLU(),
75.         )
76.         self.decoder3 = nn.Sequential(
77.             nn.Conv2d(256,256,kernel_size=3,stride=1,padding=1),
78.             nn.BatchNorm2d(256),
79.             nn.ReLU(),
80.             nn.Conv2d(256,256,kernel_size=3,stride=1,padding=1),
81.             nn.BatchNorm2d(256),
82.             nn.ReLU(),
83.             nn.Conv2d(256,128,kernel_size=3,stride=1,padding=1),
84.             nn.BatchNorm2d(128),
85.             nn.ReLU(),
86.         )
87.         self.decoder4 = nn.Sequential(
88.             nn.Conv2d(128,128,kernel_size=3,stride=1,padding=1),
89.             nn.BatchNorm2d(128),
90.             nn.ReLU(),
91.             nn.Conv2d(128,64,kernel_size=3,stride=1,padding=1),
92.             nn.BatchNorm2d(64),
93.             nn.ReLU(),
94.         )
95.         self.decoder5 = nn.Sequential(
96.             nn.Conv2d(64,64,kernel_size=3,stride=1,padding=1),
97.             nn.BatchNorm2d(64),
98.             nn.ReLU(),
99.             nn.Conv2d(64,64,kernel_size=3,stride=1,padding=1),
100.             nn.BatchNorm2d(64),
101.             nn.ReLU(),
102.             nn.Conv2d(64,num_classes,kernel_size=1),
103.         )
105.         self.max_pool = nn.MaxPool2d(2,2,return_indices=True)
106.         self.max_uppool = nn.MaxUnpool2d(2,2)
108.         self.initialize_weights()
110.     def initialize_weights(self):
111.         for m in self.modules():
112.             if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
113.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
114.                 if m.bias is not None:
115.                     nn.init.constant_(m.bias, 0)
116.             elif isinstance(m, nn.BatchNorm2d):
117.                 nn.init.constant_(m.weight, 1)
118.                 nn.init.constant_(m.bias, 0)
120.     def forward(self, x):
121.         x1 = self.encoder1(x)
122.         x,pool_indices1 = self.max_pool(x1)
123.         x2 = self.encoder2(x)
124.         x,pool_indices2 = self.max_pool(x2)
125.         x3 = self.encoder3(x)
126.         x,pool_indices3 = self.max_pool(x3)
127.         x4 = self.encoder4(x)
128.         x,pool_indices4 = self.max_pool(x4)
129.         x5 = self.encoder5(x)
130.         x,pool_indices5 = self.max_pool(x5)
132.         x = self.max_uppool(x,pool_indices5)
133.         x = crop(x, x5)
134.         x = self.decoder1(x)
135.         x = self.max_uppool(x,pool_indices4)
136.         x = crop(x, x4)
137.         x = self.decoder2(x)
138.         x = self.max_uppool(x,pool_indices3)
139.         x = crop(x, x3)
140.         x = self.decoder3(x)
141.         x = self.max_uppool(x,pool_indices2)
142.         x = crop(x, x2)
143.         x = self.decoder4(x)
144.         x = self.max_uppool(x,pool_indices1)
145.         x = crop(x, x1)
146.         x = self.decoder5(x)
148.         return x

复制代码

数据处理(dataloader)

数据集名称：CamVid
数据集下载地址：Object Recognition in Video Dataset

在这里进行下载，CamVid数据集有两种，一种是官方的就是上述的下载地址的，总共有32种类别，划分的会更加的细致。但是一般官网的太难打开了，所以我们可以通过Kaggle中的CamVid (Cambridge-Driving Labeled Video Database)进行下载。
还有一种就是11类别的(不包括背景)，会将一些语义相近的内容进行合并，就划分的没有这么细致，任务难度也会比较低一些。(如果你在网上找不到的话，可以在评论区发言或是私聊我要取)
CamVid 数据集主要用于自动驾驶场景中的语义分割，包含驾驶场景中的道路、交通标志、车辆等类别的标注图像。该数据集旨在推动自动驾驶系统在道路场景中的表现。
数据特点：

• 图像数量：包括701帧视频序列图像，分为训练集、验证集和测试集。
  
• 类别：包含32个类别(也有包含11个类别的)，包括道路、建筑物、车辆、行人等。
  
• 挑战：由于数据集主要来自城市交通场景，因此面临着动态变化的天气、光照、交通密度等挑战
  

这里我已经专门发了一篇博客对语义分割任务常用的数据集做了深入的介绍，已经具体讲解了其实现的处理代码。如果你对语义分割常用数据集有不了解的话，可以先去我的语义分割专栏中进行了解哦！！  我这里就直接附上代码了。

1. import os
2. from PIL import Image
3. import albumentations as A
4. from albumentations.pytorch.transforms import ToTensorV2
5. from torch.utils.data import Dataset, DataLoader
6. import numpy as np
7. import torch
8. # 11类
9. Cam_CLASSES = [ "Unlabelled","Sky","Building","Pole",
10.                 "Road","Sidewalk", "Tree","SignSymbol",
11.                 "Fence","Car","Pedestrian","Bicyclist"]
13. # 用于做可视化
14. Cam_COLORMAP = [
15.     [0, 0, 0],[128, 128, 128],[128, 0, 0],[192, 192, 128],
16.     [128, 64, 128],[0, 0, 192],[128, 128, 0],[192, 128, 128],
17.     [64, 64, 128],[64, 0, 128],[64, 64, 0],[0, 128, 192]
18. ]
21. # 转换RGB mask为类别id的函数
22. def mask_to_class(mask):
23.     mask_class = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int64)
24.     for idx, color in enumerate(Cam_COLORMAP):
25.         color = np.array(color)
26.         # 每个像素和当前颜色匹配
27.         matches = np.all(mask == color, axis=-1)
28.         mask_class[matches] = idx
29.     return mask_class
32. class CamVidDataset(Dataset):
33.     def __init__(self, image_dir, label_dir):
34.         self.image_dir = image_dir
35.         self.label_dir = label_dir
36.         self.transform = A.Compose([
37.             A.Resize(224, 224),
38.             A.HorizontalFlip(),
39.             A.VerticalFlip(),
40.             A.Normalize(),
41.             ToTensorV2(),
42.         ])
44.         self.images = sorted(os.listdir(image_dir))
45.         self.labels = sorted(os.listdir(label_dir))
47.         assert len(self.images) == len(self.labels), "Images and labels count mismatch!"
49.     def __len__(self):
50.         return len(self.images)
52.     def __getitem__(self, idx):
53.         img_path = os.path.join(self.image_dir, self.images[idx])
54.         label_path = os.path.join(self.label_dir, self.labels[idx])
56.         image = np.array(Image.open(img_path).convert("RGB"))
57.         label_rgb = np.array(Image.open(label_path).convert("RGB"))
59.         # RGB转类别索引
60.         mask = mask_to_class(label_rgb)
61.         #mask = torch.from_numpy(np.array(mask)).long()
63.         # Albumentations 需要 (H, W, 3) 和 (H, W)
64.         transformed = self.transform(image=image, mask=mask)
66.         return transformed['image'], transformed['mask'].long()
69. def get_dataloader(data_path, batch_size=4, num_workers=4):
70.     train_dir = os.path.join(data_path, 'train')
71.     val_dir = os.path.join(data_path, 'val')
72.     trainlabel_dir = os.path.join(data_path, 'train_labels')
73.     vallabel_dir = os.path.join(data_path, 'val_labels')
74.     train_dataset = CamVidDataset(train_dir, trainlabel_dir)
75.     val_dataset = CamVidDataset(val_dir, vallabel_dir)
77.     train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, pin_memory=True, num_workers=num_workers)
78.     val_loader = DataLoader(val_dataset, shuffle=False, batch_size=batch_size, pin_memory=True, num_workers=num_workers)
79.     return train_loader, val_loader

复制代码

评价指标(metric)

我们这里语义分割采用的评价指标为：PA(像素准确率)，CPA(类别像素准确率)，MPA(类别平均像素准确率)，IoU(交并比)，mIoU(平均交并比)，FWIoU(频率加权交并比)，mF1(平均F1分数)。
这里我已经专门发了一篇博客对这些平均指标做了深入的介绍，已经具体讲解了其实现的代码。如果你对这些评价指标有不了解的话，可以先去我的语义分割专栏中进行了解哦！！  我这里就直接附上代码了。

1. import numpy as np
3. __all__ = ['SegmentationMetric']
6. class SegmentationMetric(object):
7.     def __init__(self, numClass):
8.         self.numClass = numClass
9.         self.confusionMatrix = np.zeros((self.numClass,) * 2)
11.     def genConfusionMatrix(self, imgPredict, imgLabel):
12.         mask = (imgLabel >= 0) & (imgLabel < self.numClass)
13.         label = self.numClass * imgLabel[mask] + imgPredict[mask]
14.         count = np.bincount(label, minlength=self.numClass ** 2)
15.         confusionMatrix = count.reshape(self.numClass, self.numClass)
16.         return confusionMatrix
18.     def addBatch(self, imgPredict, imgLabel):
19.         assert imgPredict.shape == imgLabel.shape
20.         self.confusionMatrix += self.genConfusionMatrix(imgPredict, imgLabel)
21.         return self.confusionMatrix
23.     def pixelAccuracy(self):
24.         acc = np.diag(self.confusionMatrix).sum() / self.confusionMatrix.sum()
25.         return acc
27.     def classPixelAccuracy(self):
28.         denominator = self.confusionMatrix.sum(axis=1)
29.         denominator = np.where(denominator == 0, 1e-12, denominator)
30.         classAcc = np.diag(self.confusionMatrix) / denominator
31.         return classAcc
33.     def meanPixelAccuracy(self):
34.         classAcc = self.classPixelAccuracy()
35.         meanAcc = np.nanmean(classAcc)
36.         return meanAcc
38.     def IntersectionOverUnion(self):
39.         intersection = np.diag(self.confusionMatrix)
40.         union = np.sum(self.confusionMatrix, axis=1) + np.sum(self.confusionMatrix, axis=0) - np.diag(
41.             self.confusionMatrix)
42.         union = np.where(union == 0, 1e-12, union)
43.         IoU = intersection / union
44.         return IoU
46.     def meanIntersectionOverUnion(self):
47.         mIoU = np.nanmean(self.IntersectionOverUnion())
48.         return mIoU
50.     def Frequency_Weighted_Intersection_over_Union(self):
51.         denominator1 = np.sum(self.confusionMatrix)
52.         denominator1 = np.where(denominator1 == 0, 1e-12, denominator1)
53.         freq = np.sum(self.confusionMatrix, axis=1) / denominator1
54.         denominator2 = np.sum(self.confusionMatrix, axis=1) + np.sum(self.confusionMatrix, axis=0) - np.diag(
55.             self.confusionMatrix)
56.         denominator2 = np.where(denominator2 == 0, 1e-12, denominator2)
57.         iu = np.diag(self.confusionMatrix) / denominator2
58.         FWIoU = (freq[freq > 0] * iu[freq > 0]).sum()
59.         return FWIoU
61.     def classF1Score(self):
62.         tp = np.diag(self.confusionMatrix)
63.         fp = self.confusionMatrix.sum(axis=0) - tp
64.         fn = self.confusionMatrix.sum(axis=1) - tp
66.         precision = tp / (tp + fp + 1e-12)
67.         recall = tp / (tp + fn + 1e-12)
69.         f1 = 2 * precision * recall / (precision + recall + 1e-12)
70.         return f1
72.     def meanF1Score(self):
73.         f1 = self.classF1Score()
74.         mean_f1 = np.nanmean(f1)
75.         return mean_f1
77.     def reset(self):
78.         self.confusionMatrix = np.zeros((self.numClass, self.numClass))
80.     def get_scores(self):
81.         scores = {
82.             'Pixel Accuracy': self.pixelAccuracy(),
83.             'Class Pixel Accuracy': self.classPixelAccuracy(),
84.             'Intersection over Union': self.IntersectionOverUnion(),
85.             'Class F1 Score': self.classF1Score(),
86.             'Frequency Weighted Intersection over Union': self.Frequency_Weighted_Intersection_over_Union(),
87.             'Mean Pixel Accuracy': self.meanPixelAccuracy(),
88.             'Mean Intersection over Union(mIoU)': self.meanIntersectionOverUnion(),
89.             'Mean F1 Score': self.meanF1Score()
90.         }
91.         return scores

复制代码

训练流程(train)

到这里，所有的前期准备都已经就绪，我们就要开始训练我们的模型了。

1. def parse_arguments():
2.     parser = argparse.ArgumentParser()
3.     parser.add_argument('--data_root', type=str, default='../../data/CamVid/CamVid(11)', help='Dataset root path')
4.     parser.add_argument('--data_name', type=str, default='CamVid', help='Dataset class names')
5.     parser.add_argument('--model', type=str, default='Segnet', help='Segmentation model')
6.     parser.add_argument('--num_classes', type=int, default=12, help='Number of classes')
7.     parser.add_argument('--epochs', type=int, default=50, help='Epochs')
8.     parser.add_argument('--lr', type=float, default=0.005, help='Learning rate')
9.     parser.add_argument('--momentum', type=float, default=0.9, help='Momentum')
10.     parser.add_argument('--weight-decay', type=float, default=1e-4, help='Weight decay')
11.     parser.add_argument('--batch_size', type=int, default=8, help='Batch size')
12.     parser.add_argument('--checkpoint', type=str, default='./checkpoint', help='Checkpoint directory')
13.     parser.add_argument('--resume', type=str, default=None, help='Resume checkpoint path')
14.     return parser.parse_args()

复制代码

首先来看看我们的一些参数的设定，一般我们都是这样放在最前面，能够让人更加快速的了解其代码的一些核心参数设置。首先就是我们的数据集位置(data_root)，然后就是我们的数据集名称(classes_name)，这个暂时没什么用，因为我们目前只用了CamVid数据集，然后就是检测模型的选择(model)，我们选择Segnet模型，数据集的类别数(num_classes)，训练epoch数，这个你设置大一点也行，因为我们会在训练过程中保存最好结果的模型的。学习率(lr)，动量(momentum)，权重衰减(weight-decay)，这些都属于模型超参数，大家可以尝试不同的数值，多试试，就会有个大致的了解的，批量大小(batch_size)根据自己电脑性能来设置，一般都是为2的倍数，保存权重的文件夹(checkpoint)，是否继续训练(resume)。
[code]def train(args):    if not os.path.exists(args.checkpoint):        os.makedirs(args.checkpoint)    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')    n_gpu = torch.cuda.device_count()    print(f"Device: {device}, GPUs available: {n_gpu}")    # Dataloader    train_loader, val_loader = get_dataloader(args.data_root, batch_size=args.batch_size)    train_dataset_size = len(train_loader.dataset)    val_dataset_size = len(val_loader.dataset)    print(f"Train samples: {train_dataset_size}, Val samples: {val_dataset_size}")    # Model    model = get_model(num_classes=args.num_classes)    model.to(device)    # Loss + Optimizer + Scheduler    criterion = nn.CrossEntropyLoss(ignore_index=0)    #optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)    scaler = torch.cuda.amp.GradScaler()    # Resume    start_epoch = 0    best_miou = 0.0    if args.resume and os.path.isfile(args.resume):        print(f"Loading checkpoint '{args.resume}'")        checkpoint = torch.load(args.resume)        start_epoch = checkpoint['epoch']        best_miou = checkpoint['best_miou']        model.load_state_dict(checkpoint['model_state_dict'])        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])        print(f"Loaded checkpoint (epoch {start_epoch})")    # Training history    history = {        'train_loss': [],        'val_loss': [],        'pixel_accuracy': [],        'miou': []    }    print(f"
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！