Repository: microsoft/torchgeo
Pull request #2352: Added substation segementation dataset
Author: login: rijuld
Labels: name: documentation
description: Improvements or additions to documentation
name: datasets
description: Geospatial or benchmark datasets
name: testing
description: Continuous integration testing
name: datamodules
description: PyTorch Lightning datamodules
Created: 2024-10-17T15:01:05Z | Merged: 2025-03-12T16:38:09Z



Diff:
diff --git a/docs/api/datamodules.rst b/docs/api/datamodules.rst
index 3569bbfef8c..97ee3b690e7 100644
--- a/docs/api/datamodules.rst
+++ b/docs/api/datamodules.rst
@@ -218,6 +218,11 @@ SSL4EO-L Benchmark
 
 .. autoclass:: SSL4EOLBenchmarkDataModule
 
+Substation
+^^^^^^^^^^
+
+.. autoclass:: SubstationDataModule
+
 SustainBench Crop Yield
 ^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst
index 8147863f095..fc3aef563de 100644
--- a/docs/api/datasets.rst
+++ b/docs/api/datasets.rst
@@ -488,6 +488,11 @@ SSL4EO-L Benchmark
 
 .. autoclass:: SSL4EOLBenchmark
 
+Substation
+^^^^^^^^^^
+
+.. autoclass:: Substation
+
 SustainBench Crop Yield
 ^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/docs/api/datasets/non_geo_datasets.csv b/docs/api/datasets/non_geo_datasets.csv
index 1defcb032bd..23dbaffbf8f 100644
--- a/docs/api/datasets/non_geo_datasets.csv
+++ b/docs/api/datasets/non_geo_datasets.csv
@@ -54,6 +54,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-4.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
diff --git a/tests/conf/substation.yaml b/tests/conf/substation.yaml
new file mode 100644
index 00000000000..cd045d6d728
--- /dev/null
+++ b/tests/conf/substation.yaml
@@ -0,0 +1,20 @@
+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'
+    bands: [0, 1, 2]
+    num_of_timepoints: 4
+    timepoint_aggregation: 'median'
+    mask_2d: False
diff --git a/tests/data/substation/data.py b/tests/data/substation/data.py
new file mode 100644
index 00000000000..975e2cbc0f4
--- /dev/null
+++ b/tests/data/substation/data.py
@@ -0,0 +1,82 @@
+#!/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 Literal
+
+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: Literal['image', 'mask']) -> None:
+    """Generates .npz files for images or masks based on the path.
+
+    Args:
+        path: Base path for saving files.
+        value: Type of file, 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)
+        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: Base directory for dataset.
+        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}')
diff --git a/tests/data/substation/image_stack.tar.gz b/tests/data/substation/image_stack.tar.gz
new file mode 100644
index 00000000000..165a3b4f7d1
Binary files /dev/null and b/tests/data/substation/image_stack.tar.gz differ
diff --git a/tests/data/substation/image_stack/image_0.npz b/tests/data/substation/image_stack/image_0.npz
new file mode 100644
index 00000000000..9f8faabbd30
Binary files /dev/null and b/tests/data/substation/image_stack/image_0.npz differ
diff --git a/tests/data/substation/image_stack/image_1.npz b/tests/data/substation/image_stack/image_1.npz
new file mode 100644
index 00000000000..f972ad3e3af
Binary files /dev/null and b/tests/data/substation/image_stack/image_1.npz differ
diff --git a/tests/data/substation/image_stack/image_2.npz b/tests/data/substation/image_stack/image_2.npz
new file mode 100644
index 00000000000..05a7d23f576
Binary files /dev/null and b/tests/data/substation/image_stack/image_2.npz differ
diff --git a/tests/data/substation/image_stack/image_3.npz b/tests/data/substation/image_stack/image_3.npz
new file mode 100644
index 00000000000..9953c5cb12e
Binary files /dev/null and b/tests/data/substation/image_stack/image_3.npz differ
diff --git a/tests/data/substation/image_stack/image_4.npz b/tests/data/substation/image_stack/image_4.npz
new file mode 100644
index 00000000000..e88095255e9
Binary files /dev/null and b/tests/data/substation/image_stack/image_4.npz differ
diff --git a/tests/data/substation/mask.tar.gz b/tests/data/substation/mask.tar.gz
new file mode 100644
index 00000000000..c3cfb7547a7
Binary files /dev/null and b/tests/data/substation/mask.tar.gz differ
diff --git a/tests/data/substation/mask/image_0.npz b/tests/data/substation/mask/image_0.npz
new file mode 100644
index 00000000000..0a6123223e4
Binary files /dev/null and b/tests/data/substation/mask/image_0.npz differ
diff --git a/tests/data/substation/mask/image_1.npz b/tests/data/substation/mask/image_1.npz
new file mode 100644
index 00000000000..ac389454a58
Binary files /dev/null and b/tests/data/substation/mask/image_1.npz differ
diff --git a/tests/data/substation/mask/image_2.npz b/tests/data/substation/mask/image_2.npz
new file mode 100644
index 00000000000..1df3295e717
Binary files /dev/null and b/tests/data/substation/mask/image_2.npz differ
diff --git a/tests/data/substation/mask/image_3.npz b/tests/data/substation/mask/image_3.npz
new file mode 100644
index 00000000000..ccaca48de02
Binary files /dev/null and b/tests/data/substation/mask/image_3.npz differ
diff --git a/tests/data/substation/mask/image_4.npz b/tests/data/substation/mask/image_4.npz
new file mode 100644
index 00000000000..7829fd98e08
Binary files /dev/null and b/tests/data/substation/mask/image_4.npz differ
diff --git a/tests/datasets/test_substation.py b/tests/datasets/test_substation.py
new file mode 100644
index 00000000000..80bbf5af2b4
--- /dev/null
+++ b/tests/datasets/test_substation.py
@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+from pytest import MonkeyPatch
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'bands': [1, 2, 3], 'mask_2d': True},
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'concat',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'median',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'first',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'random',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+            {'bands': [1, 2, 3], 'mask_2d': False},
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'first',
+                'num_of_timepoints': 4,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'timepoint_aggregation': 'random',
+                'num_of_timepoints': 4,
+                'mask_2d': True,
+            },
+        ]
+    )
+    def dataset(self, request: pytest.FixtureRequest, tmp_path: Path) -> Substation:
+        """Fixture for the Substation with parameterization."""
+        root = os.path.join('tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(root, transforms=transforms, **request.param)
+
+    def test_getitem(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], torch.Tensor)
+        assert len(dataset) == 5
+
+        match dataset.timepoint_aggregation:
+            case 'concat':
+                assert x['image'].shape == torch.Size([12, 32, 32])
+            case 'median':
+                assert x['image'].shape == torch.Size([3, 32, 32])
+            case 'first' | 'random':
+                assert x['image'].shape == torch.Size([3, 32, 32])
+            case _:
+                assert x['image'].shape == torch.Size(
+                    [dataset.num_of_timepoints, 3, 32, 32]
+                )
+
+        if dataset.mask_2d:
+            assert x['mask'].shape == torch.Size([2, 32, 32])
+        else:
+            assert x['mask'].shape == torch.Size([32, 32])
+
+    def test_download(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> None:
+        url = os.path.join('tests', 'data', 'substation')
+        filename = Substation.filename_images
+        maskname = Substation.filename_masks
+        monkeypatch.setattr(Substation, 'url_for_images', os.path.join(url, filename))
+        monkeypatch.setattr(Substation, 'url_for_masks', os.path.join(url, maskname))
+        Substation(tmp_path, download=True)
+
+    def test_extract(self, tmp_path: Path) -> None:
+        root = os.path.join('tests', 'data', 'substation')
+        filename = Substation.filename_images
+        maskname = Substation.filename_masks
+        shutil.copyfile(os.path.join(root, filename), tmp_path / filename)
+        shutil.copyfile(os.path.join(root, maskname), tmp_path / maskname)
+        Substation(tmp_path)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
+            Substation(tmp_path)
+
+    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()
diff --git a/tests/trainers/test_segmentation.py b/tests/trainers/test_segmentation.py
index 5b5bc1f867d..278d13bd765 100644
--- a/tests/trainers/test_segmentation.py
+++ b/tests/trainers/test_segmentation.py
@@ -81,6 +81,7 @@ class TestSemanticSegmentationTask:
             'spacenet6',
             'ssl4eo_l_benchmark_cdl',
             'ssl4eo_l_benchmark_nlcd',
+            'substation',
             'vaihingen2d',
         ],
     )
diff --git a/torchgeo/datamodules/__init__.py b/torchgeo/datamodules/__init__.py
index be49fa97463..ade99c7c9e0 100644
--- a/torchgeo/datamodules/__init__.py
+++ b/torchgeo/datamodules/__init__.py
@@ -48,6 +48,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
@@ -112,6 +113,7 @@
     'SpaceNet1DataModule',
     'SpaceNet6DataModule',
     'SpaceNetBaseDataModule',
+    'SubstationDataModule',
     'SustainBenchCropYieldDataModule',
     'TreeSatAIDataModule',
     'TropicalCycloneDataModule',
diff --git a/torchgeo/datamodules/substation.py b/torchgeo/datamodules/substation.py
new file mode 100644
index 00000000000..e763905cbd2
--- /dev/null
+++ b/torchgeo/datamodules/substation.py
@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            size: Size of the input images.
+            **kwargs: Additional keyword arguments passed to
+                :class:`~torchgeo.datasets.Substation`.
+        """
+        super().__init__(Substation, batch_size, num_workers, **kwargs)
+        self.val_split_pct = val_split_pct
+        self.test_split_pct = test_split_pct
+        self.train_aug = K.AugmentationSequential(
+            K.Resize(size),
+            K.RandomHorizontalFlip(p=0.5),
+            K.RandomVerticalFlip(p=0.5),
+            data_keys=None,
+            keepdim=True,
+        )
+        self.aug = K.AugmentationSequential(
+            K.Resize(size), data_keys=None, keepdim=True
+        )
+
+    def setup(self, stage: str) -> None:
+        """Set up datasets.
+
+        Args:
+            stage: One of 'fit', 'validate', 'test', or 'predict'.
+        """
+        dataset = Substation(**self.kwargs)
+
+        generator = torch.Generator().manual_seed(0)
+        total_len = len(dataset)
+        val_len = int(total_len * self.val_split_pct)
+        test_len = int(total_len * self.test_split_pct)
+        train_len = total_len - val_len - test_len
+
+        self.train_dataset, self.val_dataset, self.test_dataset = random_split(
+            dataset, [train_len, val_len, test_len], generator
+        )
diff --git a/torchgeo/datasets/__init__.py b/torchgeo/datasets/__init__.py
index 1d644c6fc69..6239d525440 100644
--- a/torchgeo/datasets/__init__.py
+++ b/torchgeo/datasets/__init__.py
@@ -135,6 +135,7 @@
 )
 from .ssl4eo import SSL4EO, SSL4EOL, SSL4EOS12
 from .ssl4eo_benchmark import SSL4EOLBenchmark
+from .substation import Substation
 from .sustainbench_crop_yield import SustainBenchCropYield
 from .treesatai import TreeSatAI
 from .ucmerced import UCMerced
@@ -284,6 +285,7 @@
     'SpaceNet6',
     'SpaceNet7',
     'SpaceNet8',
+    'Substation',
     'SustainBenchCropYield',
     'TreeSatAI',
     'TropicalCyclone',
diff --git a/torchgeo/datasets/substation.py b/torchgeo/datasets/substation.py
new file mode 100644
index 00000000000..8adb5c97c9a
--- /dev/null
+++ b/torchgeo/datasets/substation.py
@@ -0,0 +1,266 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+from typing import Literal
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Substation dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSM) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits).
+
+    Dataset Format:
+
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+    md5_images = '948706609864d0283f74ee7015f9d032'
+    md5_masks = 'baa369ececdc2ff80e6ba2b4c7fe147c'
+
+    def __init__(
+        self,
+        root: Path = 'data',
+        bands: Sequence[int] = tuple(range(13)),
+        mask_2d: bool = True,
+        num_of_timepoints: int = 4,
+        timepoint_aggregation: Literal['concat', 'median', 'first', 'random']
+        | None = 'concat',
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        download: bool = False,
+        checksum: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            num_of_timepoints: Number of timepoints to use for each image.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            transforms: A transform takes input sample and returns a transformed version.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.num_of_timepoints = num_of_timepoints
+        self.timepoint_aggregation = timepoint_aggregation
+        self.transforms = transforms
+        self.download = download
+        self.checksum = checksum
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self._verify()
+        self.image_filenames = pd.Series(sorted(os.listdir(self.image_dir)))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if image.shape[0] < self.num_of_timepoints:
+            # Padding: cycle through existing timepoints
+            padded_images = []
+            for i in range(self.num_of_timepoints):
+                padded_images.append(image[i % image.shape[0]])
+            image = np.stack(padded_images)
+        elif image.shape[0] > self.num_of_timepoints:
+            # Removal: take the most recent timepoints
+            image = image[-self.num_of_timepoints :]
+
+        match self.timepoint_aggregation:
+            case 'concat':
+                # (num_of_timepoints*channels, h, w)
+                image = np.reshape(image, (-1, image.shape[2], image.shape[3]))
+            case 'median':
+                image = np.median(image, axis=0)
+            case 'first':
+                image = image[0]
+            case 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).long()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+        mask = mask.squeeze()
+
+        sample = {'image': image, 'mask': mask}
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(
+        self,
+        sample: dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: str | None = None,
+    ) -> Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional string to use as a suptitle
+
+        Returns:
+            A matplotlib Figure containing the rendered sample.
+        """
+        ncols = 2
+        shape_of_image = sample['image'].shape
+        if len(shape_of_image) == 4:
+            # Plot the first timepoint
+            image = sample['image'][0][:3].permute(1, 2, 0).cpu().numpy()
+        else:
+            image = sample['image'][:3].permute(1, 2, 0).cpu().numpy()
+        image = image / 255.0
+
+        if self.mask_2d:
+            mask = sample['mask'][0].squeeze(0).cpu().numpy()
+        else:
+            mask = sample['mask'].cpu().numpy()
+        showing_predictions = 'prediction' in sample
+        if showing_predictions:
+            prediction = sample['prediction'].cpu().numpy()
+            if self.mask_2d:
+                prediction = prediction[0]
+            ncols = 3
+
+        fig, axs = plt.subplots(ncols=ncols, figsize=(4 * ncols, 4))
+        axs[0].imshow(image)
+        axs[0].axis('off')
+        axs[1].imshow(mask, cmap='gray', interpolation='none')
+        axs[1].axis('off')
+
+        if show_titles:
+            axs[0].set_title('Image')
+            axs[1].set_title('Mask')
+
+        if showing_predictions:
+            axs[2].imshow(prediction, cmap='gray', interpolation='none')
+            axs[2].axis('off')
+            if show_titles:
+                axs[2].set_title('Prediction')
+
+        if suptitle:
+            fig.suptitle(suptitle)
+
+        return fig
+
+    def _extract(self) -> None:
+        """Extract the dataset."""
+        img_pathname = os.path.join(self.root, self.filename_images)
+        extract_archive(img_pathname)
+
+        mask_pathname = os.path.join(self.root, self.filename_masks)
+        extract_archive(mask_pathname)
+
+    def _verify(self) -> None:
+        """Verify the integrity of the dataset."""
+        # Check if the extracted files already exist
+        image_path = os.path.join(self.image_dir, '*.npz')
+        mask_path = os.path.join(self.mask_dir, '*.npz')
+        if glob.glob(image_path) and glob.glob(mask_path):
+            return
+
+        # Check if the tar.gz files for images and masks have already been downloaded
+        image_exists = os.path.exists(os.path.join(self.root, self.filename_images))
+        mask_exists = os.path.exists(os.path.join(self.root, self.filename_masks))
+        if image_exists and mask_exists:
+            self._extract()
+            return
+
+        # If dataset files are missing and download is not allowed, raise an error
+        if not self.download:
+            raise DatasetNotFoundError(self)
+
+        # Download and extract the dataset
+        self._download()
+        self._extract()
+
+    def _download(self) -> None:
+        """Download the dataset and extract it."""
+        # Download and verify images
+        download_url(
+            self.url_for_images,
+            self.root,
+            filename=self.filename_images,
+            md5=self.md5_images if self.checksum else None,
+        )
+        extract_archive(os.path.join(self.root, self.filename_images), self.root)
+
+        # Download and verify masks
+        download_url(
+            self.url_for_masks,
+            self.root,
+            filename=self.filename_masks,
+            md5=self.md5_masks if self.checksum else None,
+        )
+        extract_archive(os.path.join(self.root, self.filename_masks), self.root)


Reviews:
user: login: nilsleh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2024-10-24T07:48:17Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2024-10-24T12:27:57Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2024-10-30T17:02:51Z
user: login: nilsleh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2024-11-07T16:26:33Z
user: login: adamjstewart
body: Not a bad data loader, it just doesn't match a single other data loader in TorchGeo. I highly recommend looking at some of the 80+ existing data loaders and unit tests for those data loaders already builtin before adding a new one from scratch. Especially for unit testing, you can probably just copy-n-paste most of the existing test code for a similar dataset.
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2024-11-10T12:40:34Z
user: login: adamjstewart
body: Most of the download/extract logic isn't tested yet, and the datamodule also has no tests. To test the datamodule, you can add a `tests/conf/substation.yaml` file and add 1 line to `tests/trainers/test_segmentation.py`.
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-12-12T00:54:45Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-01-19T12:23:13Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-01-19T14:45:05Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-01-19T15:41:52Z
user: login: adamjstewart
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-01-26T15:03:40Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-02-05T21:42:09Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-02-06T13:22:41Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T09:08:09Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T09:09:46Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T09:19:10Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T09:19:52Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-02-27T14:49:47Z
user: login: adamjstewart
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-02-28T09:36:42Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T10:20:24Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T10:21:55Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T11:13:58Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-06T12:37:35Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T16:06:39Z
user: login: rijuld
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T16:09:04Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T16:18:53Z
user: login: adamjstewart
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-12T16:20:22Z

Review comments:
diff_hunk: @@ -0,0 +1,196 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any, NamedTuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Args(NamedTuple):
+    """NamedTuple containing configuration for the dataset."""
+    data_dir: str
+    in_channels: int
+    use_timepoints: bool
+    normalizing_type: str
+    normalizing_factor: np.ndarray[Any, np.float64]
+    means: np.ndarray[Any, np.float64]
+    stds: np.ndarray[Any, np.float64]
+    mask_2d: bool
+    model_type: str
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+    
+    It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    """
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = ''
+    url_for_masks = ''
+
+    def __init__(self, args: Args, image_files: list[str], geo_transforms: Any = None, 
+                 color_transforms: Any = None, image_resize: Any = None, mask_resize: Any = None) -> None:
+        """Initialize the dataset with the provided parameters.
+
+        Args:
+            args: Arguments that contain configuration information such as data_dir, in_channels, etc.
+            image_files: List of image filenames for the dataset.
+            geo_transforms: Geometric transformations to apply on the images and masks.
+            color_transforms: Color transformations to apply on the images.
+            image_resize: Resize transformation for images.
+            mask_resize: Resize transformation for masks.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(self.data_dir, 'image_stack')
+        self.mask_dir = os.path.join(self.data_dir, 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+        # Check if the dataset is available or needs to be downloaded
+        self._verify()
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        mask = np.load(mask_path)['arr_0']
+
+        # Standardize images
+        image = self._normalize_image(image)
+
+        # Handle multiple channels and timepoints
+        image = self._handle_channels_and_timepoints(image)
+
+        # Process mask
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+        mask = torch.from_numpy(mask).float().unsqueeze(dim=0)
+
+        # Apply geo and color transformations
+        image, mask = self._apply_transforms(image, mask)
+
+        return {"image": image, "mask": mask}
+
+    def __len__(self) -> int:
+        """Return the length of the dataset.
+
+        Returns:
+            The number of items in the dataset.
+        """
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
+        """Plot a random image and mask from the dataset."""
+        index = np.random.randint(0, self.__len__())
+        data = self.__getitem__(index)
+        image, mask = data["image"], data["mask"]
+        fig, axs = plt.subplots(1, 2, figsize=(15, 15))
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(mask.permute(1, 2, 0), alpha=0.5, cmap='gray')
+
+    def _normalize_image(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Normalize the image based on the selected normalizing type.
+
+        Args:
+            image: The image to normalize.
+
+        Returns:
+            Normalized image as a torch tensor.
+        """
+        if self.normalizing_type == 'percentile':
+            image = (image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            image = np.clip(image, 0, 1)
+        return torch.from_numpy(image)
+
+    def _handle_channels_and_timepoints(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Handle channels and timepoints in the image.
+
+        Args:
+            image: The image to process.
+
+        Returns:
+            Processed image as a torch tensor.
+        """
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            image = image[:4, :, :, :] if self.use_timepoints else image[0]
+        return torch.from_numpy(image)
+
+    def _apply_transforms(self, image: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
path: torchgeo/datasets/substation_seg.py
user: login: nilsleh
body: Hi @rijuld, thank you for contributing this dataset. As another pointer, torchgeo datasets usually have an accompanying datamodule that defines things like the train/val/test split, but also common data augmentations, like flips, color augmentations etc through the kornia package. So in essence, torchgeo datasets simply load a particular sample and the augmentations are applied on GPU over the batch. 

For example in this dataset, the `getitem` [method](https://github.com/zsyOAOA/ResShift/blob/2bc69fdd39827385c91db49cb1d155c69d5e85e8/trainer.py#L533) loads the image and mask, and then we have a corresponding [datamodule](https://github.com/microsoft/torchgeo/blob/f2ed33a4079f632a107096df8efdebe711a0bb91/torchgeo/datamodules/caffe.py#L43) where we define augmentations like resizing and others, which will automatically be applied with a lightning training setup. This helps streamlining the datasets and keep them "minimal" and also make use of existing augmentation implementations like Kornia.

Let me know if I can help with any further questions.
 
created_at: 2024-10-24T07:48:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,196 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any, NamedTuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Args(NamedTuple):
+    """NamedTuple containing configuration for the dataset."""
+    data_dir: str
+    in_channels: int
+    use_timepoints: bool
+    normalizing_type: str
+    normalizing_factor: np.ndarray[Any, np.float64]
+    means: np.ndarray[Any, np.float64]
+    stds: np.ndarray[Any, np.float64]
+    mask_2d: bool
+    model_type: str
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+    
+    It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    """
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = ''
+    url_for_masks = ''
+
+    def __init__(self, args: Args, image_files: list[str], geo_transforms: Any = None, 
+                 color_transforms: Any = None, image_resize: Any = None, mask_resize: Any = None) -> None:
+        """Initialize the dataset with the provided parameters.
+
+        Args:
+            args: Arguments that contain configuration information such as data_dir, in_channels, etc.
+            image_files: List of image filenames for the dataset.
+            geo_transforms: Geometric transformations to apply on the images and masks.
+            color_transforms: Color transformations to apply on the images.
+            image_resize: Resize transformation for images.
+            mask_resize: Resize transformation for masks.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(self.data_dir, 'image_stack')
+        self.mask_dir = os.path.join(self.data_dir, 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+        # Check if the dataset is available or needs to be downloaded
+        self._verify()
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        mask = np.load(mask_path)['arr_0']
+
+        # Standardize images
+        image = self._normalize_image(image)
+
+        # Handle multiple channels and timepoints
+        image = self._handle_channels_and_timepoints(image)
+
+        # Process mask
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+        mask = torch.from_numpy(mask).float().unsqueeze(dim=0)
+
+        # Apply geo and color transformations
+        image, mask = self._apply_transforms(image, mask)
+
+        return {"image": image, "mask": mask}
+
+    def __len__(self) -> int:
+        """Return the length of the dataset.
+
+        Returns:
+            The number of items in the dataset.
+        """
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
+        """Plot a random image and mask from the dataset."""
+        index = np.random.randint(0, self.__len__())
+        data = self.__getitem__(index)
+        image, mask = data["image"], data["mask"]
+        fig, axs = plt.subplots(1, 2, figsize=(15, 15))
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(mask.permute(1, 2, 0), alpha=0.5, cmap='gray')
+
+    def _normalize_image(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Normalize the image based on the selected normalizing type.
+
+        Args:
+            image: The image to normalize.
+
+        Returns:
+            Normalized image as a torch tensor.
+        """
+        if self.normalizing_type == 'percentile':
+            image = (image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            image = np.clip(image, 0, 1)
+        return torch.from_numpy(image)
+
+    def _handle_channels_and_timepoints(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Handle channels and timepoints in the image.
+
+        Args:
+            image: The image to process.
+
+        Returns:
+            Processed image as a torch tensor.
+        """
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            image = image[:4, :, :, :] if self.use_timepoints else image[0]
+        return torch.from_numpy(image)
+
+    def _apply_transforms(self, image: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
path: torchgeo/datasets/substation_seg.py
user: login: rijuld
body: Hi @nilsleh , thank you for the detailed explanation!

That makes perfect sense. I will try to make this minimal, implement this today and reach out if I have any further questions.

Thanks again!
created_at: 2024-10-24T12:27:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,196 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any, NamedTuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Args(NamedTuple):
+    """NamedTuple containing configuration for the dataset."""
+    data_dir: str
+    in_channels: int
+    use_timepoints: bool
+    normalizing_type: str
+    normalizing_factor: np.ndarray[Any, np.float64]
+    means: np.ndarray[Any, np.float64]
+    stds: np.ndarray[Any, np.float64]
+    mask_2d: bool
+    model_type: str
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+    
+    It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    """
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = ''
+    url_for_masks = ''
+
+    def __init__(self, args: Args, image_files: list[str], geo_transforms: Any = None, 
+                 color_transforms: Any = None, image_resize: Any = None, mask_resize: Any = None) -> None:
+        """Initialize the dataset with the provided parameters.
+
+        Args:
+            args: Arguments that contain configuration information such as data_dir, in_channels, etc.
+            image_files: List of image filenames for the dataset.
+            geo_transforms: Geometric transformations to apply on the images and masks.
+            color_transforms: Color transformations to apply on the images.
+            image_resize: Resize transformation for images.
+            mask_resize: Resize transformation for masks.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(self.data_dir, 'image_stack')
+        self.mask_dir = os.path.join(self.data_dir, 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+        # Check if the dataset is available or needs to be downloaded
+        self._verify()
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        mask = np.load(mask_path)['arr_0']
+
+        # Standardize images
+        image = self._normalize_image(image)
+
+        # Handle multiple channels and timepoints
+        image = self._handle_channels_and_timepoints(image)
+
+        # Process mask
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+        mask = torch.from_numpy(mask).float().unsqueeze(dim=0)
+
+        # Apply geo and color transformations
+        image, mask = self._apply_transforms(image, mask)
+
+        return {"image": image, "mask": mask}
+
+    def __len__(self) -> int:
+        """Return the length of the dataset.
+
+        Returns:
+            The number of items in the dataset.
+        """
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
+        """Plot a random image and mask from the dataset."""
+        index = np.random.randint(0, self.__len__())
+        data = self.__getitem__(index)
+        image, mask = data["image"], data["mask"]
+        fig, axs = plt.subplots(1, 2, figsize=(15, 15))
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(mask.permute(1, 2, 0), alpha=0.5, cmap='gray')
+
+    def _normalize_image(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Normalize the image based on the selected normalizing type.
+
+        Args:
+            image: The image to normalize.
+
+        Returns:
+            Normalized image as a torch tensor.
+        """
+        if self.normalizing_type == 'percentile':
+            image = (image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            image = np.clip(image, 0, 1)
+        return torch.from_numpy(image)
+
+    def _handle_channels_and_timepoints(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Handle channels and timepoints in the image.
+
+        Args:
+            image: The image to process.
+
+        Returns:
+            Processed image as a torch tensor.
+        """
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            image = image[:4, :, :, :] if self.use_timepoints else image[0]
+        return torch.from_numpy(image)
+
+    def _apply_transforms(self, image: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
path: torchgeo/datasets/substation_seg.py
user: login: rijuld
body: Hi @nilsleh,

Hope you're doing well! I wanted to clarify if it's essential to shift all data augmentations to the datamodule. If so, could you guide me on which specific parts of the dataset should be moved there?

I've already removed the geotransform and color transform and plan to add them to the datamodule in my next pull request. If there are other elements you’d suggest removing, I can address those too. Once these adjustments are made, would it be possible to merge this PR (pending review) without the datamodule updates?

Thank you very much for your help!
created_at: 2024-10-30T17:02:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,196 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any, NamedTuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Args(NamedTuple):
+    """NamedTuple containing configuration for the dataset."""
+    data_dir: str
+    in_channels: int
+    use_timepoints: bool
+    normalizing_type: str
+    normalizing_factor: np.ndarray[Any, np.float64]
+    means: np.ndarray[Any, np.float64]
+    stds: np.ndarray[Any, np.float64]
+    mask_2d: bool
+    model_type: str
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+    
+    It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    """
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = ''
+    url_for_masks = ''
+
+    def __init__(self, args: Args, image_files: list[str], geo_transforms: Any = None, 
+                 color_transforms: Any = None, image_resize: Any = None, mask_resize: Any = None) -> None:
+        """Initialize the dataset with the provided parameters.
+
+        Args:
+            args: Arguments that contain configuration information such as data_dir, in_channels, etc.
+            image_files: List of image filenames for the dataset.
+            geo_transforms: Geometric transformations to apply on the images and masks.
+            color_transforms: Color transformations to apply on the images.
+            image_resize: Resize transformation for images.
+            mask_resize: Resize transformation for masks.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(self.data_dir, 'image_stack')
+        self.mask_dir = os.path.join(self.data_dir, 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+        # Check if the dataset is available or needs to be downloaded
+        self._verify()
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        mask = np.load(mask_path)['arr_0']
+
+        # Standardize images
+        image = self._normalize_image(image)
+
+        # Handle multiple channels and timepoints
+        image = self._handle_channels_and_timepoints(image)
+
+        # Process mask
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+        mask = torch.from_numpy(mask).float().unsqueeze(dim=0)
+
+        # Apply geo and color transformations
+        image, mask = self._apply_transforms(image, mask)
+
+        return {"image": image, "mask": mask}
+
+    def __len__(self) -> int:
+        """Return the length of the dataset.
+
+        Returns:
+            The number of items in the dataset.
+        """
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
+        """Plot a random image and mask from the dataset."""
+        index = np.random.randint(0, self.__len__())
+        data = self.__getitem__(index)
+        image, mask = data["image"], data["mask"]
+        fig, axs = plt.subplots(1, 2, figsize=(15, 15))
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(image.permute(1, 2, 0))
+        axs[1].imshow(mask.permute(1, 2, 0), alpha=0.5, cmap='gray')
+
+    def _normalize_image(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Normalize the image based on the selected normalizing type.
+
+        Args:
+            image: The image to normalize.
+
+        Returns:
+            Normalized image as a torch tensor.
+        """
+        if self.normalizing_type == 'percentile':
+            image = (image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            image = np.clip(image, 0, 1)
+        return torch.from_numpy(image)
+
+    def _handle_channels_and_timepoints(self, image: np.ndarray[Any, np.float64]) -> torch.Tensor:
+        """Handle channels and timepoints in the image.
+
+        Args:
+            image: The image to process.
+
+        Returns:
+            Processed image as a torch tensor.
+        """
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            image = image[:4, :, :, :] if self.use_timepoints else image[0]
+        return torch.from_numpy(image)
+
+    def _apply_transforms(self, image: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
path: torchgeo/datasets/substation_seg.py
user: login: nilsleh
body: Apologies for the late response. Adam prefers having all data normalization in the datamodule for consistency, But I also don't think it is terrible to do in the dataset. If you move it to the datamodule, you can use the kornina `Normalize` module, that you can add to the augmentation series. Then it will be applied to `on_after_batch_transfer` in the `LightningDataModule`. 
created_at: 2024-11-07T16:26:33Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: All files need a copyright, see https://torchgeo.readthedocs.io/en/stable/user/contributing.html#licensing
created_at: 2024-11-10T12:29:32Z
author_association: MEMBER
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: Is there a non-segmentation version of this dataset? If not, let's just name the file `substation.py`.
created_at: 2024-11-10T12:30:02Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: These lines are likely over the 88 character line length limit.

Can you include a URL to link to the homepage and a more detailed description of the dataset? See other dataset files as examples of what things we like to document.
created_at: 2024-11-10T12:31:38Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: Can you use the real URL instead of this url defense wrapper?
created_at: 2024-11-10T12:32:01Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: These downloads also need MD5 checksums
created_at: 2024-11-10T12:32:27Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
+    url_for_masks: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_mask.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=nHMdYvxKmzwAdT2lOPoQ7-NEfjsOjAm00kHOcwC_AmU&e='
+
+    def __init__(
+        self,
+        args: Any,
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: Let's get rid of `args` and instead have specific parameters for each valid input. This helps with type checking and documenting. At the moment, there is absolutely no documentation suggesting that `args.use_time_stamp` is a required attribute of this mysterious `Any` class that is not documented anywhere.
created_at: 2024-11-10T12:33:56Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
+    url_for_masks: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_mask.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=nHMdYvxKmzwAdT2lOPoQ7-NEfjsOjAm00kHOcwC_AmU&e='
+
+    def __init__(
+        self,
+        args: Any,
+        image_files: list[str],
+        geo_transforms: Any | None = None,
+        color_transforms: Any | None = None,
+        image_resize: Any | None = None,
+        mask_resize: Any | None = None,
+    ) -> None:
+        """Initialize the SubstationDataset.
+
+        Args:
+            args (Any): Arguments containing various dataset parameters such as `data_dir`, `in_channels`, etc.
+            image_files (list[str]): A list of image file names.
+            geo_transforms (Any | None): Geometric transformations to be applied to the images and masks. Defaults to None.
+            color_transforms (Any | None): Color transformations to be applied to the images. Defaults to None.
+            image_resize (Any | None): Transformation for resizing the images. Defaults to None.
+            mask_resize (Any | None): Transformation for resizing the masks. Defaults to None.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(args.data_dir, 'substation_seg', 'image_stack')
+        self.mask_dir = os.path.join(args.data_dir, 'substation_seg', 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        # standardizing image
+        if self.normalizing_type == 'percentile':
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: Let's move preprocessing out of the data loader and into a data module
created_at: 2024-11-10T12:34:24Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
+    url_for_masks: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_mask.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=nHMdYvxKmzwAdT2lOPoQ7-NEfjsOjAm00kHOcwC_AmU&e='
+
+    def __init__(
+        self,
+        args: Any,
+        image_files: list[str],
+        geo_transforms: Any | None = None,
+        color_transforms: Any | None = None,
+        image_resize: Any | None = None,
+        mask_resize: Any | None = None,
+    ) -> None:
+        """Initialize the SubstationDataset.
+
+        Args:
+            args (Any): Arguments containing various dataset parameters such as `data_dir`, `in_channels`, etc.
+            image_files (list[str]): A list of image file names.
+            geo_transforms (Any | None): Geometric transformations to be applied to the images and masks. Defaults to None.
+            color_transforms (Any | None): Color transformations to be applied to the images. Defaults to None.
+            image_resize (Any | None): Transformation for resizing the images. Defaults to None.
+            mask_resize (Any | None): Transformation for resizing the masks. Defaults to None.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(args.data_dir, 'substation_seg', 'image_stack')
+        self.mask_dir = os.path.join(args.data_dir, 'substation_seg', 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        # standardizing image
+        if self.normalizing_type == 'percentile':
+            image = (
+                image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))
+            ) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            # means = np.array([1431, 1233, 1209, 1192, 1448, 2238, 2609, 2537, 2828, 884, 20, 2226, 1537]).reshape(-1, 1, 1)
+            # stds = np.array([157, 254, 290, 420, 363, 457, 575, 606, 630, 156, 3, 554, 523]).reshape(-1, 1, 1)
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            # clipping image to 0,1 range
+            image = np.clip(image, 0, 1)
+
+        # selecting channels
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            if self.model_type == 'swin':
+                image = image[
+                    :, [3, 2, 1, 4, 5, 6, 7, 10, 11], :, :
+                ]  # swin only takes 9 channels
+            else:
+                image = image[:, : self.in_channels, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            image = image[:4, :, :, :]
+            if self.args.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (4*channels,h,w)
+            elif self.args.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            # image = np.median(image, axis=0)
+            # image = image[0]
+            if self.args.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.args.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).float()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+
+        if self.image_resize:
+            image = self.image_resize(image)
+
+        if self.mask_resize:
+            mask = self.mask_resize(mask)
+
+        return {'image': image, 'mask': mask}
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: The plot method takes a sample as input and plots it. See every other dataset for an example of this.
created_at: 2024-11-10T12:34:56Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,170 @@
+"""This module handles the Substation segmentation dataset."""
+
+import os
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """SubstationDataset is responsible for handling the loading and transformation of substation segmentation datasets.
+
+    It extends NonGeoDataset, providing methods for dataset verification, downloading, and transformation.
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_image-5Fstack.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=zMCjuqjPMHRz5jeEWLCEufHvWxRPdlHEbPnUE7kXPrc&e='
+    url_for_masks: str = 'https://urldefense.proofpoint.com/v2/url?u=https-3A__storage.googleapis.com_tz-2Dml-2Dpublic_substation-2Dover-2D10km2-2Dcsv-2Dmain-2D444e360fd2b6444b9018d509d0e4f36e_mask.tar.gz&d=DwMFaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=ypwhORbsf5rB8FTl-SAxjfN_U0jrVqx6UDyBtJHbKQY&m=-2QXCp-gZof5HwBsLg7VwQD-pnLedAo09YCzdDCUTqCI-0t789z0-HhhgwVbYtX7&s=nHMdYvxKmzwAdT2lOPoQ7-NEfjsOjAm00kHOcwC_AmU&e='
+
+    def __init__(
+        self,
+        args: Any,
+        image_files: list[str],
+        geo_transforms: Any | None = None,
+        color_transforms: Any | None = None,
+        image_resize: Any | None = None,
+        mask_resize: Any | None = None,
+    ) -> None:
+        """Initialize the SubstationDataset.
+
+        Args:
+            args (Any): Arguments containing various dataset parameters such as `data_dir`, `in_channels`, etc.
+            image_files (list[str]): A list of image file names.
+            geo_transforms (Any | None): Geometric transformations to be applied to the images and masks. Defaults to None.
+            color_transforms (Any | None): Color transformations to be applied to the images. Defaults to None.
+            image_resize (Any | None): Transformation for resizing the images. Defaults to None.
+            mask_resize (Any | None): Transformation for resizing the masks. Defaults to None.
+        """
+        self.data_dir = args.data_dir
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.in_channels = args.in_channels
+        self.use_timepoints = args.use_timepoints
+        self.normalizing_type = args.normalizing_type
+        self.normalizing_factor = args.normalizing_factor
+        self.mask_2d = args.mask_2d
+        self.model_type = args.model_type
+        self.image_dir = os.path.join(args.data_dir, 'substation_seg', 'image_stack')
+        self.mask_dir = os.path.join(args.data_dir, 'substation_seg', 'mask')
+        self.image_filenames = image_files
+        self.args = args
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+        # standardizing image
+        if self.normalizing_type == 'percentile':
+            image = (
+                image - self.normalizing_factor[:, 0].reshape((-1, 1, 1))
+            ) / self.normalizing_factor[:, 2].reshape((-1, 1, 1))
+        elif self.normalizing_type == 'zscore':
+            # means = np.array([1431, 1233, 1209, 1192, 1448, 2238, 2609, 2537, 2828, 884, 20, 2226, 1537]).reshape(-1, 1, 1)
+            # stds = np.array([157, 254, 290, 420, 363, 457, 575, 606, 630, 156, 3, 554, 523]).reshape(-1, 1, 1)
+            image = (image - self.args.means) / self.args.stds
+        else:
+            image = image / self.normalizing_factor
+            # clipping image to 0,1 range
+            image = np.clip(image, 0, 1)
+
+        # selecting channels
+        if self.in_channels == 3:
+            image = image[:, [3, 2, 1], :, :]
+        else:
+            if self.model_type == 'swin':
+                image = image[
+                    :, [3, 2, 1, 4, 5, 6, 7, 10, 11], :, :
+                ]  # swin only takes 9 channels
+            else:
+                image = image[:, : self.in_channels, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            image = image[:4, :, :, :]
+            if self.args.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (4*channels,h,w)
+            elif self.args.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            # image = np.median(image, axis=0)
+            # image = image[0]
+            if self.args.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.args.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).float()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+
+        if self.image_resize:
+            image = self.image_resize(image)
+
+        if self.mask_resize:
+            mask = self.mask_resize(mask)
+
+        return {'image': image, 'mask': mask}
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(self) -> None:
+        """Plots a random image and mask from the dataset."""
+        index = np.random.randint(0, self.__len__())
+        data = self.__getitem__(index)
+        image = data['image']
+        mask = data['mask']
+
+        fig, axs = plt.subplots(1, 2, figsize=(15, 15))
+        axs[0].imshow(image.permute(1, 2, 0).cpu().numpy())
+        axs[1].imshow(image.permute(1, 2, 0).cpu().numpy())
+        axs[1].imshow(mask.squeeze().cpu().numpy(), alpha=0.5, cmap='gray')
+
+    def _verify(self) -> None:
+        """Checks if dataset exists, otherwise download it."""
+        image_dir_exists = os.path.exists(self.image_dir)
+        mask_dir_exists = os.path.exists(self.mask_dir)
+        if not (image_dir_exists and mask_dir_exists):
+            self._download()
path: torchgeo/datasets/substation_seg.py
user: login: adamjstewart
body: Downloading random files from the internet without checksum verification should not happen by default, the user should have to pass `download=True` if they really want things to be downloaded. This violates the principle of least surprise.
created_at: 2024-11-10T12:35:41Z
author_association: MEMBER
diff_hunk: @@ -369,6 +369,11 @@ PASTIS
 
 .. autoclass:: PASTIS
 
+SubstationDataset
path: docs/api/datasets.rst
user: login: adamjstewart
body: These should be in alphabetical order
created_at: 2024-11-10T12:36:06Z
author_association: MEMBER
diff_hunk: @@ -49,6 +49,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
+`SubstationDataset`_,S,OpenStreetMap & Sentinel-2, "CC BY-SA 2.0", 27K, 2, 228x228, 10, MSI
path: docs/api/datasets/non_geo_datasets.csv
user: login: adamjstewart
body: ```suggestion
`SubstationDataset`_,S,OpenStreetMap & Sentinel-2, "CC-BY-SA 2.0", 27K, 2, 228x228, 10, MSI
```
Should be a valid SPDX identifier
created_at: 2024-11-10T12:36:37Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,267 @@
+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 numpy as np
+import pytest
+import torch
+import torchvision.transforms as transforms
+
+from torchgeo.datasets import SubstationDataset
+
+
+class Args:
+    """Mocked arguments for testing SubstationDataset."""
+
+    def __init__(self) -> None:
+        self.data_dir: str = os.path.join(os.getcwd(), 'tests', 'data')
+        self.in_channels: int = 13
+        self.use_timepoints: bool = True
+        self.normalizing_type: str = 'percentile'
+        self.mask_2d: bool = True
+        self.model_type: str = 'vanilla_unet'
+        self.timepoint_aggregation: str = 'median'
+        self.color_transforms: bool = False
+        self.geo_transforms: bool = False
+        self.normalizing_factor: Any = np.array([[0, 0.5, 1.0]], dtype=np.float32)
+        self.means: Any = np.array(
+            [
+                [[1431]],
+                [[1233]],
+                [[1209]],
+                [[1192]],
+                [[1448]],
+                [[2238]],
+                [[2609]],
+                [[2537]],
+                [[2828]],
+                [[884]],
+                [[20]],
+                [[2226]],
+                [[1537]],
+            ],
+            dtype=np.float32,
+        )
+        self.stds: Any = np.array(
+            [
+                [[157]],
+                [[254]],
+                [[290]],
+                [[420]],
+                [[363]],
+                [[457]],
+                [[575]],
+                [[606]],
+                [[630]],
+                [[156]],
+                [[3]],
+                [[554]],
+                [[523]],
+            ],
+            dtype=np.float32,
+        )
+
+
+@pytest.fixture
+def dataset(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> Generator[SubstationDataset, None, None]:
+    """Fixture for the SubstationDataset."""
+    args = Args()
+    image_files = ['image_0.npz', 'image_1.npz']
+
+    yield SubstationDataset(args, image_files)
+
+
+@pytest.mark.parametrize(
+    'config',
+    [
+        {
+            'normalizing_type': 'percentile',
+            'in_channels': 3,
+            'use_timepoints': False,
+            'mask_2d': True,
+        },
+        {
+            'normalizing_type': 'zscore',
+            'in_channels': 9,
+            'model_type': 'swin',
+            'use_timepoints': True,
+            'timepoint_aggregation': 'concat',
+            'mask_2d': False,
+        },
+        {
+            'normalizing_type': None,
+            'in_channels': 12,
+            'use_timepoints': True,
+            'timepoint_aggregation': 'median',
+            'mask_2d': True,
+            'normalizing_factor': 1.0,
+        },
+        {
+            'normalizing_type': None,
+            'in_channels': 5,
+            'use_timepoints': True,
+            'timepoint_aggregation': 'first',
+            'mask_2d': False,
+            'normalizing_factor': 1.0,
+        },
+        {
+            'normalizing_type': None,
+            'in_channels': 4,
+            'use_timepoints': True,
+            'timepoint_aggregation': 'random',
+            'mask_2d': True,
+            'normalizing_factor': 1.0,
+        },
+        {
+            'normalizing_type': 'zscore',
+            'in_channels': 2,
+            'use_timepoints': False,
+            'mask_2d': False,
+            'color_transforms': True,
+            'geo_transforms': True,
+        },
+        {
+            'normalizing_type': None,
+            'in_channels': 5,
+            'use_timepoints': False,
+            'timepoint_aggregation': 'first',
+            'mask_2d': False,
+            'normalizing_factor': 1.0,
+        },
+        {
+            'normalizing_type': None,
+            'in_channels': 4,
+            'use_timepoints': False,
+            'timepoint_aggregation': 'random',
+            'mask_2d': True,
+            'normalizing_factor': 1.0,
+        },
+    ],
+)
+def test_getitem_semantic(config: dict[str, Any]) -> None:
+    args = Args()
+    for key, value in config.items():
+        setattr(args, key, value)  # Dynamically set arguments for each config
+
+    # Setting mock paths and creating dataset instance
+    image_files = ['image_0.npz', 'image_1.npz']
+    image_resize = transforms.Compose(
+        [transforms.Resize(228, transforms.InterpolationMode.BICUBIC)]
+    )
+    mask_resize = transforms.Compose(
+        [transforms.Resize(228, transforms.InterpolationMode.NEAREST)]
+    )
+    dataset = SubstationDataset(
+        args, image_files, image_resize=image_resize, mask_resize=mask_resize
+    )
+
+    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(dataset: SubstationDataset) -> None:
+    """Test the length of the dataset."""
+    assert len(dataset) == 2
+
+
+def test_output_shape(dataset: SubstationDataset) -> None:
+    """Test the output shape of the dataset."""
+    x = dataset[0]
+    assert x['image'].shape == torch.Size([13, 228, 228])
+    assert x['mask'].shape == torch.Size([2, 228, 228])
+
+
+def test_plot(dataset: SubstationDataset, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test the plot method of the dataset."""
+    # Mock plt.show to avoid showing the plot during the test
+    mock_show = MagicMock()
+    monkeypatch.setattr(plt, 'show', mock_show)
+
+    # Mock np.random.randint to return a fixed index (e.g., 0)
+    monkeypatch.setattr(
+        np.random, 'randint', lambda low, high: 0
+    )  # Correct the lambda to accept 2 arguments
+
+    # Mock __getitem__ to return a sample with an image (3 channels) and a mask
+    mock_image = torch.rand(3, 228, 228)  # Create a dummy 3-channel image (RGB)
+    mock_mask = torch.randint(0, 4, (228, 228))  # Create a dummy mask
+    monkeypatch.setattr(
+        dataset, '__getitem__', lambda idx: {'image': mock_image, 'mask': mock_mask}
+    )
+
+    # Call the plot method
+    dataset.plot()
+
+
+def test_already_downloaded(
+    dataset: SubstationDataset, 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_seg', 'image_stack.tar.gz'
+    )
+    url_for_masks = os.path.join('tests', 'data', 'substation_seg', '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_verify(dataset: SubstationDataset, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test the _verify method of the dataset."""
+    # Mock os.path.exists to return False for the image and mask directories
+    monkeypatch.setattr(os.path, 'exists', lambda path: False)
+
+    # Mock the _download method to avoid actually downloading the dataset
+    mock_download = MagicMock()
+    monkeypatch.setattr(dataset, '_download', mock_download)
+
+    # Call the _verify method
+    dataset._verify()
+
+    # Check that the _download method was called
+    mock_download.assert_called_once()
+
+
+def test_download(
+    dataset: SubstationDataset, 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_seg.download_url', mock_download_url
+    )
+    monkeypatch.setattr(
+        'torchgeo.datasets.substation_seg.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__])
path: tests/datasets/test_substation_seg.py
user: login: adamjstewart
body: This is not needed, the file is run by pytest, not the other way around.
created_at: 2024-11-10T12:38:07Z
author_association: MEMBER
diff_hunk: @@ -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 = 228  # Image dimensions
path: tests/data/substation/data.py
user: login: adamjstewart
body: We usually use much smaller fake images (32 x 32) to make the tests run faster
created_at: 2024-12-12T00:48:41Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,157 @@
+import os
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Needs the same copyright header
created_at: 2024-12-12T00:49:03Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,157 @@
+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 SubstationDataset
+
+
+class Args:
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Let's get rid of this
created_at: 2024-12-12T00:49:17Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,157 @@
+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 SubstationDataset
+
+
+class Args:
+    """Mocked arguments for testing SubstationDataset."""
+
+    def __init__(self) -> None:
+        self.data_dir: str = os.path.join(os.getcwd(), 'tests', 'data')
+        self.in_channels: int = 13
+        self.use_timepoints: bool = True
+        self.mask_2d: bool = True
+        self.timepoint_aggregation: str = 'median'
+
+
+@pytest.fixture
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Can we put all of these in a class like our other tests?
created_at: 2024-12-12T00:49:36Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: It's not a base class, it's a class
created_at: 2024-12-12T00:50:50Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
+
+    This dataset is responsible for handling the loading and transformation of
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Is there a homepage for this dataset where I can find more information?
created_at: 2024-12-12T00:51:03Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: We don't usually put "Dataset" in the class name
created_at: 2024-12-12T00:51:13Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
+
+    This dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory: str = 'Substation'
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: You don't need to specify the types here, mypy can automatically infer them
created_at: 2024-12-12T00:51:45Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
+
+    This dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        data_dir: str,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: We usually call this variable `root`
created_at: 2024-12-12T00:51:57Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
+
+    This dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        data_dir: str,
+        in_channels: int,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: We usually call this `bands` and allow the user to specify a list of bands, which may not be in order
created_at: 2024-12-12T00:52:30Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,241 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class SubstationDataset(NonGeoDataset):
+    """Base class for Substation Dataset.
+
+    This dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory: str = 'Substation'
+    filename_images: str = 'image_stack.tar.gz'
+    filename_masks: str = 'mask.tar.gz'
+    url_for_images: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks: str = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        data_dir: str,
+        in_channels: int,
+        use_timepoints: bool,
+        image_files: list[str],
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: This is usually automatically detected by the dataset class based on the downloaded data
created_at: 2024-12-12T00:53:11Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: str,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: This should be `.utils.Path`, you can copy this from nearby files. This will solve the mypy issue.
created_at: 2025-01-19T12:23:13Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: str,
path: torchgeo/datasets/substation.py
user: login: rijuld
body: Thanks @adamjstewart. My datamodule uses geotransforms and color transforms as arguments, how can we add those in substation yaml file to test?
created_at: 2025-01-19T14:45:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: str,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: It should be possible to include complex arguments using class_path, init_args, dict_kwargs, etc. Or you could simplify your datamodule.
created_at: 2025-01-19T15:41:52Z
author_association: MEMBER
diff_hunk: @@ -53,6 +53,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
path: docs/api/datasets/non_geo_datasets.csv
user: login: adamjstewart
body: Where did you find the license? I don't see anything on GitHub or in the paper.
created_at: 2025-01-26T14:47:27Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,224 @@
+# 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 DatasetNotFoundError, 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', 'substation')
+
+        yield Substation(
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: We don't normally use yield when creating our fixtures, any reason you're using a generator?
created_at: 2025-01-26T14:49:25Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,224 @@
+# 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 DatasetNotFoundError, 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', 'substation')
+
+        yield Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+        )
+
+    @pytest.mark.parametrize(
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Shouldn't you parametrize the fixture instead of the unit test? Then all other unit tests will also be parametrized.
created_at: 2025-01-26T14:50:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,224 @@
+# 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
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: We don't use `unittest`, we use `pytest`. Can you convert these to `MonkeyPatch` which is builtin to pytest?
created_at: 2025-01-26T14:51:17Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,162 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import torch
+from torch.utils.data import Subset, random_split
+from tqdm import tqdm
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: tqdm isn't a dependency and we don't really want to add new deps unless we have to
created_at: 2025-01-26T14:53:39Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,162 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import torch
+from torch.utils.data import Subset, random_split
+from tqdm import tqdm
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        normalizing_type: str = 'percentile',
+        normalizing_factor: Any | None = None,
+        means: Any | None = None,
+        stds: Any | None = None,
+        bands: list[int] = [1, 2, 3],
+        num_of_timepoints: int = 4,
+        model_type: str = 'default',
+        geo_transforms: Any | None = None,
+        color_transforms: Any | None = None,
+        image_resize: Any | None = None,
+        mask_resize: Any | None = None,
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: We try to avoid dynamic typing, can you be more specific than `Any` which disables all type checking?
created_at: 2025-01-26T14:54:16Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,162 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import torch
+from torch.utils.data import Subset, random_split
+from tqdm import tqdm
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        normalizing_type: str = 'percentile',
+        normalizing_factor: Any | None = None,
+        means: Any | None = None,
+        stds: Any | None = None,
+        bands: list[int] = [1, 2, 3],
+        num_of_timepoints: int = 4,
+        model_type: str = 'default',
+        geo_transforms: Any | None = None,
+        color_transforms: Any | None = None,
+        image_resize: Any | None = None,
+        mask_resize: Any | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            root: Path to the dataset directory.
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            normalizing_type: Normalization type ('percentile', 'zscore', or 'default').
+            normalizing_factor: Normalization factor for percentile normalization.
+            means: Mean values for z-score normalization.
+            stds: Standard deviation values for z-score normalization.
+            bands: Number of input channels to use.
+            model_type: Type of model being used (e.g., 'swin' for specific channel selection).
+            geo_transforms: Geometric transformations to apply to the data.
+            color_transforms: Color transformations to apply to the image.
+            image_resize: Resizing function for the image.
+            mask_resize: Resizing function for the mask.
+            num_of_timepoints: Number of timepoints to use in the dataset.
+            **kwargs: Additional arguments passed to Substation.
+        """
+        super().__init__(Substation, batch_size, num_workers, **kwargs)
+        self.root = root
+        self.val_split_pct = val_split_pct
+        self.test_split_pct = test_split_pct
+        self.normalizing_type = normalizing_type
+        self.normalizing_factor = normalizing_factor
+        self.means = means
+        self.stds = stds
+        self.bands = bands
+        self.model_type = model_type
+        self.geo_transforms = geo_transforms
+        self.color_transforms = color_transforms
+        self.image_resize = image_resize
+        self.mask_resize = mask_resize
+        self.num_of_timepoints = num_of_timepoints
+        self.geo_transforms = (
+            geo_transforms if geo_transforms is not None else self._identity
+        )
+        self.color_transforms = (
+            color_transforms if color_transforms is not None else self._identity
+        )
+        self.image_resize = image_resize if image_resize is not None else self._identity
+        self.mask_resize = mask_resize if mask_resize is not None else self._identity
+
+    def _identity(self, x: torch.Tensor) -> torch.Tensor:
+        """Identity function for default transformations."""
+        return x
+
+    def setup(self, stage: str) -> None:
+        """Set up datasets.
+
+        Args:
+            stage: One of 'fit', 'validate', 'test', or 'predict'.
+        """
+        dataset = Substation(
+            root=self.root,
+            bands=self.bands,
+            use_timepoints=True,
+            mask_2d=False,
+            num_of_timepoints=self.num_of_timepoints,
+            timepoint_aggregation='median',
+            download=True,
+            checksum=False,
+        )
+
+        generator = torch.Generator().manual_seed(0)
+        total_len = len(dataset)
+        val_len = int(total_len * self.val_split_pct)
+        test_len = int(total_len * self.test_split_pct)
+        train_len = total_len - val_len - test_len
+
+        self.train_dataset, self.val_dataset, self.test_dataset = random_split(
+            dataset, [train_len, val_len, test_len], generator
+        )
+
+        if stage == 'fit':
+            self.train_dataset = self._apply_transforms(self.train_dataset)
+            self.val_dataset = self._apply_transforms(self.val_dataset)
+        elif stage == 'test':
+            self.test_dataset = self._apply_transforms(self.test_dataset)
+
+    def _apply_transforms(self, dataset: Subset[Any]) -> Subset[Any]:
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: You should use `self.aug`, `self.train_aug`, etc. to store the augmentations you want to perform. These are then automatically run on the GPU for you, no need to run them manually. See https://torchgeo.readthedocs.io/en/latest/tutorials/contribute_datamodule.html, `torchgeo/datamodules/geo.py`, and every other builtin datamodule for examples of this.
created_at: 2025-01-26T14:56:17Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
```
created_at: 2025-01-26T14:56:54Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: No need to document base class, this is automatically added to the docs. Instead, describe the data itself.
created_at: 2025-01-26T14:57:28Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Needs an empty line before and after, otherwise it's all a single paragraph. See https://torchgeo--2352.org.readthedocs.build/en/2352/api/datasets.html#substation
created_at: 2025-01-26T14:58:32Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Needs empty line after, see https://torchgeo--2352.org.readthedocs.build/en/2352/api/datasets.html#substation
created_at: 2025-01-26T14:58:54Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: list[int],
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Are tuples also allowed? In this case you would use Sequence instead of list
created_at: 2025-01-26T14:59:39Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: list[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            num_of_timepoints: Number of timepoints to use for each image.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+        """
+        self.root = root
+        self.bands = bands
+        self.use_timepoints = use_timepoints
+        self.timepoint_aggregation = timepoint_aggregation
+        self.mask_2d = mask_2d
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self.download = download
+        self.checksum = checksum
+
+        self._verify()
+        self.load_image_filenames()
+
+    def load_image_filenames(self) -> None:
+        """Load image filenames from the image directory."""
+        self.image_filenames = os.listdir(self.image_dir)
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
        self.image_filenames = sorted(os.listdir(self.image_dir))
```
No need for a helper function for a single line of code. Also, sorting this list makes the dataset reproducible.
created_at: 2025-01-26T15:00:46Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: list[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            num_of_timepoints: Number of timepoints to use for each image.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: These are in a different order than the parameters. Also, can you add a transforms parameter to match all of our other datasets? It should be used at the end of `__getitem__`.
created_at: 2025-01-26T15:01:34Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: list[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            num_of_timepoints: Number of timepoints to use for each image.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+        """
+        self.root = root
+        self.bands = bands
+        self.use_timepoints = use_timepoints
+        self.timepoint_aggregation = timepoint_aggregation
+        self.mask_2d = mask_2d
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self.download = download
+        self.checksum = checksum
+
+        self._verify()
+        self.load_image_filenames()
+
+    def load_image_filenames(self) -> None:
+        """Load image filenames from the image directory."""
+        self.image_filenames = os.listdir(self.image_dir)
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            image = image[: self.num_of_timepoints, :, :, :]
+            if self.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (num_of_timepoints*channels,h,w)
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            # image = np.median(image, axis=0)
+            # image = image[0]
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).long()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+        mask = mask.squeeze()
+
+        return {'image': image, 'mask': mask}
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(
+        self,
+        sample: dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: str | None = None,
+    ) -> Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional string to use as a suptitle
+
+        Returns:
+            A matplotlib Figure containing the rendered sample.
+        """
+        ncols = 2
+
+        image = sample['image'][:3].permute(1, 2, 0).cpu().numpy()
+        image = image / 255.0  # Normalize image
+
+        if self.mask_2d:
+            mask = sample['mask'][0].squeeze(0).cpu().numpy()
+        else:
+            mask = sample['mask'].cpu().numpy()
+        showing_predictions = 'prediction' in sample
+        if showing_predictions:
+            prediction = sample['prediction'].cpu().numpy()
+            if self.mask_2d:
+                prediction = prediction[0]
+            print(prediction.shape)
+            ncols = 3
+
+        print(mask.shape, image.shape)
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Can remove debugging print statements
created_at: 2025-01-26T15:02:20Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: list[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            num_of_timepoints: Number of timepoints to use for each image.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+        """
+        self.root = root
+        self.bands = bands
+        self.use_timepoints = use_timepoints
+        self.timepoint_aggregation = timepoint_aggregation
+        self.mask_2d = mask_2d
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self.download = download
+        self.checksum = checksum
+
+        self._verify()
+        self.load_image_filenames()
+
+    def load_image_filenames(self) -> None:
+        """Load image filenames from the image directory."""
+        self.image_filenames = os.listdir(self.image_dir)
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            image = image[: self.num_of_timepoints, :, :, :]
+            if self.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (num_of_timepoints*channels,h,w)
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            # image = np.median(image, axis=0)
+            # image = image[0]
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).long()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+        mask = mask.squeeze()
+
+        return {'image': image, 'mask': mask}
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(
+        self,
+        sample: dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: str | None = None,
+    ) -> Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional string to use as a suptitle
+
+        Returns:
+            A matplotlib Figure containing the rendered sample.
+        """
+        ncols = 2
+
+        image = sample['image'][:3].permute(1, 2, 0).cpu().numpy()
+        image = image / 255.0  # Normalize image
+
+        if self.mask_2d:
+            mask = sample['mask'][0].squeeze(0).cpu().numpy()
+        else:
+            mask = sample['mask'].cpu().numpy()
+        showing_predictions = 'prediction' in sample
+        if showing_predictions:
+            prediction = sample['prediction'].cpu().numpy()
+            if self.mask_2d:
+                prediction = prediction[0]
+            print(prediction.shape)
+            ncols = 3
+
+        print(mask.shape, image.shape)
+
+        fig, axs = plt.subplots(ncols=ncols, figsize=(4 * ncols, 4))
+        axs[0].imshow(image)
+        axs[0].axis('off')
+        axs[1].imshow(mask, cmap='gray', interpolation='none')
+        axs[1].axis('off')
+
+        if show_titles:
+            axs[0].set_title('Image')
+            axs[1].set_title('Mask')
+
+        if showing_predictions:
+            axs[2].imshow(prediction, cmap='gray', interpolation='none')
+            axs[2].axis('off')
+            if show_titles:
+                axs[2].set_title('Prediction')
+
+        if suptitle:
+            fig.suptitle(suptitle)
+
+        return fig
+
+    def _extract(self) -> None:
+        """Extract the dataset."""
+        img_pathname = os.path.join(self.root, self.filename_images)
+        extract_archive(img_pathname)
+
+        mask_pathname = os.path.join(self.root, self.filename_masks)
+        extract_archive(mask_pathname)
+
+    def _verify(self) -> None:
+        """Verify the integrity of the dataset."""
+        # Check if the extracted files already exist
+        image_path = os.path.join(self.image_dir, '*.npz')
+        mask_path = os.path.join(self.mask_dir, '*.npz')
+        if glob.glob(image_path) and glob.glob(mask_path):
+            return
+
+        # Check if the tar.gz files for images and masks have already been downloaded
+        image_exists = os.path.exists(os.path.join(self.root, self.filename_images))
+        mask_exists = os.path.exists(os.path.join(self.root, self.filename_masks))
+        if image_exists and mask_exists:
+            self._extract()
+            return
+
+        # If dataset files are missing and download is not allowed, raise an error
+        if not getattr(self, 'download', True):
+            raise DatasetNotFoundError(self)
+
+        # Download and extract the dataset
+        self._download()
+        self._extract()
+
+    def _download(self) -> None:
+        """Download the dataset and extract it."""
+        # Download and verify images
+        download_url(
+            self.url_for_images,
+            self.root,
+            filename=self.filename_images,
+            md5='INSERT_IMAGES_MD5_HASH' if self.checksum else None,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: TODO: add MD5s
created_at: 2025-01-26T15:03:09Z
author_association: MEMBER
diff_hunk: @@ -53,6 +53,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
path: docs/api/datasets/non_geo_datasets.csv
user: login: rijuld
body: I will change the links to hugging face links; we have updated the license there.
created_at: 2025-02-05T21:42:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -53,6 +53,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
path: docs/api/datasets/non_geo_datasets.csv
user: login: rijuld
body: https://huggingface.co/datasets/neurograce/SubstationDataset/tree/main
created_at: 2025-02-06T13:22:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,224 @@
+# 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 DatasetNotFoundError, 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', 'substation')
+
+        yield Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+        )
+
+    @pytest.mark.parametrize(
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Still waiting on this comment
created_at: 2025-02-24T09:08:08Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,214 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture
+    def dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    @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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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', 'substation')
+        dataset = Substation(root=root, **config)
+
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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."""
+        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')
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        monkeypatch.setattr(dataset, '_download', lambda: None)
+        dataset._download()
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', lambda *args, **kwargs: None
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', lambda *args, **kwargs: None
+        )
+
+        dataset._download()
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
+            Substation(
+                bands=[1, 2, 3],
+                use_timepoints=True,
+                mask_2d=True,
+                timepoint_aggregation='median',
+                num_of_timepoints=4,
+                root=tmp_path,
+            )
+
+    def test_not_downloaded_with_download(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        filename = 'image_stack'
+        maskname = 'mask'
+        source_image_path = os.path.join('tests', 'data', 'substation', filename)
+        source_mask_path = os.path.join('tests', 'data', 'substation', maskname)
+        target_image_path = tmp_path / filename
+        target_mask_path = tmp_path / maskname
+
+        def mock_download(self: Substation) -> None:
+            shutil.copytree(source_image_path, target_image_path)
+            shutil.copytree(source_mask_path, target_mask_path)
+
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.Substation._download', mock_download
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.Substation._extract', lambda self: None
+        )
+
+        Substation(
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            root=tmp_path,
+            download=True,
+        )
+
+    def test_extract(self, tmp_path: Path) -> None:
+        filename = Substation.filename_images
+        maskname = Substation.filename_masks
+        shutil.copyfile(
+            os.path.join('tests', 'data', 'substation', filename), tmp_path / filename
+        )
+        shutil.copyfile(
+            os.path.join('tests', 'data', 'substation', maskname), tmp_path / maskname
+        )
+        Substation(
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            root=tmp_path,
+        )
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: You can delete these lines. This test can be run with `pytest tests/datasets/test_substation.py`
created_at: 2025-02-24T09:09:46Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        bands: list[int] = [1, 2, 3],
+        num_of_timepoints: int = 4,
+        timepoint_aggregation: str = 'median',
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: Any of these parameters that just get passed straight to the dataset can be removed, they are already covered in `kwargs`.
created_at: 2025-02-24T09:11:27Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        bands: list[int] = [1, 2, 3],
+        num_of_timepoints: int = 4,
+        timepoint_aggregation: str = 'median',
+        model_type: str = 'default',
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: `model_type` doesn't seem to be used
created_at: 2025-02-24T09:11:38Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
"""Substation segmentation dataset."""
```
created_at: 2025-02-24T09:12:05Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
    """Substation dataset.
```
created_at: 2025-02-24T09:12:27Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSMF) and
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
    available data repositories, including OpenStreetMap (OSM) and
```
created_at: 2025-02-24T09:13:15Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSMF) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits)
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
    (i.e., revisits).
```
created_at: 2025-02-24T09:13:43Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSMF) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits)
+
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: Sequence[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: ```suggestion
            bands: Channels to use from the image.
```
created_at: 2025-02-24T09:15:16Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSMF) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits)
+
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: Sequence[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = sorted(os.listdir(self.image_dir))
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: We've seen reports that storing lists as dataset attributes results in memory leaks in Python multiprocessing. It may be better to store these as pandas Series.
created_at: 2025-02-24T09:16:41Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSMF) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits)
+
+    Dataset Format:
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: Sequence[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Number of channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = sorted(os.listdir(self.image_dir))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            if image.shape[0] < self.num_of_timepoints:
+                # Padding: cycle through existing timepoints
+                padded_images = []
+                for i in range(self.num_of_timepoints):
+                    padded_images.append(image[i % image.shape[0]])
+                image = np.stack(padded_images)
+            elif image.shape[0] > self.num_of_timepoints:
+                # Removal: take the most recent timepoints
+                image = image[-self.num_of_timepoints :]
+
+            if self.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (num_of_timepoints*channels,h,w)
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            # image = np.median(image, axis=0)
+            # image = image[0]
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Can remove commented out code
created_at: 2025-02-24T09:17:51Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,247 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""This module handles the Substation segmentation dataset."""
+
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Class for Substation Dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg/blob/main/README.md>`__
+    dataset is responsible for handling the loading and transformation of
+    substation segmentation datasets. It extends NonGeoDataset, providing methods
+    for dataset verification, downloading, and transformation.
+    Dataset Format:
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Needs empty line after
created_at: 2025-02-24T09:19:52Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,224 @@
+# 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 DatasetNotFoundError, 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', 'substation')
+
+        yield Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+        )
+
+    @pytest.mark.parametrize(
path: tests/datasets/test_substation.py
user: login: rijuld
body: I have made a change, could you check if that correct?
created_at: 2025-02-27T14:49:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: None of these are actually being used at the moment. We just rerun the same code 10x without changes. See https://docs.pytest.org/en/stable/how-to/parametrize.html for how to use this properly. Also, if it makes the code simpler, you could use `itertools.product` to create the cross product of multiple parameter options.
created_at: 2025-02-28T09:25:55Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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 dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    def test_getitem_semantic(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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:
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Could probably combine this and `test_getitem_semantic`
created_at: 2025-02-28T09:26:48Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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 dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    def test_getitem_semantic(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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."""
+        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')
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        monkeypatch.setattr(dataset, '_download', lambda: None)
+        dataset._download()
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', lambda *args, **kwargs: None
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', lambda *args, **kwargs: None
+        )
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: This isn't needed, we already monkeypatch these in `conftest.py`
created_at: 2025-02-28T09:27:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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 dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    def test_getitem_semantic(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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."""
+        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')
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        monkeypatch.setattr(dataset, '_download', lambda: None)
+        dataset._download()
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', lambda *args, **kwargs: None
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', lambda *args, **kwargs: None
+        )
+
+        dataset._download()
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
+            Substation(
+                bands=[1, 2, 3],
+                use_timepoints=True,
+                mask_2d=True,
+                timepoint_aggregation='median',
+                num_of_timepoints=4,
+                root=tmp_path,
+            )
+
+    def test_not_downloaded_with_download(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        filename = 'image_stack'
+        maskname = 'mask'
+        source_image_path = os.path.join('tests', 'data', 'substation', filename)
+        source_mask_path = os.path.join('tests', 'data', 'substation', maskname)
+        target_image_path = tmp_path / filename
+        target_mask_path = tmp_path / maskname
+
+        def mock_download(self: Substation) -> None:
+            shutil.copytree(source_image_path, target_image_path)
+            shutil.copytree(source_mask_path, target_mask_path)
+
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.Substation._download', mock_download
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.Substation._extract', lambda self: None
+        )
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Same here, see other dataset test code for how we usually test this.
created_at: 2025-02-28T09:28:10Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            root: Path to the dataset directory.
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            model_type: Type of model being used (e.g., 'swin' for specific channel selection).
+            size: Size of the input images.
+
+            **kwargs: Additional arguments passed to Substation.
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: ```suggestion
            **kwargs: Additional keyword arguments passed to 
                :class:`~torchgeo.dataset.Substation`.
```
created_at: 2025-02-28T09:29:42Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            root: Path to the dataset directory.
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            model_type: Type of model being used (e.g., 'swin' for specific channel selection).
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: model_type is not used and can be removed
created_at: 2025-02-28T09:30:12Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            root: Path to the dataset directory.
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: root can be removed, it's already covered by `kwargs`
created_at: 2025-02-28T09:32:41Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        root: str,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            root: Path to the dataset directory.
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            model_type: Type of model being used (e.g., 'swin' for specific channel selection).
+            size: Size of the input images.
+
+            **kwargs: Additional arguments passed to Substation.
+        """
+        super().__init__(Substation, batch_size, num_workers, **kwargs)
+        self.root = root
+        self.val_split_pct = val_split_pct
+        self.test_split_pct = test_split_pct
+        self.train_aug = K.AugmentationSequential(
+            K.Resize(size),
+            K.RandomHorizontalFlip(p=0.5),
+            K.RandomVerticalFlip(p=0.5),
+            data_keys=None,
+            keepdim=True,
+        )
+        self.aug = K.AugmentationSequential(
+            K.Resize(size), data_keys=None, keepdim=True
+        )
+
+    def setup(self, stage: str) -> None:
+        """Set up datasets.
+
+        Args:
+            stage: One of 'fit', 'validate', 'test', or 'predict'.
+        """
+        dataset = Substation(
+            root=self.root,
+            use_timepoints=True,
+            mask_2d=False,
+            download=True,
+            checksum=False,
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: All of these can be removed. We never want to default to downloading without user consent, and we never want to download in `setup`, only in `prepare_data`. All other args should be controllable by `kwargs` for maximum flexibility.
created_at: 2025-02-28T09:33:58Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,263 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Substation dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSM) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits).
+
+    Dataset Format:
+
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+
+    def __init__(
+        self,
+        root: Path,
+        bands: Sequence[int],
+        mask_2d: bool,
+        timepoint_aggregation: str = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = pd.Series(sorted(os.listdir(self.image_dir)))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            if image.shape[0] < self.num_of_timepoints:
+                # Padding: cycle through existing timepoints
+                padded_images = []
+                for i in range(self.num_of_timepoints):
+                    padded_images.append(image[i % image.shape[0]])
+                image = np.stack(padded_images)
+            elif image.shape[0] > self.num_of_timepoints:
+                # Removal: take the most recent timepoints
+                image = image[-self.num_of_timepoints :]
+
+            if self.timepoint_aggregation == 'concat':
+                image = np.reshape(
+                    image, (-1, image.shape[2], image.shape[3])
+                )  # (num_of_timepoints*channels,h,w)
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
+
+        mask = np.load(mask_path)['arr_0']
+        mask[mask != 3] = 0
+        mask[mask == 3] = 1
+
+        image = torch.from_numpy(image)
+        mask = torch.from_numpy(mask).long()
+        mask = mask.unsqueeze(dim=0)
+
+        if self.mask_2d:
+            mask_0 = 1.0 - mask
+            mask = torch.concat([mask_0, mask], dim=0)
+        mask = mask.squeeze()
+
+        sample = {'image': image, 'mask': mask}
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample
+
+    def __len__(self) -> int:
+        """Returns the number of items in the dataset."""
+        return len(self.image_filenames)
+
+    def plot(
+        self,
+        sample: dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: str | None = None,
+    ) -> Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional string to use as a suptitle
+
+        Returns:
+            A matplotlib Figure containing the rendered sample.
+        """
+        ncols = 2
+
+        image = sample['image'][:3].permute(1, 2, 0).cpu().numpy()
+        image = image / 255.0  # Normalize image
+
+        if self.mask_2d:
+            mask = sample['mask'][0].squeeze(0).cpu().numpy()
+        else:
+            mask = sample['mask'].cpu().numpy()
+        showing_predictions = 'prediction' in sample
+        if showing_predictions:
+            prediction = sample['prediction'].cpu().numpy()
+            if self.mask_2d:
+                prediction = prediction[0]
+            ncols = 3
+
+        fig, axs = plt.subplots(ncols=ncols, figsize=(4 * ncols, 4))
+        axs[0].imshow(image)
+        axs[0].axis('off')
+        axs[1].imshow(mask, cmap='gray', interpolation='none')
+        axs[1].axis('off')
+
+        if show_titles:
+            axs[0].set_title('Image')
+            axs[1].set_title('Mask')
+
+        if showing_predictions:
+            axs[2].imshow(prediction, cmap='gray', interpolation='none')
+            axs[2].axis('off')
+            if show_titles:
+                axs[2].set_title('Prediction')
+
+        if suptitle:
+            fig.suptitle(suptitle)
+
+        return fig
+
+    def _extract(self) -> None:
+        """Extract the dataset."""
+        img_pathname = os.path.join(self.root, self.filename_images)
+        extract_archive(img_pathname)
+
+        mask_pathname = os.path.join(self.root, self.filename_masks)
+        extract_archive(mask_pathname)
+
+    def _verify(self) -> None:
+        """Verify the integrity of the dataset."""
+        # Check if the extracted files already exist
+        image_path = os.path.join(self.image_dir, '*.npz')
+        mask_path = os.path.join(self.mask_dir, '*.npz')
+        if glob.glob(image_path) and glob.glob(mask_path):
+            return
+
+        # Check if the tar.gz files for images and masks have already been downloaded
+        image_exists = os.path.exists(os.path.join(self.root, self.filename_images))
+        mask_exists = os.path.exists(os.path.join(self.root, self.filename_masks))
+        if image_exists and mask_exists:
+            self._extract()
+            return
+
+        # If dataset files are missing and download is not allowed, raise an error
+        if not getattr(self, 'download', True):
+            raise DatasetNotFoundError(self)
+
+        # Download and extract the dataset
+        self._download()
+        self._extract()
+
+    def _download(self) -> None:
+        """Download the dataset and extract it."""
+        # Download and verify images
+        download_url(
+            self.url_for_images,
+            self.root,
+            filename=self.filename_images,
+            md5='bc0e3d2565f30b1c84a0d4cf37d44be6' if self.checksum else None,
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Can we store these md5 at the class level just like url?
created_at: 2025-02-28T09:36:20Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,23 @@
+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'
+    bands: [0, 1, 2]
+    num_of_timepoints: 4
+    timepoint_aggregation: 'median'
+    mask_2d: False
+    download: True,
path: tests/conf/substation.yaml
user: login: adamjstewart
body: We should not download anything during testing, you can remove this line
created_at: 2025-03-06T10:20:24Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+import torch
+from torch.utils.data import random_split
+
+from ..datasets import Substation
+from .geo import NonGeoDataModule
+
+
+class SubstationDataModule(NonGeoDataModule):
+    """Substation Data Module with train-test split and transformations.
+
+    .. versionadded:: 0.7
+    """
+
+    def __init__(
+        self,
+        batch_size: int = 64,
+        num_workers: int = 0,
+        val_split_pct: float = 0.2,
+        test_split_pct: float = 0.2,
+        size: int = 256,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new SubstationDataModule instance.
+
+        Args:
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for data loading.
+            val_split_pct: Percentage of data to use for validation.
+            test_split_pct: Percentage of data to use for testing.
+            size: Size of the input images.
+
path: torchgeo/datamodules/substation.py
user: login: adamjstewart
body: Remove this empty line
created_at: 2025-03-06T10:21:55Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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 dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    def test_getitem_semantic(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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."""
+        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')
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        monkeypatch.setattr(dataset, '_download', lambda: None)
+        dataset._download()
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', lambda *args, **kwargs: None
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', lambda *args, **kwargs: None
+        )
path: tests/datasets/test_substation.py
user: login: rijuld
body: Hi @adamjstewart I tried removing this and running the tests, but the tests failed after removing this.
created_at: 2025-03-06T11:13:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+
+from torchgeo.datasets import DatasetNotFoundError, Substation
+
+
+class TestSubstation:
+    @pytest.fixture(
+        params=[
+            {'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': None,
+                'num_of_timepoints': 3,
+                'mask_2d': False,
+            },
+            {
+                'bands': [1, 2, 3],
+                'use_timepoints': True,
+                'timepoint_aggregation': None,
+                'num_of_timepoints': 5,
+                '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 dataset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Substation:
+        """Fixture for the Substation."""
+        root = os.path.join(os.getcwd(), 'tests', 'data', 'substation')
+        transforms = nn.Identity()
+        return Substation(
+            root=root,
+            bands=[1, 2, 3],
+            use_timepoints=True,
+            mask_2d=True,
+            timepoint_aggregation='median',
+            num_of_timepoints=4,
+            transforms=transforms,
+        )
+
+    def test_getitem_semantic(self, dataset: Substation) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x['image'], torch.Tensor)
+        assert isinstance(x['mask'], 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."""
+        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')
+        shutil.copy(url_for_images, tmp_path)
+        shutil.copy(url_for_masks, tmp_path)
+
+        monkeypatch.setattr(dataset, '_download', lambda: None)
+        dataset._download()
+
+    def test_download(
+        self, dataset: Substation, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Test the _download method of the dataset."""
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.download_url', lambda *args, **kwargs: None
+        )
+        monkeypatch.setattr(
+            'torchgeo.datasets.substation.extract_archive', lambda *args, **kwargs: None
+        )
path: tests/datasets/test_substation.py
user: login: adamjstewart
body: Probably because the data has already been downloaded and extracted. Try creating a new dataset with tmp_path as root and download the data from the test location (requires monkeypatching the URL). See all of our other dataset test files for examples of how we normally test this.
created_at: 2025-03-06T12:37:35Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,270 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+from typing import Literal
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Substation dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSM) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits).
+
+    Dataset Format:
+
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+    md5_images = '948706609864d0283f74ee7015f9d032'
+    md5_masks = 'baa369ececdc2ff80e6ba2b4c7fe147c'
+
+    def __init__(
+        self,
+        root: Path = 'data',
+        bands: Sequence[int] = tuple(range(13)),
+        mask_2d: bool = True,
+        timepoint_aggregation: Literal['concat', 'median', 'first', 'random']
+        | None = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = pd.Series(sorted(os.listdir(self.image_dir)))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            if image.shape[0] < self.num_of_timepoints:
+                # Padding: cycle through existing timepoints
+                padded_images = []
+                for i in range(self.num_of_timepoints):
+                    padded_images.append(image[i % image.shape[0]])
+                image = np.stack(padded_images)
+            elif image.shape[0] > self.num_of_timepoints:
+                # Removal: take the most recent timepoints
+                image = image[-self.num_of_timepoints :]
+
+            if self.timepoint_aggregation == 'concat':
+                # (num_of_timepoints*channels, h, w)
+                image = np.reshape(image, (-1, image.shape[2], image.shape[3]))
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: I don't fully understand this. It looks like `use_timepoints=True, timepoint_aggregation='first'` or `use_timepoints=False, timepoint_aggregation='concat'` are ignored? If we're just grouping concat/median and first/random together, maybe we can drop that argument and have a single if/elif/else for `timepoint_aggregation`?
created_at: 2025-03-12T16:06:39Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,270 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+from typing import Literal
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Substation dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSM) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits).
+
+    Dataset Format:
+
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+    md5_images = '948706609864d0283f74ee7015f9d032'
+    md5_masks = 'baa369ececdc2ff80e6ba2b4c7fe147c'
+
+    def __init__(
+        self,
+        root: Path = 'data',
+        bands: Sequence[int] = tuple(range(13)),
+        mask_2d: bool = True,
+        timepoint_aggregation: Literal['concat', 'median', 'first', 'random']
+        | None = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = pd.Series(sorted(os.listdir(self.image_dir)))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            if image.shape[0] < self.num_of_timepoints:
+                # Padding: cycle through existing timepoints
+                padded_images = []
+                for i in range(self.num_of_timepoints):
+                    padded_images.append(image[i % image.shape[0]])
+                image = np.stack(padded_images)
+            elif image.shape[0] > self.num_of_timepoints:
+                # Removal: take the most recent timepoints
+                image = image[-self.num_of_timepoints :]
+
+            if self.timepoint_aggregation == 'concat':
+                # (num_of_timepoints*channels, h, w)
+                image = np.reshape(image, (-1, image.shape[2], image.shape[3]))
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
path: torchgeo/datasets/substation.py
user: login: rijuld
body: > I don't fully understand this. It looks like `use_timepoints=True, timepoint_aggregation='first'` or `use_timepoints=False, timepoint_aggregation='concat'` are ignored? If we're just grouping concat/median and first/random together, maybe we can drop that argument and have a single if/elif/else for `timepoint_aggregation`?

Yeah that would be better
created_at: 2025-03-12T16:09:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,270 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Substation segmentation dataset."""
+
+import glob
+import os
+from collections.abc import Callable, Sequence
+from typing import Literal
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from .errors import DatasetNotFoundError
+from .geo import NonGeoDataset
+from .utils import Path, download_url, extract_archive
+
+
+class Substation(NonGeoDataset):
+    """Substation dataset.
+
+    The `Substation <https://github.com/Lindsay-Lab/substation-seg>`__
+    dataset is curated by TransitionZero and sourced from publicly
+    available data repositories, including OpenSreetMap (OSM) and
+    Copernicus Sentinel data. The dataset consists of Sentinel-2
+    images from 27k+ locations; the task is to segment power-substations,
+    which appear in the majority of locations in the dataset.
+    Most locations have 4-5 images taken at different timepoints
+    (i.e., revisits).
+
+    Dataset Format:
+
+    * .npz file for each datapoint
+
+    Dataset Features:
+
+    * 26,522 image-mask pairs stored as numpy files.
+    * Data from 5 revisits for most locations.
+    * Multi-temporal, multi-spectral images (13 channels) paired with masks,
+      with a spatial resolution of 228x228 pixels
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.48550/arXiv.2409.17363
+    """
+
+    directory = 'Substation'
+    filename_images = 'image_stack.tar.gz'
+    filename_masks = 'mask.tar.gz'
+    url_for_images = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/image_stack.tar.gz'
+    url_for_masks = 'https://storage.googleapis.com/tz-ml-public/substation-over-10km2-csv-main-444e360fd2b6444b9018d509d0e4f36e/mask.tar.gz'
+    md5_images = '948706609864d0283f74ee7015f9d032'
+    md5_masks = 'baa369ececdc2ff80e6ba2b4c7fe147c'
+
+    def __init__(
+        self,
+        root: Path = 'data',
+        bands: Sequence[int] = tuple(range(13)),
+        mask_2d: bool = True,
+        timepoint_aggregation: Literal['concat', 'median', 'first', 'random']
+        | None = 'concat',
+        download: bool = False,
+        checksum: bool = False,
+        transforms: Callable[[dict[str, Tensor]], dict[str, Tensor]] | None = None,
+        num_of_timepoints: int = 4,
+        use_timepoints: bool = False,
+    ) -> None:
+        """Initialize the Substation.
+
+        Args:
+            root: Path to the directory containing the dataset.
+            bands: Channels to use from the image.
+            mask_2d: Whether to use a 2D mask.
+            timepoint_aggregation: How to aggregate multiple timepoints.
+            download: Whether to download the dataset if it is not found.
+            checksum: Whether to verify the dataset after downloading.
+            transforms: A transform takes input sample and returns a transformed version.
+            num_of_timepoints: Number of timepoints to use for each image.
+            use_timepoints: Whether to use multiple timepoints for each image.
+        """
+        self.root = root
+        self.bands = bands
+        self.mask_2d = mask_2d
+        self.timepoint_aggregation = timepoint_aggregation
+        self.download = download
+        self.use_timepoints = use_timepoints
+        self.checksum = checksum
+        self.transforms = transforms
+        self.image_dir = os.path.join(root, 'image_stack')
+        self.mask_dir = os.path.join(root, 'mask')
+        self.num_of_timepoints = num_of_timepoints
+        self._verify()
+        self.image_filenames = pd.Series(sorted(os.listdir(self.image_dir)))
+
+    def __getitem__(self, index: int) -> dict[str, Tensor]:
+        """Get an item from the dataset by index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            A dictionary containing the image and corresponding mask.
+        """
+        image_filename = self.image_filenames[index]
+        image_path = os.path.join(self.image_dir, image_filename)
+        mask_path = os.path.join(self.mask_dir, image_filename)
+
+        image = np.load(image_path)['arr_0']
+
+        # selecting channels
+        image = image[:, self.bands, :, :]
+
+        # handling multiple images across timepoints
+        if self.use_timepoints:
+            if image.shape[0] < self.num_of_timepoints:
+                # Padding: cycle through existing timepoints
+                padded_images = []
+                for i in range(self.num_of_timepoints):
+                    padded_images.append(image[i % image.shape[0]])
+                image = np.stack(padded_images)
+            elif image.shape[0] > self.num_of_timepoints:
+                # Removal: take the most recent timepoints
+                image = image[-self.num_of_timepoints :]
+
+            if self.timepoint_aggregation == 'concat':
+                # (num_of_timepoints*channels, h, w)
+                image = np.reshape(image, (-1, image.shape[2], image.shape[3]))
+            elif self.timepoint_aggregation == 'median':
+                image = np.median(image, axis=0)
+        else:
+            if self.timepoint_aggregation == 'first':
+                image = image[0]
+            elif self.timepoint_aggregation == 'random':
+                image = image[np.random.randint(image.shape[0])]
path: torchgeo/datasets/substation.py
user: login: adamjstewart
body: Done, code looks **much** cleaner now.
created_at: 2025-03-12T16:18:53Z
author_association: MEMBER

PR comments:
user: login: adamjstewart
created_at: 2024-10-22T11:54:13Z
body: Hi @rijuld, thanks for the contribution! If you're new to creating PyTorch datasets, I highly recommend reading the following tutorials:

* https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
* https://pytorch.org/tutorials/beginner/data_loading_tutorial.html

The only difference between datasets in torchvision and NonGeoDatasets in TorchGeo is that our `__getitem__` returns a dictionary instead of a tuple. Other than that, they share all the same basic components.

Most of your issues seem to be due to the use of `args`. I think you just need to remove this and explicitly list all parameters in the function signature. This will also simplify your testing code. Take a look at other existing datasets, we have about 75 examples to choose from. If you find one that is similar to your dataset, it shouldn't actually require that many changes to get them working.
author_association: MEMBER
user: login: rijuld
created_at: 2024-10-22T12:05:16Z
body: Hi @adamjstewart , thanks a ton for the feedback! I will go through this tutorial. 
author_association: CONTRIBUTOR
user: login: rijuld
created_at: 2024-10-30T17:54:24Z
body: @microsoft-github-policy-service agree
author_association: CONTRIBUTOR
user: login: adamjstewart
created_at: 2024-11-26T12:35:23Z
body: Can you resolve the merge conflict? The order of `__all__` changed in the main branch.
author_association: MEMBER
user: login: adamjstewart
created_at: 2024-12-02T23:15:15Z
body: Still a lot of missing test coverage, will review in more detail after AGU.
author_association: MEMBER
user: login: rijuld
created_at: 2025-01-09T15:07:12Z
body: @adamjstewart I'm unsure about how to test the configuration file locally. Additionally, would testing the conf file cover the datamodule as well, or would I need to write separate test cases for the datamodule to ensure adequate code coverage?


author_association: CONTRIBUTOR
user: login: adamjstewart
created_at: 2025-01-11T18:32:52Z
body: > I'm unsure about how to test the configuration file locally.

You can get the test name from CI and then run it like so:
```console
> pytest tests/trainers/test_segmentation.py::TestSemanticSegmentationTask::test_trainer[True-substation]
```

> Additionally, would testing the conf file cover the datamodule as well

Yes
author_association: MEMBER
user: login: adamjstewart
created_at: 2025-03-12T11:46:55Z
body: We're preparing a 0.7.0 release in the next week or so and I would love to get this dataset in. Do you mind if I push changes directly to your branch to fix the remaining test comments?
author_association: MEMBER
user: login: rijuld
created_at: 2025-03-12T12:00:58Z
body: > We're preparing a 0.7.0 release in the next week or so and I would love to get this dataset in. Do you mind if I push changes directly to your branch to fix the remaining test comments?

Sure! Feel free to push changes directly to my branch. Thanks!

author_association: CONTRIBUTOR