作者：Adam Stewart（伊利诺伊大学厄巴纳-香槟分校）、Caleb Robinson（微软人工智能向善研究实验室）、Isaac Corley（德克萨斯大学圣安东尼奥分校）

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

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

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

挑战

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

遥感影像并非如此统一。卫星往往捕获的图像不是简单的 RGB 图像，而是多光谱（Landsat 8 有 11 个光谱波段）甚至是高光谱（Hyperion 有 242 个光谱波段）。这些图像捕获了更广泛波长范围（400 nm–15 µm）的信息，远超出可见光谱。不同卫星的空间分辨率也差异很大—GOES 的分辨率为 4 公里/像素，Maxar 影像为 30 厘米/像素，而无人机影像分辨率可高达 7 毫米/像素。这些数据集几乎总是包含时间维度，卫星重访周期为每天、每周或每两周。图像通常与数据集中的其他图像存在重叠，并且需要根据地理元数据拼接在一起。这些图像往往非常大（例如 10K x 10K 像素），因此不可能将整张图像通过神经网络。这些数据以数百种不同的栅格和矢量文件格式（如 GeoTIFF 和 ESRI Shapefile）分布，需要像 GDAL 这样的专业库来加载。

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

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

即使您在索引时正确进行了地理参考，如果不将图像投影到共同的 CRS，最终会得到带有无数据值旋转的图像，且图像不会像素对齐。

解决方案

目前，如果没有在这两个截然不同的领域都具备专业知识，同时处理深度学习模型和地理空间数据可能会相当具有挑战性。为了解决这些挑战，我们构建了 TorchGeo，这是一个用于处理地理空间数据的 PyTorch 领域库。TorchGeo 旨在简化：

机器学习专家处理地理空间数据的过程，以及
遥感专家探索机器学习解决方案的过程。

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

$ pip install torchgeo

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

地理空间数据集和采样器

示例应用：结合 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 定义了许多采样器。在此示例中，我们将使用一个随机采样器，它返回 256 x 256 像素的图像，每个 epoch 返回 10,000 个样本。我们还将使用自定义的 collation 函数将每个样本字典组合成一个样本小批量。

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）
组合不同地理位置的数据集（例如，切萨皮克湾 NY 和 PA）

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

同时组合图像和目标标签并从中采样（例如，Landsat 和 CDL）
组合来自多个图像源的数据集用于多模态学习或数据融合（例如，Landsat 和 Sentinel）

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

多光谱和地理空间变换

在深度学习中，通常会对数据进行增强和变换，以便模型对输入空间的变化具有鲁棒性。地理空间数据可能存在季节性变化和扭曲效应等变化，以及云层覆盖和大气畸变等图像处理和捕获问题。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()

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

基准数据集

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

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

如果您之前使用过 torchvision，应该熟悉这类数据集。在此示例中，我们将为西北工业大学 (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 Tensor 键的字典。

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

使用 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 数据集上达到最先进的性能。这类基准结果对于评估在处理遥感数据问题时不同建模选择的贡献非常重要。

未来工作与贡献

为了使 TorchGeo 尽可能易于使用，尤其对于没有深度学习经验的用户，仍有大量工作要做。我们计划实现这一目标的方式之一是扩展我们的教程，使其包含“编写自定义数据集”和“迁移学习”等主题，或“土地覆盖映射”和“目标检测”等任务。

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

除了这些较大的项目，我们一直在寻求添加新的数据集、数据增强变换和采样策略。如果您精通 Python 并有兴趣为 TorchGeo 做出贡献，我们非常欢迎！TorchGeo 在 MIT 许可下开源，因此您几乎可以在任何项目中使用它。

外部链接

主页：https://github.com/microsoft/torchgeo
文档：https://torchgeo.readthedocs.io/
PyPI：https://pypi.ac.cn/project/torchgeo/
论文：https://arxiv.org/abs/2111.08872

如果您喜欢 TorchGeo，请在 GitHub 上给我们点赞！如果您在工作中使用了 TorchGeo，请引用我们的论文。

致谢

我们衷心感谢所有 TorchGeo 贡献者为创建该库所做的努力，感谢微软人工智能向善计划的支持，以及 PyTorch 团队的指导。这项研究是 Blue Waters 持续千万亿级计算项目的一部分，该项目由美国国家科学基金会 (奖项 OCI-0725070 和 ACI-1238993)、伊利诺伊州以及自 2019 年 12 月起的美国国家地理空间情报局提供支持。Blue Waters 是伊利诺伊大学厄巴纳-香槟分校及其国家超级计算应用中心的联合项目。该研究部分得到 NSF 奖项 IIS-1908104、OAC-1934634 和 DBI-2021898 的支持。

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

挑战