Added substation segementation dataset

docs/api/datamodules.rst

@@ -213,6 +213,11 @@ SSL4EO-L Benchmark
 
 .. autoclass:: SSL4EOLBenchmarkDataModule
 
+Substation
+^^^^^^^^^^
+
+.. autoclass:: SubstationDataModule
+
 SustainBench Crop Yield
 ^^^^^^^^^^^^^^^^^^^^^^^
 

docs/api/datasets.rst

@@ -469,6 +469,11 @@ SSL4EO-L Benchmark
 
 .. autoclass:: SSL4EOLBenchmark
 
+Substation
+^^^^^^^^^^
+
+.. autoclass:: Substation
+
 SustainBench Crop Yield
 ^^^^^^^^^^^^^^^^^^^^^^^
 

docs/api/datasets/non_geo_datasets.csv

@@ -52,6 +52,7 @@ Dataset,Task,Source,License,# Samples,# Classes,Size (px),Resolution (m),Bands
 `SSL4EO`_-S12,T,Sentinel-1/2,"CC-BY-4.0",1M,-,264x264,10,"SAR, MSI"
 `SSL4EO-L Benchmark`_,S,Lansat & CDL,"CC0-1.0",25K,134,264x264,30,MSI
 `SSL4EO-L Benchmark`_,S,Lansat & NLCD,"CC0-1.0",25K,17,264x264,30,MSI
+`Substation`_,S,OpenStreetMap & Sentinel-2, "CC-BY-SA 2.0", 27K, 2, 228x228, 10, MSI
 `SustainBench Crop Yield`_,R,MODIS,"CC-BY-SA-4.0",11k,-,32x32,-,MSI
 `TreeSatAI`_,"C, R, S","Aerial, Sentinel-1/2",CC-BY-4.0,50K,"12, 15, 20","6, 20, 304","0.2, 10","CIR, MSI, SAR"
 `Tropical Cyclone`_,R,GOES 8--16,"CC-BY-4.0","108,110",-,256x256,4K--8K,MSI

tests/conf/substation.yaml

@@ -0,0 +1,16 @@
+model:
+  class_path: SemanticSegmentationTask
+  init_args:
+    loss: 'ce'
+    model: 'unet'
+    backbone: 'resnet18'
+    in_channels: 3
+    num_classes: 7
+    num_filters: 1
+    ignore_index: null
+data:
+  class_path: SubstationDataModule
+  init_args:
+    batch_size: 1
+  dict_kwargs:
+    root: 'tests/data/substation'

tests/data/substation/data.py

@@ -0,0 +1,83 @@
+#!/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
+
+# Parameters
+SIZE = 32  # Image dimensions
+NUM_SAMPLES = 5  # Number of samples
+np.random.seed(0)
+
+# Define directory hierarchy
+FILENAME_HIERARCHY = dict[str, 'FILENAME_HIERARCHY'] | list[str]
+
+filenames: FILENAME_HIERARCHY = {'image_stack': ['image'], 'mask': ['mask']}
+
+
+def create_file(path: str, value: str) -> None:
+    """
+    Generates .npz files for images or masks based on the path.
+
+    Args:
+    - path (str): Base path for saving files (either 'image' or 'mask').
+    """
+    for i in range(NUM_SAMPLES):
+        new_path = f'{path}_{i}.npz'
+
+        if value == 'image':
+            # Generate image data with shape (4, 13, SIZE, SIZE) for timepoints and channels
+            data = np.random.rand(4, 13, SIZE, SIZE).astype(
+                np.float32
+            )  # 4 timepoints, 13 channels
+        elif value == 'mask':
+            # Generate mask data with shape (SIZE, SIZE) with 4 classes
+            data = np.random.randint(0, 4, size=(SIZE, SIZE)).astype(np.uint8)
+
+        np.savez_compressed(new_path, arr_0=data)
+
+
+def create_directory(directory: str, hierarchy: FILENAME_HIERARCHY) -> None:
+    """
+    Recursively creates directory structure based on hierarchy and populates with data files.
+
+    Args:
+    - directory (str): Base directory for dataset.
+    - hierarchy (FILENAME_HIERARCHY): Directory and file structure.
+    """
+    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_directory(path, value)
+    else:
+        # Base case
+        for value in hierarchy:
+            path = os.path.join(directory, 'image')
+            create_file(path, value)
+
+
+if __name__ == '__main__':
+    # Generate directory structure and data
+    create_directory('.', filenames)
+
+    # Create zip archives of dataset folders
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    shutil.make_archive('image_stack', 'gztar', '.', 'image_stack')
+    shutil.make_archive('mask', 'gztar', '.', 'mask')
+
+    # Compute and print MD5 checksums for data validation
+    with open(filename_images, 'rb') as f:
+        md5_images = hashlib.md5(f.read()).hexdigest()
+        print(f'{filename_images}: {md5_images}')
+
+    with open(filename_masks, 'rb') as f:
+        md5_masks = hashlib.md5(f.read()).hexdigest()
+        print(f'{filename_masks}: {md5_masks}')

tests/datasets/test_substation.py

@@ -0,0 +1,161 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from collections.abc import Generator
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+
+from torchgeo.datasets import Substation
+
+
+class TestSubstation:
+    @pytest.fixture
+    def dataset(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> Generator[Substation, None, None]:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data')
+
+        yield Substation(
+            root=root,
+            bands=[1,2,3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+        )
+
+    @pytest.mark.parametrize(
+        'config',
+        [
+            {'bands': [1,2,3], 'use_timepoints': False, 'mask_2d': True},
+            {
+                'bands': [1,2,3],
+                'use_timepoints': True,
+                'timepoint_aggregation': 'concat',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1,2,3],
+                'use_timepoints': True,
+                'timepoint_aggregation': 'median',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+            {
+                'bands': [1,2,3],
+                'use_timepoints': True,
+                'timepoint_aggregation': 'first',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1,2,3],
+                'use_timepoints': True,
+                'timepoint_aggregation': 'random',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+            {'bands': [1,2,3], 'use_timepoints': False, 'mask_2d': False},
+            {
+                'bands': [1,2,3],
+                'use_timepoints': False,
+                'timepoint_aggregation': 'first',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1,2,3],
+                'use_timepoints': False,
+                'timepoint_aggregation': 'random',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+        ],
+    )
+    def test_getitem_semantic(self, config: dict[str, Any]) -> None:
+        root = os.path.join(os.getcwd(), 'tests', 'data')
+        dataset = Substation(root=root, **config)
+
+        x = dataset[0]
+        assert isinstance(x, dict), f'Expected dict, got {type(x)}'
+        assert isinstance(
+            x['image'], torch.Tensor
+        ), 'Expected image to be a torch.Tensor'
+        assert isinstance(x['mask'], torch.Tensor), 'Expected mask to be a torch.Tensor'
+
+    def test_len(self, dataset: Substation) -> None:
+        """Test the length of the dataset."""
+        assert len(dataset) == 5
+
+    def test_output_shape(self, dataset: Substation) -> None:
+        """Test the output shape of the dataset."""
+        x = dataset[0]
+        assert x['image'].shape == torch.Size([3, 32, 32])
+        assert x['mask'].shape == torch.Size([2, 32, 32])
+
+    def test_plot(self, dataset: Substation) -> None:
+        sample = dataset[0]
+        dataset.plot(sample, suptitle='Test')
+        plt.close()
+        dataset.plot(sample, show_titles=False)
+        plt.close()
+        sample['prediction'] = sample['mask'].clone()
+        dataset.plot(sample)
+        plt.close()
+
+    def test_already_downloaded(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test that the dataset doesn't re-download if already present."""
+        # Simulating that files are already present by copying them to the target directory
+        url_for_images = os.path.join(
+            'tests', 'data', 'substation', 'image_stack.tar.gz'
+        )
+        url_for_masks = os.path.join('tests', 'data', 'substation', 'mask.tar.gz')
+
+        # Copy files to the temporary directory to simulate already downloaded files
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        # No download should be attempted, since the files are already present
+        # Mock the _download method to simulate the behavior
+        monkeypatch.setattr(dataset, '_download', MagicMock())
+        dataset._download()  # This will now call the mocked method
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        # Mock the download_url and extract_archive functions
+        mock_download_url = MagicMock()
+        mock_extract_archive = MagicMock()
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', mock_download_url
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', mock_extract_archive
+        )
+
+        # Call the _download method
+        dataset._download()
+
+        # Check that download_url was called twice
+        mock_download_url.assert_called()
+        assert mock_download_url.call_count == 2
+
+        # Check that extract_archive was called twice
+        mock_extract_archive.assert_called()
+        assert mock_extract_archive.call_count == 2
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/trainers/test_segmentation.py

@@ -80,6 +80,7 @@ class TestSemanticSegmentationTask:
             'spacenet6',
             'ssl4eo_l_benchmark_cdl',
             'ssl4eo_l_benchmark_nlcd',
+            'substation',
             'vaihingen2d',
         ],
     )

torchgeo/datamodules/__init__.py

@@ -47,6 +47,7 @@
 from .spacenet import SpaceNet1DataModule, SpaceNet6DataModule, SpaceNetBaseDataModule
 from .ssl4eo import SSL4EOLDataModule, SSL4EOS12DataModule
 from .ssl4eo_benchmark import SSL4EOLBenchmarkDataModule
+from .substation import SubstationDataModule
 from .sustainbench_crop_yield import SustainBenchCropYieldDataModule
 from .treesatai import TreeSatAIDataModule
 from .ucmerced import UCMercedDataModule
@@ -110,6 +111,7 @@
     'SpaceNet1DataModule',
     'SpaceNet6DataModule',
     'SpaceNetBaseDataModule',
+    'SubstationDataModule',
     'SustainBenchCropYieldDataModule',
     'TreeSatAIDataModule',
     'TropicalCycloneDataModule',