Xarray Dataset Support

requirements/datasets.txt

@@ -9,4 +9,5 @@ radiant-mlhub==0.4.1
 rarfile==4.0
 scikit-image==0.21.0
 scipy==1.11.1
+xarray
 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/rioxarray/data.py

@@ -0,0 +1,65 @@
+# 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), (-60, -55), (-85, 80)]
+
+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(lats), len(lons))
+
+    # Create the xarray dataset
+    ds = xr.Dataset(
+        data_vars={var_name: (("time", "x", "y"), data)},
+        coords={"x": lats, "y": lons, "time": times},
+    )
+    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_rioxarray.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", "rioxarray", "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

@@ -82,6 +82,7 @@
 from .potsdam import Potsdam2D
 from .reforestree import ReforesTree
 from .resisc45 import RESISC45
+from .rioxarray import RioXarrayDataset
 from .seco import SeasonalContrastS2
 from .sen12ms import SEN12MS
 from .sentinel import Sentinel, Sentinel1, Sentinel2
@@ -229,6 +230,7 @@
     "NonGeoClassificationDataset",
     "NonGeoDataset",
     "RasterDataset",
+    "RioXarrayDataset",
     "UnionDataset",
     "VectorDataset",
     # Utilities

torchgeo/datasets/rioxarray.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""In-memory geographical xarray.DataArray."""
+
+import glob
+import os
+import re
+import sys
+from datetime import datetime
+from typing import Any, Callable, Optional, cast
+
+import netCDF4  # noqa: F401
+import numpy as np
+import torch
+import xarray as xr
+from rasterio.crs import CRS
+from rtree.index import Index, Property
+
+from .geo import GeoDataset
+from .utils import BoundingBox
+
+
+class RioXarrayDataset(GeoDataset):
+    """Wrapper for geographical datasets stored as Xarray Datasets.
+
+    Relies on rioxarray.
+    """
+
+    filename_glob = "*"
+    filename_regex = ".*"
+
+    def __init__(
+        self,
+        root: str,
+        data_variables: list[str],
+        crs: Optional[CRS] = None,
+        res: Optional[float] = None,
+        transforms: Optional[Callable[[dict[str, Any]], dict[str, Any]]] = None,
+    ) -> None:
+        """Initialize a new Dataset instance.
+
+        Args:
+            root: directory with nc files
+            data_variables: data variables that should be gathered from the xr_datasets
+            crs: :term:`coordinate reference system (CRS)` to warp to
+                (defaults to the CRS of dataarray)
+            res: resolution of the dataset in units of CRS
+                (defaults to the resolution of the dataarray)
+            transforms: a function/transform that takes an input sample
+                and returns a transformed version
+
+        Raises:
+            FileNotFoundError: if no files are found in ``root``
+        """
+        super().__init__(transforms)
+
+        self.root = root
+        self.data_variables = data_variables
+        self.transforms = transforms
+
+        # Create an R-tree to index the dataset
+        self.index = Index(interleaved=False, properties=Property(dimension=3))
+
+        # Populate the dataset index
+        i = 0
+        pathname = os.path.join(root, self.filename_glob)
+        filename_regex = re.compile(self.filename_regex, re.VERBOSE)
+        for filepath in glob.iglob(pathname, recursive=True):
+            match = re.match(filename_regex, os.path.basename(filepath))
+            if match is not None:
+                with xr.open_dataset(filepath, decode_times=True) as ds:
+                    if crs is None:
+                        crs = ds.rio.crs
+                    if res is None:
+                        res = ds.rio.resolution()[0]
+
+                    (minx, miny, maxx, maxy) = ds.rio.bounds()
+
+                if hasattr(ds, "time"):
+                    try:
+                        indices = ds.indexes["time"].to_datetimeindex()
+                    except AttributeError:
+                        indices = ds.indexes["time"]
+
+                    mint = indices.min().to_pydatetime().timestamp()
+                    maxt = indices.max().to_pydatetime().timestamp()
+                else:
+                    mint = 0
+                    maxt = sys.maxsize
+                coords = (minx, maxx, miny, maxy, mint, maxt)
+                self.index.insert(i, coords, filepath)
+                i += 1
+
+        if i == 0:
+            msg = f"No {self.__class__.__name__} data was found in `root='{self.root}'`"
+            raise FileNotFoundError(msg)
+
+        if not crs:
+            self._crs = "EPSG:4326"
+        else:
+            self._crs = cast(CRS, crs)
+        self.res = cast(float, res)
+
+    def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
+        """Retrieve image/mask and metadata indexed by query.
+
+        Args:
+            query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
+
+        Returns:
+            sample of image/mask and metadata at that index
+
+        Raises:
+            IndexError: if query is not found in the index
+        """
+        hits = self.index.intersection(tuple(query), objects=True)
+        items = [hit.object for hit in hits]
+
+        if not items:
+            raise IndexError(
+                f"query: {query} not found in index with bounds: {self.bounds}"
+            )
+
+        data_arrays: list["np.typing.NDArray[np.float32]"] = []
+        for item in items:
+            with xr.open_dataset(item, decode_cf=True) as ds:
+                if not ds.rio.crs:
+                    ds.rio.write_crs(self._crs, inplace=True)
+
+                # clip box ignores time dimension
+                clipped = ds.rio.clip_box(
+                    minx=query.minx, miny=query.miny, maxx=query.maxx, maxy=query.maxy
+                )
+                # select time dimension
+                if hasattr(ds, "time"):
+                    try:
+                        clipped["time"] = clipped.indexes["time"].to_datetimeindex()
+                    except AttributeError:
+                        clipped["time"] = clipped.indexes["time"]
+                    clipped = clipped.sel(
+                        time=slice(
+                            datetime.fromtimestamp(query.mint),
+                            datetime.fromtimestamp(query.maxt),
+                        )
+                    )
+                for variable in self.data_variables:
+                    if hasattr(clipped, variable):
+                        data_arrays.append(clipped[variable].data.squeeze())
+
+        sample = {"image": torch.from_numpy(np.stack(data_arrays)), "bbox": query}
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample