使用TorchGeo进行地理空间深度学习

前言

TorchGeo是一个PyTorch域库，提供特定于地理空间数据的数据集、采样器、转换和预训练模型。

https://github.com/microsoft/torchgeo

几十年来，地球观测卫星、飞机和最近的无人机平台一直在收集越来越多的地球表面图像。有了关于季节和长期趋势的信息，遥感图像对于解决人类面临的一些最大挑战，包括气候变化适应、自然灾害监测、水资源管理以及日益增长的全球人口的粮食安全，可能是非常宝贵的。从计算机视觉的角度来看，这包括土地覆盖测绘（语义分割）、森林砍伐和洪水监测（变化检测）、冰川流（像素跟踪）、飓风跟踪和强度估计（回归）以及建筑和道路检测（对象检测、实例分割）等应用。通过利用深度学习架构、更便宜、更强大的GPU以及数PB的免费可用卫星图像数据集的最新进展，我们可以更接近于解决这些重要问题。

2005年8月28日拍摄的美国国家海洋和大气管理局卡特里娜飓风卫星图像（来源）。像TorchGeo这样的地理空间机器学习库可以用于检测、跟踪和预测飓风和其他自然灾害的未来轨迹。

1. 挑战

在传统的计算机视觉数据集中，例如 ImageNet，图像文件本身往往相当简单且易于使用。 大多数图像具有 3 个光谱波段 (RGB)，以 PNG 或 JPEG 等常见文件格式存储，并且可以使用 PIL 或 OpenCV 等流行软件库轻松加载。 这些数据集中的每个图像通常都足够小，可以直接传递到神经网络。 此外，大多数这些数据集包含有限数量的精心策划的图像，这些图像被假定为独立且分布相同，从而使训练-验证-测试拆分变得简单明了。 由于这种相对同质性，相同的预训练模型（例如，在 ImageNet 上预训练的 CNN）已显示在使用迁移学习方法的广泛视觉任务中有效。 现有的库，如 torchvision，可以很好地处理这些简单的情况，并在过去十年中被用于在视觉任务方面取得重大进展。

遥感图像并不是那么统一。 与简单的 RGB 图像不同，卫星倾向于捕捉多光谱（Landsat 8 有 11 个光谱带）甚至高光谱（Hyperion 有 242 个光谱带）的图像。 这些图像捕获的信息波长范围更广 (400 nm–15 µm)，远在可见光谱之外。 不同的卫星也有非常不同的空间分辨率——GOES 的分辨率为 4 km/px，Maxar 图像为 30 cm/px，无人机图像分辨率可高达 7 mm/px。 这些数据集几乎总是有时间成分，卫星重访是每天、每周或每两周一次。 图像通常与数据集中的其他图像有重叠，需要根据地理元数据拼接在一起。 这些图像往往非常大（例如 10K x 10K 像素），因此不可能将整个图像传递给神经网络。 这些数据以数百种不同的光栅和矢量文件格式（如 GeoTIFF 和 ESRI Shapefile）分布，需要加载 GDAL 等专业库。

从左到右：墨卡托、阿尔伯斯等面积和中断的古德等距投影（来源）。 地理空间数据与许多不同类型的参考系统之一相关联，这些参考系统将 3D 地球投影到 2D 表示上。 合并来自不同来源的数据通常涉及重新投影到一个公共参考系统，以确保所有层都对齐。

虽然每个图像都是 2D 的，但地球本身是 3D 的。 为了将图像拼接在一起，首先需要将它们投影到地球的二维表示上，称为坐标参考系统 (CRS)。 大多数人都熟悉像墨卡托这样扭曲区域大小的等角表示（格陵兰看起来比非洲大，尽管非洲大 15 倍），但还有许多其他常用的 CRS。 每个数据集可能使用不同的 CRS，单个数据集中的每个图像也可能在唯一的 CRS 中。 为了使用来自多层的数据，它们必须共享一个通用的 CRS，否则数据将无法正确对齐。 对于那些不熟悉遥感数据的人来说，这可能是一项艰巨的任务。

即使您在索引期间正确地对图像进行了地理配准，但如果您不将它们投影到公共 CRS，您最终会得到周围没有数据值的旋转图像，并且图像不会像素对齐。

2、解决方案

目前，如果不具备这两个截然不同领域的专业知识，同时使用深度学习模型和地理空间数据可能会非常具有挑战性。 为了应对这些挑战，我们构建了 TorchGeo，这是一个用于处理地理空间数据的 PyTorch 域库。 TorchGeo 旨在使其变得简单：

• 让机器学习专家处理地理空间数据，以及

• 供遥感专家探索机器学习解决方案。

TorchGeo 不仅是一个研究项目，还是一个生产质量的库，它使用持续集成来测试在一系列平台（Linux、macOS、Windows）上使用一系列 Python 版本进行的每次提交。 它可以通过任何您喜欢的包管理器轻松安装，包括 pip、conda 和 spack：

pip install torchgeo

TorchGeo 旨在与其他 PyTorch 域库（如 torchvision、torchtext 和 torchaudio）具有相同的 API。 如果您已经在计算机视觉数据集的工作流程中使用 torchvision，则只需更改几行代码即可切换到 TorchGeo。 所有 TorchGeo 数据集和采样器都与 PyTorch DataLoader 类兼容，这意味着您可以利用 PyTorch Lightning 等包装库进行分布式训练。 在以下部分中，我们将探索 TorchGeo 的可能用例，以展示它的使用有多么简单。

3、地理空间数据集和采样器

我们将 A) 来自 Landsat 8 的场景和 B) 农田数据层标签组合在一起的示例应用程序，即使这些文件位于不同的 EPSG 投影中。 我们想使用地理空间边界框作为索引从这些数据集中对分块 C) 和 D) 进行采样。

许多遥感应用涉及使用地理空间数据集——具有地理元数据的数据集。 在 TorchGeo 中，我们定义了一个 GeoDataset 类来表示这些类型的数据集。 每个 GeoDataset 不是按整数索引，而是按时空边界框索引，这意味着可以智能组合覆盖不同地理范围的两个或多个数据集。

在此示例中，我们展示了使用 TorchGeo 使用地理空间数据以及从 Landsat 和农田数据层 (CDL) 数据的组合中采样小图像块是多么容易。 首先，我们假设用户下载了 Landsat 7 和 8 图像。 由于 Landsat 8 的光谱波段比 Landsat 7 多，我们将只使用两颗卫星共有的波段。 我们将通过合并这两个数据集来创建一个包含 Landsat 7 和 8 数据的所有图像的数据集。

from torch.utils.data import DataLoader
from torchgeo.datasets import CDL, Landsat7, Landsat8, stack_samples
from torchgeo.samplers import RandomGeoSampler

landsat7 = Landsat7(root="...")
landsat8 = Landsat8(root="...", bands=Landsat8.all_bands[1:-2])
landsat = landsat7 | landsat8

接下来，我们获取此数据集和 CDL 数据集之间的交集。 我们希望采用交集而不是并集，以确保我们只从同时拥有 Landsat 和 CDL 数据的区域进行采样。 请注意，我们可以自动下载和校验 CDL 数据。 另请注意，每个数据集可能包含不同 CRS 或分辨率的文件，但 TorchGeo 会自动确保使用匹配的 CRS 和分辨率。

cdl = CDL(root="...", download=True, checksum=True)
dataset = landsat & cdl

该数据集现在可以与 PyTorch 数据加载器一起使用。 与基准数据集不同，地理空间数据集通常包含非常大的图像。 例如，CDL 数据集由覆盖整个美国本土的单个图像组成。 为了使用地理空间坐标从这些数据集中进行采样，TorchGeo 定义了许多采样器。 在这个例子中，我们将使用一个随机采样器，每个 epoch 返回 256 x 256 像素图像和 10,000 个样本。 我们还将使用自定义整理功能将每个样本字典组合成一个小批量样本。

sampler = RandomGeoSampler(dataset, size=256, length=10000)
dataloader = DataLoader(dataset, batch_size=128, sampler=sampler, collate_fn=stack_samples)

此数据加载器现在可以用于您的正常训练/评估流程。

for batch in dataloader:
    image = batch["image"]
    mask = batch["mask"]

    # train a model, or make predictions using a pre-trained model

许多应用程序都涉及基于像这样的地理空间元数据智能地组合数据集。 例如，用户可能想要：

• 合并多个图像源的数据集并将它们视为等效（例如，Landsat 7 和 8）

• 合并不同地理空间位置的数据集（例如，纽约州切萨皮克和宾夕法尼亚州）

这些组合要求所有查询都存在于至少一个数据集中，并且可以使用 UnionDataset 创建。 同样，用户可能希望：

• 结合图像和目标标签并同时从两者中采样（例如，Landsat 和 CDL）

• 结合多个图像源的数据集以进行多模态学习或数据融合（例如 Landsat 和 Sentinel）

这些组合要求所有查询都存在于两个数据集中，并且可以使用 IntersectionDataset 创建。 当您使用交集 (&) 和联合 (|) 运算符时，TorchGeo 会自动为您组合这些数据集。

4、多光谱和地理空间变换

在深度学习中，通常会增加和转换数据，以便模型对输入空间的变化具有鲁棒性。 地理空间数据可能存在季节性变化和扭曲效应等变化，以及云层覆盖和大气失真等图像处理和捕获问题。 TorchGeo 利用 Kornia 库的增强和转换，支持 GPU 加速并支持具有 3 个以上通道的多光谱图像。

传统的地理空间分析计算和可视化光谱指数，这些指数是多光谱波段的组合。 光谱指数旨在突出显示与某些应用相关的多光谱图像中感兴趣的区域，例如植被健康、人为变化或城市化进程加快的区域或积雪。 TorchGeo 支持多种变换，这些变换可以计算常见的光谱指数并将它们作为附加波段附加到多光谱图像张量中。

下面，我们展示了一个简单示例，我们在 Sentinel-2 图像上计算归一化差异植被指数 (NDVI)。 NDVI 衡量植被的存在和植被健康状况，并计算为红色和近红外 (NIR) 光谱带之间的归一化差异。 光谱索引变换对从 TorchGeo 数据集返回的样本字典进行操作，并将生成的光谱索引附加到图像通道维度。

首先，我们实例化一个 Sentinel-2 数据集并加载一个示例图像。 然后，我们绘制此数据的真实颜色 (RGB) 表示以查看我们正在查看的区域。

import matplotlib.pyplot as plt
from torchgeo.datasets import Sentinel2
from torchgeo.transforms import AppendNDVI

dataset = Sentinel2(root="...")
sample = dataset[...]
fig = dataset.plot(sample)
plt.show()

接下来，我们实例化并计算 NDVI 变换，将这个新通道附加到图像的末尾。 Sentinel-2 图像的红色波段使用索引 0，近红外波段使用索引 3。 为了可视化数据，我们还对图像进行了归一化处理。 NDVI 值的范围可以从 -1 到 1，但我们希望使用 0 到 1 的范围进行绘图。

transform = AppendNDVI(index_red=0, index_nir=3)
sample = transform(sample)
sample["image"][-1] = (sample["image"][-1] + 1) / 2
plt.imshow(sample["image"][-1], cmap="RdYlGn_r")
plt.show()

2018 年 11 月 16 日由 Sentinel-2 卫星拍摄的得克萨斯山地区的真彩色（左）和 NDVI（右）。 在NDVI图像中，红色表示水体，黄色表示土壤贫瘠，浅绿色表示不健康的植被，深绿色表示健康的植被。

5、基准数据集

计算机视觉进步背后的驱动因素之一是标准化基准数据集（如 ImageNet 和 MNIST）的存在。 使用这些数据集，研究人员可以直接比较不同模型和训练程序的性能，以确定哪个模型表现最好。 在遥感领域，有很多这样的数据集，但由于上述处理这些数据的困难以及缺乏用于加载这些数据集的现有库，许多研究人员选择使用他们自己的自定义数据集。

TorchGeo 的目标之一是为这些现有数据集提供易于使用的数据加载器。 TorchGeo 包括许多基准数据集——包括输入图像和目标标签的数据集。 这包括用于图像分类、回归、语义分割、对象检测、实例分割、变化检测等任务的数据集。

如果您以前使用过 torchvision，那么这些类型的数据集应该很熟悉。 在此示例中，我们将为Northwestern Polytechnical University (NWPU)超高分辨率十级 (VHR-10) 地理空间对象检测数据集创建一个数据集。 这个数据集可以自动下载、校验和提取，就像 torchvision 一样。

from torch.utils.data import DataLoader
from torchgeo.datasets import VHR10

dataset = VHR10(root="...", download=True, checksum=True)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=4)

for batch in dataloader:
    image = batch["image"]
    label = batch["label"]

    # train a model, or make predictions using a pre-trained model

所有 TorchGeo 数据集都与 PyTorch 数据加载器兼容，使其易于集成到现有的培训工作流程中。 TorchGeo 中的基准数据集与 torchvision 中的类似数据集之间的唯一区别是每个数据集返回一个字典，其中包含每个 PyTorch 张量的键。

来自在 NWPU VHR-10 数据集上训练的 Mask R-CNN 模型的示例预测。 该模型预测具有高置信度分数的所有对象的清晰边界框和掩码。

6、PyTorch Lightning 的可重复性

TorchGeo 的另一个关键目标是再现性。 对于许多这些基准数据集，没有预定义的训练-验证-测试拆分，或者预定义的拆分存在类不平衡或地理分布问题。 因此，文献中报告的性能指标要么无法重现，要么不能表明预训练模型在不同地理位置的工作情况。

为了促进文献中发表的结果之间的直接比较，并进一步减少在 TorchGeo 中使用数据集运行实验所需的样板代码，我们创建了 PyTorch Lightning 数据模块，其中包含定义明确的训练-验证-测试拆分和用于分类等各种任务的训练器 、回归和语义分割。 这些数据模块展示了如何合并来自 kornia 库的增强，包括预处理转换（使用预先计算的通道统计数据），并让用户轻松地试验与数据本身相关的超参数（与建模过程相反）。 在 Inria 航拍图像标签数据集上训练语义分割模型就像导入几行代码和四行代码一样简单。

from pytorch_lightning import Trainer
from torchgeo.datamodules import InriaAerialImageLabelingDataModule
from torchgeo.trainers import SemanticSegmentationTask

datamodule = InriaAerialImageLabelingDataModule(root_dir="...", batch_size=64, num_workers=6)
task = SemanticSegmentationTask(segmentation_model="unet", encoder_weights="imagenet", learning_rate=0.1)
trainer = Trainer(gpus=1, default_root_dir="...")

trainer.fit(model=task, datamodule=datamodule)

由在 Inria 航拍图像标记数据集上训练的 U-Net 模型生成的建筑物分割。 重现这些结果就像几个导入和四行代码一样简单，使不同模型和训练技术的比较简单易行。

在我们的预印本中，我们展示了一组结果，这些结果使用上述数据模块和训练器来对 TorchGeo 中的几个数据集的简单建模方法进行基准测试。 例如，我们发现一个简单的 ResNet-50 可以在 So2Sat 数据集上实现最先进的性能。 在处理遥感数据问题时，这些类型的基线结果对于评估不同建模选择的贡献非常重要。

7、未来工作

为了使 TorchGeo 尽可能易于使用，尤其是对于没有深度学习经验的用户，还有很多剩余的工作要做。 我们计划实现这一目标的方法之一是扩展我们的教程，以包括诸如“编写自定义数据集”和“迁移学习”之类的主题，或诸如“土地覆盖测绘”和“对象检测”之类的任务。

我们正在进行的另一个重要项目是预训练模型。 大多数遥感研究人员使用非常小的标记数据集，并且可以从预训练模型和迁移学习方法中受益。 TorchGeo 是第一个提供多光谱图像预训练模型的深度学习库。 我们的目标是为不同的图像模态（光学、SAR、多光谱）和特定平台（Landsat、Sentinel、MODIS）提供模型，以及显示它们在不同训练数据量下的性能的基准结果。 自监督学习是训练此类模型的一种很有前途的方法。 卫星图像数据集通常包含数 PB 的图像，但准确标记的数据集更难获得。 自监督学习方法将使我们能够直接在原始图像上进行训练，而无需大型标记数据集。

除了这些更大的项目之外，我们一直在寻求添加新的数据集、数据增强转换和采样策略。 如果您精通 Python 并且有兴趣为 TorchGeo 做出贡献，我们很乐意看到您的贡献！ TorchGeo 在 MIT 许可证下是开源的，因此您几乎可以在任何项目中使用它。

致谢

Adam Stewart (University of Illinois at Urbana-Champaign), Caleb Robinson (Microsoft AI for Good Research Lab), Isaac Corley (University of Texas at San Antonio)