SSL4EO-S12: add additional metadata

Author: adamjstewart

torchgeo/datasets/sentinel.py

@@ -4,7 +4,7 @@
 """Sentinel datasets."""
 
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any
+from typing import Any, ClassVar
 
 import matplotlib.pyplot as plt
 import torch
@@ -28,6 +28,9 @@ class Sentinel(RasterDataset):
     * https://asf.alaska.edu/datasets/daac/sentinel-1/
     """
 
+    date_format = '%Y%m%dT%H%M%S'
+    separate_files = True
+
 
 class Sentinel1(Sentinel):
     r"""Sentinel-1 dataset.
@@ -136,9 +139,11 @@ class Sentinel1(Sentinel):
         _(?P<band>[VH]{2})
         \.
     """
-    date_format = '%Y%m%dT%H%M%S'
+
+    # https://sentiwiki.copernicus.eu/web/s1-mission
     all_bands = ('HH', 'HV', 'VV', 'VH')
-    separate_files = True
+    # Central wavelength (μm)
+    wavelength = 55500
 
     def __init__(
         self,
@@ -274,10 +279,9 @@ class Sentinel2(Sentinel):
         (?:_(?P<resolution>{}m))?
         \..*$
     """
-    date_format = '%Y%m%dT%H%M%S'
 
-    # https://gisgeography.com/sentinel-2-bands-combinations/
-    all_bands: tuple[str, ...] = (
+    # https://sentiwiki.copernicus.eu/web/s2-mission
+    all_bands = (
         'B01',
         'B02',
         'B03',
@@ -293,8 +297,32 @@ class Sentinel2(Sentinel):
         'B12',
     )
     rgb_bands = ('B04', 'B03', 'B02')
-
-    separate_files = True
+    # Central wavelength (μm)
+    wavelengths: ClassVar[dict[str, float]] = {
+        'B01': 0.4427,
+        'B02': 0.4927,
+        'B03': 0.5598,
+        'B04': 0.6646,
+        'B05': 0.7041,
+        'B06': 0.7405,
+        'B07': 0.7828,
+        'B08': 0.8328,
+        'B8A': 0.8647,
+        'B09': 0.9451,
+        'B10': 1.3735,
+        'B11': 1.6137,
+        'B12': 2.2024,
+        # For compatibility with other dataset naming conventions
+        'B1': 0.4427,
+        'B2': 0.4927,
+        'B3': 0.5598,
+        'B4': 0.6646,
+        'B5': 0.7041,
+        'B6': 0.7405,
+        'B7': 0.7828,
+        'B8': 0.8328,
+        'B9': 0.9451,
+    }
 
     def __init__(
         self,

torchgeo/datasets/ssl4eo.py

@@ -6,8 +6,9 @@
 import glob
 import os
 import random
+import re
 from collections.abc import Callable
-from typing import ClassVar, TypedDict
+from typing import Any, ClassVar, TypedDict
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -18,7 +19,15 @@
 
 from .errors import DatasetNotFoundError
 from .geo import NonGeoDataset
-from .utils import Path, check_integrity, download_url, extract_archive
+from .sentinel import Sentinel, Sentinel1, Sentinel2
+from .utils import (
+    BoundingBox,
+    Path,
+    check_integrity,
+    disambiguate_timestamp,
+    download_url,
+    extract_archive,
+)
 
 
 class SSL4EO(NonGeoDataset):
@@ -321,7 +330,7 @@ def plot(
         return fig
 
 
-class SSL4EOS12(NonGeoDataset):
+class SSL4EOS12(SSL4EO):
     """SSL4EO-S12 dataset.
 
     `Sentinel-1/2 <https://github.com/zhu-xlab/SSL4EO-S12>`_ version of SSL4EO.
@@ -362,6 +371,7 @@ class _Metadata(TypedDict):
             'filename': 's1.tar.gz',
             'md5': '51ee23b33eb0a2f920bda25225072f3a',
             'bands': ['VV', 'VH'],
+            'filename_regex': r'^S1[AB]_(?P<mode>SM|IW|EW|WV)_.{9}_(?P<date>\d{8}T\d{6})',
         },
         's2c': {
             'filename': 's2_l1c.tar.gz',
@@ -381,6 +391,7 @@ class _Metadata(TypedDict):
                 'B11',
                 'B12',
             ],
+            'filename_regex': r'^(?P<date>\d{8}T\d{6})',
         },
         's2a': {
             'filename': 's2_l2a.tar.gz',
@@ -399,6 +410,7 @@ class _Metadata(TypedDict):
                 'B11',
                 'B12',
             ],
+            'filename_regex': r'^(?P<date>\d{8}T\d{6})',
         },
     }
 
@@ -439,7 +451,7 @@ def __init__(
 
         self._verify()
 
-    def __getitem__(self, index: int) -> dict[str, Tensor]:
+    def __getitem__(self, index: int) -> dict[str, Any]:
         """Return an index within the dataset.
 
         Args:
@@ -451,17 +463,37 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
         root = os.path.join(self.root, self.split, f'{index:07}')
         subdirs = os.listdir(root)
         subdirs = random.sample(subdirs, self.seasons)
+        filename_regex = self.metadata[self.split]['filename_regex']
 
         images = []
+        bounds = []
+        wavelengths = []
         for subdir in subdirs:
             directory = os.path.join(root, subdir)
-            for band in self.bands:
-                filename = os.path.join(directory, f'{band}.tif')
-                with rasterio.open(filename) as f:
-                    image = f.read(out_shape=(1, self.size, self.size))
-                    images.append(torch.from_numpy(image.astype(np.float32)))
-
-        sample = {'image': torch.cat(images)}
+            if match := re.match(filename_regex, subdir):
+                date_str = match.group('date')
+                mint, maxt = disambiguate_timestamp(date_str, Sentinel.date_format)
+                for band in self.bands:
+                    match self.split:
+                        case 's1':
+                            wavelengths.append(Sentinel1.wavelength)
+                        case 's2c' | 's2a':
+                            wavelengths.append(Sentinel2.wavelengths[band])
+
+                    filename = os.path.join(directory, f'{band}.tif')
+                    with rasterio.open(filename) as f:
+                        minx, maxx = f.bounds.left, f.bounds.right
+                        miny, maxy = f.bounds.bottom, f.bounds.top
+                        image = f.read(out_shape=(1, self.size, self.size))
+                        images.append(torch.from_numpy(image.astype(np.float32)))
+                bounds.append(BoundingBox(minx, maxx, miny, maxy, mint, maxt))
+
+        sample = {
+            'image': torch.cat(images),
+            'bounds': bounds,
+            'wavelengths': wavelengths,
+            'gsd': 10,
+        }
 
         if self.transforms is not None:
             sample = self.transforms(sample)