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

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

美国国家海洋和大气管理局 (NOAA) 拍摄的卡特里娜飓风卫星图像，摄于 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 km/px，Maxar 影像的分辨率为 30 cm/px，而无人机影像的分辨率可高达 7 mm/px。这些数据集几乎总是有时间分量，卫星重访频率为每日、每周或每两周一次。图像通常与数据集中的其他图像有重叠，需要根据地理元数据进行拼接。这些图像通常非常大（例如 10K x 10K 像素），因此无法将整张图像输入神经网络。这些数据分布在数百种不同的栅格和矢量文件格式中，如 GeoTIFF 和 ESRI Shapefile，需要专门的库（如 GDAL）来加载。

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

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

即使你在索引过程中正确地对图像进行了地理参考，如果你没有将它们投影到一个公共 CRS，你最终会得到周围有无数据值 (nodata) 的旋转图像，并且图像无法进行像素级对齐。

解决方案

目前，要在没有这两个截然不同领域专业知识的情况下同时处理深度学习模型和地理空间数据，可能会非常困难。为了解决这些挑战，我们构建了 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 的潜在用例，以展示它的使用是多么简单。

地理空间数据集和采样器

示例应用：我们结合了 A) Landsat 8 的场景和 B) 耕地数据层 (Cropland Data Layer) 的标签，尽管这些文件处于不同的 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 function) 将每个样本字典组合成一个小批次 (mini-batch) 的样本。

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）
合并不同地理位置的数据集（例如 Chesapeake 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 作为 NIR 波段。为了可视化数据，我们还要归一化图像。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 图像中，红色表示水体，黄色表示贫瘠土壤，浅绿色表示植被不健康，深绿色表示植被健康。

基准数据集

推动计算机视觉进步的驱动因素之一是标准化基准数据集（如 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 数据模块 (datamodules)，以及用于各种任务（如分类、回归和语义分割）的 训练器 (trainers)。这些数据模块展示了如何结合 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 贡献者、为我们提供支持的微软 AI for Good 项目，以及给予我们指导的 PyTorch 团队。这项研究是 Blue Waters 可持续拍字节级计算项目的一部分，该项目得到了美国国家科学基金会（奖项 OCI-0725070 和 ACI-1238993）、伊利诺伊州以及截至 2019 年 12 月的国家地理空间情报局的支持。Blue Waters 是伊利诺伊大学厄巴纳-香槟分校及其国家超级计算应用中心共同努力的成果。该研究部分得到了 NSF 拨款 IIS-1908104、OAC-1934634 和 DBI-2021898 的支持。

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

挑战