在上一篇 文章中,我们Fast.ai 在卫星图像中检测云轮廓,检测物体轮廓被称为语义分割。虽然我们用几行代码就能达到 96% 的准确率,但该模型无法考虑数据集中提供的所有输入通道(红、绿、蓝和近红外)。问题在于,深度学习框架(如 Keras、Fast.ai 甚至 PyTorch)中的大多数语义分割模型都是为处理 RGB 图像而设计的,并带有预先训练好的权重。除此之外,这些库中的视觉模块也仅限于 RGB 文件。这就是我们在前面的故事中忽略近红外通道,只使用 RGB patchs的原因。
这促使我采用了一种完全不同的方法,从头开始构建我自己的 U-Net,而不是使用视觉库来准备数据集。本篇我将介绍如何使用简单的 PyTorch 函数创建数据集,而下一篇将继续介绍如何创建一个简单的U-NET分割模型,并对其进行训练以获得结果。
1、数据准备
在上一篇文章中我们使用的数据集,是由 8400 个大小为 384x384 的训练片段组成(适合深度学习)。还有一组 9201 个图像块patchs用于测试,但我们暂时不处理它们。这些patchs分别位于红、绿、蓝和近红外(NIR)通道的目录中,另外还有一个用于参考掩膜(真实值ground truth*_gt)的目录。
2、PyTorch 数据集类
在上一篇文章中,我们在磁盘中创建了 rgb_patch*.tif 文件,并使用 PIL 将条带合并为 384x`384x3