本文来源公众号“OpenCV与AI深度学习”,仅用于学术分享,侵权删,干货满满。
原文链接:基于PyTorch语义分割实现洪水识别(数据集 + 源码)
背 景
全世界有数百万人因洪水而流离失所。确实,通过使用深度学习 + 计算机视觉,我们无法总是预测下一次洪水何时会来袭。但我们可以在洪水灾区的图像上训练语义分割算法。这样的模型可以帮助分析和决策未来的情况。为了尽我们的微薄之力,我们将在本文中使用 PyTorch训练一个用于洪水识别的语义分割模型。
洪水分割数据集
洪水分割看起来可能与卫星图像中的水体分割类似。我们在尝试在自定义数据集上训练 PyTorch DeepLabV3 模型时涵盖了水体分割项目。
但在尝试划分洪水区域时,情况会变得更加复杂。让我们先探索数据集,然后再进一步讨论。
我们将使用 Kaggle 的洪水区域分割的修改版本。这是一个非常好且干净的数据集。但它存在一些图像扩展问题,并且其中一个图像已损坏。
https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation
我准备了这个数据集的修改版本。这是Kaggle 上的洪水分割数据集训练/验证分割。我们将使用此版本的数据集。以下是更改:
-
-
数据集现在包含训练和验证部分。
-
有 257 张训练图像和蒙版以及 32 张验证图像和蒙版。
-
现在数据集中没有损坏的图像。
-
下载并提取数据集后,您将发现以下结构。
flood-area-segmentation
├── train_images
├── train_masks
├── valid_images
└── valid_masks
所有地面实况图像均为 RGB 格式。所有蒙版均为灰度图像,洪水区域为白色像素(255),其他区域为黑色像素(0)。因此,这是一个二元语义分割数据集。以下是一些图像。
洪水识别的语义分割为何如此困难?
让我们来谈谈为什么根据这个数据集进行洪水分割很困难。
首先,这个数据集仅包含 257 张训练图像和掩码。当然,这不足以训练最先进的语义分割模型。我们需要足够的增强技术才能减少损失。
其次,洪水区域之间会有很多小的背景像素,包括房屋、农田,甚至是部分被淹没的车辆。
以下列举几个相同的例子。
模型必须正确学习这些背景类别,而不是对其进行分割。对于较小的模型来说,这可能特别具有挑战性。
考虑到上述问题,我们将尽力训练出最好的模型。
PyTorch安装与项目目录
该项目的代码使用TORCH 1.12.0 和 TORCHVISION 0.13.0。您也可以使用较新的版本。
此外,Albumentations是图像和蒙版增强的要求之一。
请务必安装以上两个库。
洪水识别语义分割的目录结构:
.
├── input
│ ├── flood-area-segmentation
│ │ ├── train_images
│ │ ├── train_masks
│ │ ├── valid_images
│ │ └── valid_masks
│ └── inference_data
│ └── video_1.mp4
├── outputs
│ ├── inference_results
│ │ ├── 1005.jpg
│ │ ...
│ │ └── 38.jpg
│ ├── inference_results_video
│ │ └── video_1.mp4
│ ├── valid_preds
│ │ ├── e0_b7.jpg
│ │ ...
│ │ └── e9_b7.jpg
│ ├── accuracy.png
│ ├── best_model.pth
│ ├── loss.png
│ └── model.pth
└── src
├── config.py
├── datasets.py
├── engine.py
├── inference_image.py
├── inference_video.py
├── metrics.py
├── model.py
├── train.py
└── utils.py
input目录包含洪水区域分割数据集和推理数据。
outputs目录中,我们拥有所有与训练和推理相关的输出。这包括模型权重文件、推理后的图像结果以及验证循环的分割输出。
src目录包含所有 Python 源代码文件。由于有 9 个 Python 文件,我们无法详细介绍每个文件。但我们将在下一节中探讨代码的重要部分。
使用DeepLabV3 ResNet50 进行洪水识别的语义分割
从现在开始,我们将讨论项目的技术方面。在开始训练之前,我们将讨论以下主题:
- 配置文件。
- 用于洪水识别语义分割的深度学习模型。
- 数据集准备策略。
- 训练集的数据增强。
- 训练超参数和参数。这些包括优化器、学习率以及根据 VRAM 可用性选择批处理大小。我们还将讨论训练图像大小的影响。
配置文件
我们将在其余 Python 文件中使用一些配置。我们不必反复重写这些配置,只需config.py归档并存储这些。
ALL_CLASSES = ['background', 'flood']
LABEL_COLORS_LIST = [
(0, 0, 0), # Background.
(255, 255, 255), # Flood.
]
VIS_LABEL_MAP = [
(0, 0, 0), # Background.
(255, 0, 0), # Flood.
]
ALL_CLASSED列表包含数据集中的类的名称。
LABEL_COLORS_LIST列表包含数据集中的 RGB 颜色元组。在数据集中,背景类别为黑色,淹没区域为白色。
最后,VIS_LABEL_MAP包含我们将用于可视化的 RGB 颜色元组。为了清晰起见,我们将在推理过程中使用红色来划分被淹没的区域。
DeepLabV3 ResNet50 模型
我们将在此项目中使用优质、可靠、经典的DeepLabV3 和 ResNet50主干进行语义分割。
准备用于自定义模型训练的模型非常简单。事实上,我们只需使用以下几行代码即可准备用于自定义语义分割训练的模型。
import torch.nn as nn
from torchvision.models.segmentation import deeplabv3_resnet50
def prepare_model(num_classes=2):
model = deeplabv3_resnet50(weights='DEFAULT')
model.classifier[4] = nn.Conv2d(256, num_classes, 1)
model.aux_classifier[4] = nn.Conv2d(256, num_classes, 1)
return model
为了获得更好的性能,我们将使用预训练的权重。另外,我们需要修改分类器和辅助分类器根据数据集中的类别数量,将 head 分为两类:背景类和洪水泛滥区类。
选择带有 ResNet 的 DeepLabV3 的主要原因是速度和性能。在 PyTorch 中,我们还有另外两个 DeepLabV3 选项:
- 采用 ResNet101 主干。虽然性能良好,但需要更多内存进行训练,推理速度也较慢。
- 还有以 MobileNet 为骨干的 DeepLabV3。速度很快,但性能不太好。
为此,我们选择 DeepLabV3 ResNet50 来兼顾两全其美。
数据集准备
虽然我们不会介绍数据集准备的整个代码,但这里有几件重要的事情需要注意。
其中之一就是数据增强策略。
请记住,我们只有 257 张图像和蒙版用于训练。这不足以训练出一个好的模型。但数据增强可以帮助我们做到这一点。我们将使用 Albumentations 库轻松地将增强功能应用于图像和蒙版。
我们将对洪水分割数据集的训练图像和蒙版应用以下增强功能。
- 水平翻转
- 隨機亮度對比
- 随机太阳光晕
- 随机雾
- 图像压缩
除了水平翻转(以 0.5 的概率应用)之外,我们将以 0.2 的概率应用所有其他增强。
这是随机应用这些增强功能后图像的样子。
如您所见,仅对蒙版应用了空间增强(翻转)。影响颜色、色调或饱和度的其他像素级操作不应用于蒙版。使用 Albumentations 很容易实现这一点。请查看dataset.py文件以获取更多详细信息。
使用脚本、指标以及训练和验证代码
utils.py文件包含许多辅助函数和实用程序类。其中包括用于保存最佳模型的类、用于保存准确率和损失图的函数,以及用于在图像上叠加分割图的函数。
metrics.py包含pix_acc函数。这将返回标记像素的总数和正确分类的像素。我们将使用这些来计算engine.py文件。
因此,engine.py包含我们将在每个时期调用的训练和验证函数。
在洪水分割数据集上训练 DeepLabV3
最后,我们可以在洪水分割数据集上训练 DeepLabV3 ResNet50 模型了。
注意:所有训练和推理实验均在具有 10GB RTX 3080 GPU、第 10 代 i7 CPU 和 32 GB RAM 的机器上进行。
train.py是我们将用来启动训练的驱动脚本。
在开始训练之前,让我们先来看看训练脚本支持的参数解析器。
--epoches:我们想要训练模型的时期数。
--lr:优化器的学习率,默认为0.0001。
--batch:数据加载器的批次大小。
--imgsz:我们想要用来训练模型的图像大小。语义分割模型通常需要更高分辨率的图像才能取得良好的效果。我们将在 512×512 图像上训练模型,这也是默认值。
--scheduler:这是一个布尔标志,表示我们是否要使用学习率调度程序。如果我们传递此标志,则学习率将在 25 个时期后降低 10 倍。
如果在训练时遇到 OOM(内存不足)错误,请考虑减少图像大小或批量大小。
我们可以通过在终端中执行以下命令来开始训练源码作为当前工作目录。
python train.py --epochs 50 --batch 4 --lr 0.0001 --scheduler
我们对模型进行 50 个epoch的训练,batch size为 4,起始学习率为 0.0001,并应用学习率调度程序。
以下是终端的截断输出。
python train.py --epochs 50 --batch 4 --lr 0.0001 --scheduler
Namespace(epochs=50, lr=0.0001, batch=4, scheduler=True)
41,994,308 total parameters.
41,994,308 training parameters.
Adjusting learning rate of group 0 to 1.0000e-04.
EPOCH: 1
Training
Loss: 0.4229 | PixAcc: 83.23: 100%|████████████████████| 64/64 [00:50<00:00, 1.26it/s]
Validating
Loss: 0.3060 | PixAcc: 89.33: 100%|████████████████████| 8/8 [00:02<00:00, 3.57it/s]
Best validation loss: 0.38311174884438515
Saving best model for epoch: 1
Train Epoch Loss: 0.5018, Train Epoch PixAcc: 64.9515
Valid Epoch Loss: 0.3831, Valid Epoch PixAcc: 86.8635
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
EPOCH: 2
Training
Loss: 0.3776 | PixAcc: 81.60: 100%|████████████████████| 64/64 [00:49<00:00, 1.30it/s]
Validating
Loss: 0.4301 | PixAcc: 91.83: 100%|████████████████████| 8/8 [00:02<00:00, 3.84it/s]
Train Epoch Loss: 0.3887, Train Epoch PixAcc: 72.3201
Valid Epoch Loss: 0.4537, Valid Epoch PixAcc: 94.0050
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
EPOCH: 3
Training
Loss: 0.4062 | PixAcc: 68.83: 100%|████████████████████| 64/64 [00:49<00:00, 1.29it/s]
Validating
Loss: 0.2974 | PixAcc: 93.10: 100%|████████████████████| 8/8 [00:02<00:00, 3.86it/s]
Train Epoch Loss: 0.3698, Train Epoch PixAcc: 75.0250
Valid Epoch Loss: 0.4060, Valid Epoch PixAcc: 94.7763
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
.
.
.
EPOCH: 50
Training
Loss: 0.1831 | PixAcc: 80.59: 100%|████████████████████| 128/128 [00:39<00:00, 3.26it/s]
Validating
Loss: 0.1339 | PixAcc: 98.66: 100%|████████████████████| 16/16 [00:02<00:00, 7.16it/s]
Train Epoch Loss: 0.1459, Train Epoch PixAcc: 91.4595
Valid Epoch Loss: 0.2437, Valid Epoch PixAcc: 88.4480
Adjusting learning rate of group 0 to 1.0000e-06.
--------------------------------------------------
TRAINING COMPLETE
分析 DeepLabV3 洪水分割结果
有趣的是,我们在 epoch 33 上获得了最小的验证损失 0.222,而最高的像素准确度出现在 epoch 3。这种情况有点奇怪,因为最佳模型是基于最小损失保存的。让我们来看看准确度和损失图。
我们可以看到在初始阶段和整个训练过程中出现了很多波动,直到学习率下降。这意味着我们需要学习率调度程序。也许我们可以用更低的初始学习率进行训练,以缓解这种不稳定的训练。
无论如何,我们现在有一个训练好的模型。让我们对图像和视频进行一些推理实验。
图像洪水分割推理
对于图像推理,我们将从数据集中选择相同的验证图像。
为了对图像进行推理,我们将使用inference_image.py脚本。
python inference_image.py --input ../input/flood-area-segmentation/valid_images/
我们使用--input标志提供验证图像目录的路径。
您可以在里面找到outputs/inference_results目录。
以下是一些输出。第一列显示模型的预测,右列显示真实值掩码。
虽然结果可能并不完美,但效果相当不错。该模型仍然难以区分细长的陆地线和被淹没的区域。但我们可以通过更多的训练数据和更好的训练技术来纠正这个问题。
视频洪水分割推理
python inference_video.py --input ../input/inference_data/ video.mp4
在 RTX 3080 GPU 上,我们平均获得 21 FPS。这可能并不像你想象的那么好。我们在将帧输入模型之前不会调整它们的大小。将它们调整为 512×512 分辨率将稍微提高 FPS。
以下是我们得到的输出。
总体来说,效果还不错。如果将其部署到无人机上,即可实时观察洪水情况了。
源码下载:
链接: https://pan.baidu.com/s/1Dg9q3ofvUKlm5AAYnoNnJA 提取码: z7u3
数据集下载:
https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation
THE END !
文章结束,感谢阅读。您的点赞,收藏,评论是我继续更新的动力。大家有推荐的公众号可以评论区留言,共同学习,一起进步。