SSL4EO Landsat Downstream Dataset/module CDL, NLCD (#1338)

* outline downstream dataset

* revert commit

* outline of downstream dataset

* some work

* add test data

* run black

* add plot method

* new naming scheme

* fix unit tests

* add datamodule

* requested changes

* remove comments

* EOF

* battle against red circle

* fix failing tests

* add train val test split

* huggingface support

* fix failing trainer test

* running version of datamodule with trainer

* requested changes

* rename datamodules

* add colormap

* outline downstream dataset

* revert commit

* requested changes

* pass tests

* requested changes

tests/conf/ssl4eo_l_benchmark_cdl.yaml

@@ -0,0 +1,20 @@
+module:
+  _target_: torchgeo.trainers.SemanticSegmentationTask
+  loss: "ce"
+  model: "unet"
+  backbone: "resnet18"
+  weights: null
+  learning_rate: 1e-3
+  learning_rate_schedule_patience: 6
+  in_channels: 7
+  num_classes: 134
+  num_filters: 1
+  ignore_index: 0
+
+datamodule:
+  _target_: torchgeo.datamodules.SSL4EOLBenchmarkDataModule
+  root: "tests/data/ssl4eo_benchmark_landsat"
+  input_sensor: "tm_toa"
+  mask_product: "cdl"
+  batch_size: 2
+  num_workers: 0

tests/conf/ssl4eo_l_benchmark_nlcd.yaml

@@ -0,0 +1,20 @@
+module:
+  _target_: torchgeo.trainers.SemanticSegmentationTask
+  loss: "ce"
+  model: "unet"
+  backbone: "resnet18"
+  weights: null
+  learning_rate: 1e-3
+  learning_rate_schedule_patience: 6
+  in_channels: 6
+  num_classes: 17
+  num_filters: 1
+  ignore_index: 0
+
+datamodule:
+  _target_: torchgeo.datamodules.SSL4EOLBenchmarkDataModule
+  root: "tests/data/ssl4eo_benchmark_landsat"
+  input_sensor: "etm_sr"
+  mask_product: "nlcd"
+  batch_size: 2
+  num_workers: 0

tests/data/ssl4eo_benchmark_landsat/data.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import hashlib
+import os
+import shutil
+from typing import Union
+
+import numpy as np
+import rasterio
+from rasterio import Affine
+from rasterio.crs import CRS
+
+SIZE = 264
+
+np.random.seed(0)
+
+FILENAME_HIERARCHY = Union[dict[str, "FILENAME_HIERARCHY"], list[str]]
+
+filenames: FILENAME_HIERARCHY = {
+    "tm_toa": {
+        "0000001": {"LT05_172030_20010526": ["all_bands.tif"]},
+        "0000002": {"LT05_223084_20010413": ["all_bands.tif"]},
+        "0000003": {"LT05_172034_20020902": ["all_bands.tif"]},
+        "0000004": {"LT05_172034_20020903": ["all_bands.tif"]},
+        "0000005": {"LT05_172034_20020904": ["all_bands.tif"]},
+        "0000006": {"LT05_172034_20020905": ["all_bands.tif"]},
+        "0000007": {"LT05_172034_20020906": ["all_bands.tif"]},
+        "0000008": {"LT05_172034_20020907": ["all_bands.tif"]},
+        "0000009": {"LT05_172034_20020908": ["all_bands.tif"]},
+        "0000010": {"LT05_172034_20020909": ["all_bands.tif"]},
+    },
+    "etm_sr": {
+        "0000001": {"LE07_172030_20010526": ["all_bands.tif"]},
+        "0000002": {"LE07_223084_20010413": ["all_bands.tif"]},
+        "0000003": {"LE07_172034_20020902": ["all_bands.tif"]},
+        "0000004": {"LE07_172034_20020903": ["all_bands.tif"]},
+        "0000005": {"LE07_172034_20020904": ["all_bands.tif"]},
+        "0000006": {"LE07_172034_20020905": ["all_bands.tif"]},
+        "0000007": {"LE07_172034_20020906": ["all_bands.tif"]},
+        "0000008": {"LE07_172034_20020907": ["all_bands.tif"]},
+        "0000009": {"LE07_172034_20020908": ["all_bands.tif"]},
+        "0000010": {"LE07_172034_20020909": ["all_bands.tif"]},
+    },
+    "etm_toa": {
+        "0000001": {"LE07_172030_20010526": ["all_bands.tif"]},
+        "0000002": {"LE07_223084_20010413": ["all_bands.tif"]},
+        "0000003": {"LE07_172034_20020902": ["all_bands.tif"]},
+        "0000004": {"LE07_172034_20020903": ["all_bands.tif"]},
+        "0000005": {"LE07_172034_20020904": ["all_bands.tif"]},
+        "0000006": {"LE07_172034_20020905": ["all_bands.tif"]},
+        "0000007": {"LE07_172034_20020906": ["all_bands.tif"]},
+        "0000008": {"LE07_172034_20020907": ["all_bands.tif"]},
+        "0000009": {"LE07_172034_20020908": ["all_bands.tif"]},
+        "0000010": {"LE07_172034_20020909": ["all_bands.tif"]},
+    },
+    "oli_tirs_toa": {
+        "0000001": {"LC08_172030_20010526": ["all_bands.tif"]},
+        "0000002": {"LC08_223084_20010413": ["all_bands.tif"]},
+        "0000003": {"LC08_172034_20020902": ["all_bands.tif"]},
+        "0000004": {"LC08_172034_20020903": ["all_bands.tif"]},
+        "0000005": {"LC08_172034_20020904": ["all_bands.tif"]},
+        "0000006": {"LC08_172034_20020905": ["all_bands.tif"]},
+        "0000007": {"LC08_172034_20020906": ["all_bands.tif"]},
+        "0000008": {"LC08_172034_20020907": ["all_bands.tif"]},
+        "0000009": {"LC08_172034_20020908": ["all_bands.tif"]},
+        "0000010": {"LC08_172034_20020909": ["all_bands.tif"]},
+    },
+    "oli_sr": {
+        "0000001": {"LC08_172030_20010526": ["all_bands.tif"]},
+        "0000002": {"LC08_223084_20010413": ["all_bands.tif"]},
+        "0000003": {"LC08_172034_20020902": ["all_bands.tif"]},
+        "0000004": {"LC08_172034_20020903": ["all_bands.tif"]},
+        "0000005": {"LC08_172034_20020904": ["all_bands.tif"]},
+        "0000006": {"LC08_172034_20020905": ["all_bands.tif"]},
+        "0000007": {"LC08_172034_20020906": ["all_bands.tif"]},
+        "0000008": {"LC08_172034_20020907": ["all_bands.tif"]},
+        "0000009": {"LC08_172034_20020908": ["all_bands.tif"]},
+        "0000010": {"LC08_172034_20020909": ["all_bands.tif"]},
+    },
+}
+
+num_bands = {"tm_toa": 7, "etm_sr": 6, "etm_toa": 9, "oli_tirs_toa": 11, "oli_sr": 7}
+years = {"tm": 2011, "etm": 2019, "oli": 2019}
+
+
+def create_image(path: str) -> None:
+    profile = {
+        "driver": "GTiff",
+        "dtype": "uint8",
+        "nodata": None,
+        "width": SIZE,
+        "height": SIZE,
+        "count": num_bands["_".join(path.split(os.sep)[1].split("_")[2:][:-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": "pixel",
+    }
+
+    Z = np.random.randint(low=0, high=255, size=(SIZE, SIZE))
+
+    with rasterio.open(path, "w", **profile) as src:
+        for i in src.indexes:
+            src.write(Z, i)
+
+
+def create_mask(path: str) -> None:
+    profile = {
+        "driver": "GTiff",
+        "dtype": "uint8",
+        "nodata": None,
+        "width": SIZE,
+        "height": SIZE,
+        "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)
+
+
+def create_img_directory(directory: str, hierarchy: FILENAME_HIERARCHY) -> None:
+    if isinstance(hierarchy, dict):
+        # Recursive case
+        for key, value in hierarchy.items():
+            if any([x in key for x in filenames.keys()]):
+                key = f"ssl4eo_l_{key}_benchmark"
+            path = os.path.join(directory, key)
+            os.makedirs(path, exist_ok=True)
+            create_img_directory(path, value)
+    else:
+        # Base case
+        for value in hierarchy:
+            path = os.path.join(directory, value)
+            create_image(path)
+
+
+def create_mask_directory(
+    directory: str, hierarchy: FILENAME_HIERARCHY, mask_product: str
+) -> None:
+    if isinstance(hierarchy, dict):
+        # Recursive case
+        for key, value in hierarchy.items():
+            path = os.path.join(directory, key)
+            os.makedirs(path, exist_ok=True)
+            create_mask_directory(path, value, mask_product)
+    else:
+        # Base case
+        for value in hierarchy:
+            path = os.path.join(directory, value)
+            year = years[path.split(os.sep)[1].split("_")[2]]
+            create_mask(path.replace("all_bands", f"{mask_product}_{year}"))
+
+
+def create_tarballs(directories) -> None:
+    for directory in directories:
+        # Create tarballs
+        shutil.make_archive(directory, "gztar", ".", directory)
+
+        # Compute checksums
+        with open(f"{directory}.tar.gz", "rb") as f:
+            md5 = hashlib.md5(f.read()).hexdigest()
+            print(directory, md5)
+
+
+if __name__ == "__main__":
+    # image directories
+    create_img_directory(".", filenames)
+    directories = filenames.keys()
+    directories = [f"ssl4eo_l_{key}_benchmark" for key in directories]
+    create_tarballs(directories)
+
+    # mask directory cdl
+    mask_keep = ["tm_toa", "etm_sr", "oli_sr"]
+    mask_filenames = {
+        f"ssl4eo_l_{key.split('_')[0]}_cdl": val
+        for key, val in filenames.items()
+        if key in mask_keep
+    }
+    create_mask_directory(".", mask_filenames, "cdl")
+    directories = mask_filenames.keys()
+    create_tarballs(directories)
+
+    # mask directory nlcd
+    mask_filenames = {
+        f"ssl4eo_l_{key.split('_')[0]}_nlcd": val
+        for key, val in filenames.items()
+        if key in mask_keep
+    }
+    create_mask_directory(".", mask_filenames, "nlcd")
+    directories = mask_filenames.keys()
+    create_tarballs(directories)