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 不仅仅是一个研究项目，它还是一个生产质量库，它使用持续集成来测试每个提交，并在各种平台（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 和农作物数据层 (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）
组合来自不同地理位置的数据集（例如，切萨皮克纽约和宾夕法尼亚州）

这些组合要求所有查询都存在于*至少一个*数据集中，并且可以使用 `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 数据模块，这些模块具有明确定义的训练-验证-测试拆分，以及用于分类、回归和语义分割等各种任务的 训练器。这些数据模块展示了如何整合 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 贡献者为创建该库所做的努力，感谢微软人工智能向善计划提供的支持，以及 PyTorch 团队的指导。这项研究是“蓝水”持续百亿亿次级计算项目的一部分，该项目由美国国家科学基金会（奖项 OCI-0725070 和 ACI-1238993）、伊利诺伊州以及自 2019 年 12 月起由国家地理空间情报局资助。“蓝水”是伊利诺伊大学厄巴纳-香槟分校及其国家超级计算应用中心共同努力的成果。这项研究部分得到了 NSF 拨款 IIS-1908104、OAC-1934634 和 DBI-2021898 的支持。

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

挑战