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

https://github.com/microsoft/torchgeo

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

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

挑战

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

遥感图像并非如此统一。卫星倾向于捕获多光谱（Landsat 8 有 11 个光谱带）甚至高光谱（Hyperion 有 242 个光谱带）图像，而不是简单的 RGB 图像。这些图像捕获了更广泛波长（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 旨在简化

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

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) 农田数据层标签，尽管这些文件位于不同的 EPSG 投影中。我们希望使用地理空间边界框作为索引，从这些数据集中采样 C) 和 D) 补丁。

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

在此示例中，我们展示了使用 TorchGeo 处理地理空间数据并从 Landsat 和 Cropland Data Layer (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 个样本。我们还将使用自定义的归类函数将每个样本字典组合成一个迷你批次的样本。

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()

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 数据模块，以及用于分类、回归和语义分割等各种任务的训练器。这些数据模块展示了如何整合 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)

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

在我们的预印本中，我们展示了一组结果，这些结果使用上述数据模块和训练器对 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 的部分支持。