Xarray Dataset Support

pyproject.toml

@@ -243,6 +243,8 @@ filterwarnings = [
     "ignore:ANTIALIAS is deprecated and will be removed in Pillow 10:DeprecationWarning:tensorboardX.summary",
     # https://github.com/Lightning-AI/lightning/issues/16756
     "ignore:Deprecated call to `pkg_resources.declare_namespace:DeprecationWarning",
+    # https://github.com/pydata/xarray/issues/7259
+    "ignore: numpy.ndarray size changed, may indicate binary incompatibility. Expected 16 from C header, got 96 from PyObject",
     "ignore:pkg_resources is deprecated as an API.:DeprecationWarning:lightning_utilities.core.imports",
     "ignore:Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated:DeprecationWarning:jsonargparse",
     # https://github.com/pytorch/pytorch/issues/110549

requirements/datasets.txt

@@ -7,5 +7,9 @@ pyvista==0.42.3
 radiant-mlhub==0.4.1
 rarfile==4.1
 scikit-image==0.22.0
+xarray==2023.7.0
+rioxarray==0.14.1
+xarray
+netCDF4
 scipy==1.11.3
 zipfile-deflate64==0.2.0

requirements/min-reqs.old

@@ -31,6 +31,7 @@ radiant-mlhub==0.3.0
 rarfile==4.0
 scikit-image==0.18.0
 scipy==1.6.2
+xarray
 zipfile-deflate64==0.2.0
 
 # docs

tests/data/rioxr/data.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+
+import cftime
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+SIZE = 32
+
+LATS: list[tuple[float]] = [(40, 42), (60, 62), (80, 82)]
+
+LONS: list[tuple[float]] = [(-55, -50), (-5, 5), (80, 85)]
+
+VAR_NAMES = ["zos", "tos"]
+
+DIR = "data"
+
+CF_TIME = [True, False, True]
+
+NUM_TIME_STEPS = 3
+
+
+def create_rioxr_dataset(
+    lat_min: float,
+    lat_max: float,
+    lon_min: float,
+    lon_max: float,
+    cf_time: bool,
+    var_name: str,
+    filename: str,
+):
+    # Generate x and y coordinates
+    lats = np.linspace(lat_min, lat_max, SIZE)
+    lons = np.linspace(lon_min, lon_max, SIZE)
+
+    if cf_time:
+        times = [cftime.datetime(2000, 1, i + 1) for i in range(NUM_TIME_STEPS)]
+    else:
+        times = pd.date_range(start="2000-01-01", periods=NUM_TIME_STEPS, freq="D")
+
+    # data with shape (time, x, y)
+    data = np.random.rand(len(times), len(lons), len(lats))
+
+    # Create the xarray dataset
+    ds = xr.Dataset(
+        data_vars={var_name: (("time", "x", "y"), data)},
+        coords={"x": lons, "y": lats, "time": times},
+    )
+    ds["x"].attrs["units"] = "degrees_east"
+    ds["x"].attrs["crs"] = "EPSG:4326"
+    ds["y"].attrs["units"] = "degrees_north"
+    ds["y"].attrs["crs"] = "EPSG:4326"
+    ds.to_netcdf(path=filename)
+
+
+if __name__ == "__main__":
+    if os.path.isdir(DIR):
+        shutil.rmtree(DIR)
+    os.makedirs(DIR)
+    for var_name in VAR_NAMES:
+        for lats, lons, cf_time in zip(LATS, LONS, CF_TIME):
+            path = os.path.join(DIR, f"{var_name}_{lats}_{lons}.nc")
+            create_rioxr_dataset(
+                lats[0], lats[1], lons[0], lons[1], cf_time, var_name, path
+            )

tests/datasets/test_rioxr.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+
+import pytest
+import torch
+
+from torchgeo.datasets import (
+    BoundingBox,
+    IntersectionDataset,
+    RioXarrayDataset,
+    UnionDataset,
+)
+
+pytest.importorskip("rioxarray")
+
+
+class TestRioXarrayDataset:
+    @pytest.fixture(scope="class")
+    def dataset(self) -> RioXarrayDataset:
+        root = os.path.join("tests", "data", "rioxr", "data")
+        return RioXarrayDataset(root=root, data_variables=["zos", "tos"])
+
+    def test_getitem(self, dataset: RioXarrayDataset) -> None:
+        x = dataset[dataset.bounds]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+
+    def test_and(self, dataset: RioXarrayDataset) -> None:
+        ds = dataset & dataset
+        assert isinstance(ds, IntersectionDataset)
+
+    def test_or(self, dataset: RioXarrayDataset) -> None:
+        ds = dataset | dataset
+        assert isinstance(ds, UnionDataset)
+
+    def test_invalid_query(self, dataset: RioXarrayDataset) -> None:
+        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        with pytest.raises(
+            IndexError, match="query: .* not found in index with bounds:"
+        ):
+            dataset[query]

torchgeo/datasets/__init__.py

@@ -85,6 +85,7 @@
 from .potsdam import Potsdam2D
 from .reforestree import ReforesTree
 from .resisc45 import RESISC45
+from .rioxr import RioXarrayDataset
 from .rwanda_field_boundary import RwandaFieldBoundary
 from .seasonet import SeasoNet
 from .seco import SeasonalContrastS2
@@ -240,6 +241,7 @@
     "NonGeoClassificationDataset",
     "NonGeoDataset",
     "RasterDataset",
+    "RioXarrayDataset",
     "UnionDataset",
     "VectorDataset",
     # Utilities