OpenCV与AI深度学习 | 基于PyTorch语义分割实现洪水识别（数据集 + 源码）

本文来源公众号“OpenCV与AI深度学习”，仅用于学术分享，侵权删，干货满满。

背景

全世界有数百万人因洪水而流离失所。确实，通过使用深度学习 + 计算机视觉，我们无法总是预测下一次洪水何时会来袭。但我们可以在洪水灾区的图像上训练语义分割算法。这样的模型可以帮助分析和决策未来的情况。为了尽我们的微薄之力，我们将在本文中使用 PyTorch训练一个用于洪水识别的语义分割模型。

洪水分割数据集

洪水分割看起来可能与卫星图像中的水体分割类似。我们在尝试在自定义数据集上训练 PyTorch DeepLabV3 模型时涵盖了水体分割项目。

但在尝试划分洪水区域时，情况会变得更加复杂。让我们先探索数据集，然后再进一步讨论。

我们将使用 Kaggle 的洪水区域分割的修改版本。这是一个非常好且干净的数据集。但它存在一些图像扩展问题，并且其中一个图像已损坏。

https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation

我准备了这个数据集的修改版本。这是Kaggle 上的洪水分割数据集训练/验证分割。我们将使用此版本的数据集。以下是更改：

下载并提取数据集后，您将发现以下结构。

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_