Add EnMAP dataset (#2543)

* Add EnMAP dataset

* Fix doc link

* Add EnMAP license

docs/api/datasets.rst

@@ -71,6 +71,11 @@ EDDMapS
 
 .. autoclass:: EDDMapS
 
+EnMAP
+^^^^^
+
+.. autoclass:: EnMAP
+
 EnviroAtlas
 ^^^^^^^^^^^
 

docs/api/datasets/geo_datasets.csv

@@ -8,6 +8,7 @@ Dataset,Type,Source,License,Size (px),Resolution (m)
 `Global Mangrove Distribution`_,Masks,"Remote Sensing, In Situ Measurements","public domain",-,3
 `Cropland Data Layer`_,Masks,Landsat,"public domain",-,30
 `EDDMapS`_,Points,Citizen Scientists,-,-,-
+`EnMAP`_,Imagery,EnMAP,`EnMAP Data License <https://www.enmap.org/data/resources/EnMAP_Data_License.pdf>`_,"1,200x1,200",30
 `EnviroAtlas`_,"Imagery, Masks","NAIP, NLCD, OpenStreetMap","CC-BY-4.0",-,1
 `Esri2020`_,Masks,Sentinel-2,"CC-BY-4.0",-,10
 `EU-DEM`_,DEM,"Aster, SRTM, Russian Topomaps","CSCDA-ESA",-,25

tests/data/enmap/data.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+
+import numpy as np
+import rasterio
+from rasterio import Affine
+from rasterio.crs import CRS
+
+SIZE = 32
+DTYPE = 'int16'
+
+np.random.seed(0)
+
+wkt = """
+PROJCS["WGS 84 / UTM zone 40N",
+    GEOGCS["WGS 84",
+        DATUM["WGS_1984",
+            SPHEROID["WGS 84",6378137,298.257223563,
+                AUTHORITY["EPSG","7030"]],
+            AUTHORITY["EPSG","6326"]],
+        PRIMEM["Greenwich",0,
+            AUTHORITY["EPSG","8901"]],
+        UNIT["degree",0.0174532925199433,
+            AUTHORITY["EPSG","9122"]],
+        AUTHORITY["EPSG","4326"]],
+    PROJECTION["Transverse_Mercator"],
+    PARAMETER["latitude_of_origin",0],
+    PARAMETER["central_meridian",57],
+    PARAMETER["scale_factor",0.9996],
+    PARAMETER["false_easting",500000],
+    PARAMETER["false_northing",0],
+    UNIT["metre",1,
+        AUTHORITY["EPSG","9001"]],
+    AXIS["Easting",EAST],
+    AXIS["Northing",NORTH],
+    AUTHORITY["EPSG","32640"]]
+"""
+
+profile = {
+    'driver': 'GTiff',
+    'dtype': DTYPE,
+    'nodata': -32768.0,
+    'width': SIZE,
+    'height': SIZE,
+    'count': 224,
+    'crs': CRS.from_wkt(wkt),
+    'transform': Affine(30.0, 0.0, 283455.0, 0.0, -30.0, 2786715.0),
+}
+
+filename = 'ENMAP01-____L2A-DT0000001053_20220611T072305Z_002_V010400_20231221T134421Z-SPECTRAL_IMAGE_COG.tiff'
+
+Z = np.random.randint(
+    np.iinfo(DTYPE).min, np.iinfo(DTYPE).max, size=(SIZE, SIZE), dtype=DTYPE
+)
+with rasterio.open(filename, 'w', **profile) as src:
+    for i in range(1, profile['count'] + 1):
+        src.write(Z, i)

tests/datasets/test_enmap.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+from rasterio.crs import CRS
+
+from torchgeo.datasets import (
+    BoundingBox,
+    DatasetNotFoundError,
+    EnMAP,
+    IntersectionDataset,
+    RGBBandsMissingError,
+    UnionDataset,
+)
+
+
+class TestEnMAP:
+    @pytest.fixture
+    def dataset(self) -> EnMAP:
+        root = os.path.join('tests', 'data', 'enmap')
+        transforms = nn.Identity()
+        return EnMAP(root, transforms=transforms)
+
+    def test_getitem(self, dataset: EnMAP) -> None:
+        x = dataset[dataset.bounds]
+        assert isinstance(x, dict)
+        assert isinstance(x['crs'], CRS)
+        assert isinstance(x['image'], torch.Tensor)
+
+    def test_len(self, dataset: EnMAP) -> None:
+        assert len(dataset) == 1
+
+    def test_and(self, dataset: EnMAP) -> None:
+        ds = dataset & dataset
+        assert isinstance(ds, IntersectionDataset)
+
+    def test_or(self, dataset: EnMAP) -> None:
+        ds = dataset | dataset
+        assert isinstance(ds, UnionDataset)
+
+    def test_plot(self, dataset: EnMAP) -> None:
+        x = dataset[dataset.bounds]
+        dataset.plot(x, suptitle='Test')
+        plt.close()
+
+    def test_plot_wrong_bands(self, dataset: EnMAP) -> None:
+        bands = ('B1', 'B2', 'B3')
+        ds = EnMAP(dataset.paths, bands=bands)
+        x = dataset[dataset.bounds]
+        with pytest.raises(
+            RGBBandsMissingError, match='Dataset does not contain some of the RGB bands'
+        ):
+            ds.plot(x)
+
+    def test_no_data(self, tmp_path: Path) -> None:
+        with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
+            EnMAP(tmp_path)
+
+    def test_invalid_query(self, dataset: EnMAP) -> 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

@@ -38,6 +38,7 @@
 from .dfc2022 import DFC2022
 from .digital_typhoon import DigitalTyphoon
 from .eddmaps import EDDMapS
+from .enmap import EnMAP
 from .enviroatlas import EnviroAtlas
 from .errors import DatasetNotFoundError, DependencyNotFoundError, RGBBandsMissingError
 from .esri2020 import Esri2020
@@ -209,6 +210,7 @@
     'DependencyNotFoundError',
     'DigitalTyphoon',
     'EDDMapS',
+    'EnMAP',
     'EnviroAtlas',
     'Esri2020',
     'EuroCrops',