Source code for torchgeo.datamodules.seco
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
"""Seasonal Contrast datamodule."""
from typing import Any
import kornia.augmentation as K
import torch
from einops import repeat
from ..datasets import SeasonalContrastS2
from ..transforms import AugmentationSequential
from .geo import NonGeoDataModule
[docs]class SeasonalContrastS2DataModule(NonGeoDataModule):
"""LightningDataModule implementation for the Seasonal Contrast dataset.
.. versionadded:: 0.5
"""
[docs] def __init__(
self, batch_size: int = 64, num_workers: int = 0, **kwargs: Any
) -> None:
"""Initialize a new SeasonalContrastS2DataModule instance.
Args:
batch_size: Size of each mini-batch.
num_workers: Number of workers for parallel data loading.
**kwargs: Additional keyword arguments passed to
:class:`~torchgeo.datasets.SeasonalContrastS2`.
"""
super().__init__(SeasonalContrastS2, batch_size, num_workers, **kwargs)
bands = kwargs.get('bands', SeasonalContrastS2.rgb_bands)
seasons = kwargs.get('seasons', 1)
# Normalization only available for RGB dataset, defined here:
# https://github.com/ServiceNow/seasonal-contrast/blob/8285173ec205b64bc3e53b880344dd6c3f79fa7a/datasets/seco_dataset.py
if bands == SeasonalContrastS2.rgb_bands:
_min = torch.tensor([3, 2, 0])
_max = torch.tensor([88, 103, 129])
_mean = torch.tensor([0.485, 0.456, 0.406])
_std = torch.tensor([0.229, 0.224, 0.225])
_min = repeat(_min, 'c -> (t c)', t=seasons)
_max = repeat(_max, 'c -> (t c)', t=seasons)
_mean = repeat(_mean, 'c -> (t c)', t=seasons)
_std = repeat(_std, 'c -> (t c)', t=seasons)
self.aug = AugmentationSequential(
K.Normalize(mean=_min, std=_max - _min),
K.Normalize(mean=torch.tensor(0), std=1 / torch.tensor(255)),
K.Normalize(mean=_mean, std=_std),
data_keys=['image'],
)
[docs] def setup(self, stage: str) -> None:
"""Set up datasets.
Args:
stage: Either 'fit', 'validate', 'test', or 'predict'.
"""
self.dataset = SeasonalContrastS2(**self.kwargs)