SSL4EO Landsat Downstream Dataset/module CDL, NLCD

tests/data/ssl4eo_downstream_landsat/l5-l1/data.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import hashlib
+import os
+import shutil
+
+import numpy as np
+import rasterio
+from rasterio import Affine
+from rasterio.crs import CRS
+
+IMG_DIR = "ssl4eo-l5-l1-conus"
+MASK_DIR = "l5-*-2011"
+MASKS = ["cdl", "nlcd"]
+
+SUBDIRS = [("0000000", "LT05_045030_20110723"), ("0000001", "LT05_040032_20110805")]
+
+NUM_BANDS = 7
+SIZE = 32
+
+
+def create_image(path: str) -> None:
+    profile = {
+        "driver": "GTiff",
+        "dtype": "uint8",
+        "nodata": None,
+        "width": 264,
+        "height": 264,
+        "count": 7,
+        "crs": CRS.from_epsg(4326),
+        "transform": Affine(
+            0.00037672803497508636,
+            0.0,
+            -109.07063613660262,
+            0.0,
+            -0.0002554026278261721,
+            47.49838726154881,
+        ),
+        "blockysize": 1,
+        "tiled": False,
+        "compress": "lzw",
+        "interleave": "pixel",
+    }
+
+    Z = np.random.randint(low=0, high=255, size=(NUM_BANDS, SIZE, SIZE))
+
+    with rasterio.open(path, "w", **profile) as src:
+        src.write(Z)
+
+
+def create_mask(path: str) -> None:
+    profile = {
+        "driver": "GTiff",
+        "dtype": "uint8",
+        "nodata": None,
+        "width": 264,
+        "height": 264,
+        "count": 1,
+        "crs": CRS.from_epsg(4326),
+        "transform": Affine(
+            0.00037672803497508636,
+            0.0,
+            -109.07063613660262,
+            0.0,
+            -0.0002554026278261721,
+            47.49838726154881,
+        ),
+        "blockysize": 1,
+        "tiled": False,
+        "compress": "lzw",
+        "interleave": "band",
+    }
+
+    Z = np.random.randint(low=0, high=10, size=(1, SIZE, SIZE))
+
+    with rasterio.open(path, "w", **profile) as src:
+        src.write(Z)
+
+
+if __name__ == "__main__":
+    # create images
+    if os.path.isdir(IMG_DIR):
+        shutil.rmtree(IMG_DIR)
+    os.makedirs(IMG_DIR, exist_ok=True)
+
+    for subdir in SUBDIRS:
+        img_dir = os.path.join(os.getcwd(), IMG_DIR, subdir[0], subdir[1])
+        os.makedirs(img_dir)
+        create_image(os.path.join(img_dir, "all_bands.tif"))
+
+    shutil.make_archive(IMG_DIR, "gztar", ".", IMG_DIR)
+
+    with open(f"{IMG_DIR}.tar.gz", "rb") as f:
+        md5 = hashlib.md5(f.read()).hexdigest()
+        print(IMG_DIR, md5)
+
+    # create masks
+    for mask_name in MASKS:
+        mask_dir = MASK_DIR.replace("*", mask_name)
+        if os.path.isdir(mask_dir):
+            shutil.rmtree(mask_dir)
+        os.makedirs(mask_dir, exist_ok=True)
+
+        for subdir in SUBDIRS:
+            mask_path = os.path.join(os.getcwd(), mask_dir, subdir[0], subdir[1])
+            os.makedirs(mask_path)
+            create_mask(os.path.join(mask_path, f"{mask_name}_2011.tif"))
+
+        shutil.make_archive(mask_dir, "gztar", ".", mask_dir)
+
+        with open(f"{mask_dir}.tar.gz", "rb") as f:
+            md5 = hashlib.md5(f.read()).hexdigest()
+            print(mask_dir, md5)

tests/datasets/test_ssl4eo_downstream_landsat.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from itertools import product
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+from _pytest.fixtures import SubRequest
+from pytest import MonkeyPatch
+from torch.utils.data import ConcatDataset
+
+import torchgeo.datasets.utils
+from torchgeo.datasets import SSL4EODownstream
+
+
+def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
+    shutil.copy(url, root)
+
+
+class TestSSL4EODownstream:
+    @pytest.fixture(params=product([("l5-l1", 2011)], ["cdl", "nlcd"]))
+    def dataset(
+        self, monkeypatch: MonkeyPatch, tmp_path: Path, request: SubRequest
+    ) -> SSL4EODownstream:
+        root = str(tmp_path)
+        sensor_year, mask_product = request.param
+        input_sensor, year = sensor_year
+
+        img_dir = os.path.join(
+            "tests",
+            "data",
+            "ssl4eo_downstream_landsat",
+            input_sensor,
+            f"ssl4eo-{input_sensor}-conus",
+        )
+        mask_dir = os.path.join(
+            "tests",
+            "data",
+            "ssl4eo_downstream_landsat",
+            input_sensor,
+            f"{input_sensor.split('-')[0]}-{mask_product}-{year}",
+        )
+
+        shutil.copy(img_dir, root)
+        shutil.copy(mask_dir, root)
+
+        transforms = nn.Identity()
+        return SSL4EODownstream(
+            root=root,
+            input_sensor=input_sensor,
+            mask_product=mask_product,
+            split=split,
+            download=True,
+            checksum=True,
+            transforms=transforms,
+        )
+
+    def test_getitem(self, dataset: SSL4EODownstream) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+        assert isinstance(x["mask"], torch.Tensor)
+
+    def test_invalid_split(self) -> None:
+        with pytest.raises(AssertionError):
+            SSL4EODownstream(split="foo")
+
+    def test_invalid_input_sensor(self) -> None:
+        with pytest.raises(AssertionError):
+            SSL4EODownstream(split="foo")
+
+    def test_invalid_mask_product(self) -> None:
+        with pytest.raises(AssertionError):
+            SSL4EODownstream(split="foo")
+
+    def test_add(self, dataset: SSL4EODownstream) -> None:
+        ds = dataset + dataset
+        assert isinstance(ds, ConcatDataset)

torchgeo/datasets/__init__.py

@@ -105,6 +105,7 @@
     time_series_split,
 )
 from .ssl4eo import SSL4EO, SSL4EOL, SSL4EOS12
+from .ssl4eo_downstream_landsat import SSL4EODownstream
 from .sustainbench_crop_yield import SustainBenchCropYield
 from .ucmerced import UCMerced
 from .usavars import USAVars
@@ -210,6 +211,7 @@
     "SpaceNet6",
     "SpaceNet7",
     "SSL4EO",
+    "SSL4EODownstream",
     "SSL4EOL",
     "SSL4EOS12",
     "SustainBenchCropYield",