Add multiple versions of the So2Sat dataset

tests/conf/so2sat_all.yaml

@@ -13,4 +13,5 @@ datamodule:
   root: "tests/data/so2sat"
   batch_size: 1
   num_workers: 0
+  version: "2"
   band_set: "all"

tests/conf/so2sat_rgb.yaml

@@ -0,0 +1,18 @@
+module:
+  _target_: torchgeo.trainers.ClassificationTask
+  loss: "ce"
+  model: "resnet18"
+  learning_rate: 1e-3
+  learning_rate_schedule_patience: 6
+  weights: null
+  in_channels: 3
+  num_classes: 17
+
+datamodule:
+  _target_: torchgeo.datamodules.So2SatDataModule
+  root: "tests/data/so2sat"
+  batch_size: 1
+  num_workers: 0
+  version: "3_random"
+  band_set: "rgb"
+  validation_pct: 0.5

tests/conf/so2sat_s1.yaml

@@ -13,4 +13,5 @@ datamodule:
   root: "tests/data/so2sat"
   batch_size: 1
   num_workers: 0
+  version: "2"
   band_set: "s1"

tests/data/so2sat/data.py

@@ -5,13 +5,14 @@
 
 import hashlib
 import os
+import shutil
 
 import h5py
 import numpy as np
 
-SIZE = 64  # image width/height
+SIZE = 32  # image width/height
 NUM_CLASSES = 17
-NUM_SAMPLES = 1
+NUM_SAMPLES = 2
 
 np.random.seed(0)
 
@@ -28,16 +29,21 @@
     ]
 
     # Random images
-    sen1 = np.random.randint(256, size=(NUM_SAMPLES, SIZE, SIZE, 8), dtype=np.uint8)
-    sen2 = np.random.randint(256, size=(NUM_SAMPLES, SIZE, SIZE, 10), dtype=np.uint8)
+    sen1 = np.random.randint(2, size=(NUM_SAMPLES, SIZE, SIZE, 8), dtype=np.uint8)
+    sen2 = np.random.randint(2, size=(NUM_SAMPLES, SIZE, SIZE, 10), dtype=np.uint8)
 
     # Create datasets
     with h5py.File(filename, "w") as f:
-        f.create_dataset("label", data=label)
-        f.create_dataset("sen1", data=sen1)
-        f.create_dataset("sen2", data=sen2)
+        f.create_dataset("label", data=label, compression="gzip", compression_opts=9)
+        f.create_dataset("sen1", data=sen1, compression="gzip", compression_opts=9)
+        f.create_dataset("sen2", data=sen2, compression="gzip", compression_opts=9)
 
     # Compute checksums
     with open(filename, "rb") as f:
         md5 = hashlib.md5(f.read()).hexdigest()
         print(repr(split.replace("ing", "")) + ":", repr(md5) + ",")
+
+for version in ["random", "block", "culture_10"]:
+    os.makedirs(version, exist_ok=True)
+    shutil.copyfile("training.h5", os.path.join(version, "training.h5"))
+    shutil.copyfile("testing.h5", os.path.join(version, "testing.h5"))

tests/datasets/test_so2sat.py

@@ -21,13 +21,15 @@
 class TestSo2Sat:
     @pytest.fixture(params=["train", "validation", "test"])
     def dataset(self, monkeypatch: MonkeyPatch, request: SubRequest) -> So2Sat:
-        md5s = {
-            "train": "82e0f2d51766b89cb905dbaf8275eb5b",
-            "validation": "bf292ae4737c1698b1a3c6f5e742e0e1",
-            "test": "9a3bbe181b038d4e51f122c4be3c569e",
+        md5s_by_version = {
+            "2": {
+                "train": "56e6fa0edb25b065124a3113372f76e5",
+                "validation": "940c95a737bd2fcdcc46c9a52b31424d",
+                "test": "e97a6746aadc731a1854097f32ab1755",
+            }
         }
 
-        monkeypatch.setattr(So2Sat, "md5s", md5s)
+        monkeypatch.setattr(So2Sat, "md5s_by_version", md5s_by_version)
         root = os.path.join("tests", "data", "so2sat")
         split = request.param
         transforms = nn.Identity()
@@ -51,13 +53,13 @@ def test_getitem(self, dataset: So2Sat) -> None:
         assert isinstance(x["label"], torch.Tensor)
 
     def test_len(self, dataset: So2Sat) -> None:
-        assert len(dataset) == 1
+        assert len(dataset) == 2
 
     def test_out_of_bounds(self, dataset: So2Sat) -> None:
         # h5py at version 2.10.0 raises a ValueError instead of an IndexError so we
         # check for both here
         with pytest.raises((IndexError, ValueError)):
-            dataset[1]
+            dataset[2]
 
     def test_invalid_split(self) -> None:
         with pytest.raises(AssertionError):

tests/trainers/test_classification.py

@@ -77,6 +77,7 @@ class TestClassificationTask:
             "so2sat_all",
             "so2sat_s1",
             "so2sat_s2",
+            "so2sat_rgb",
             "ucmerced",
         ],
     )

torchgeo/datamodules/so2sat.py

@@ -6,6 +6,8 @@
 from typing import Any
 
 import torch
+from torch import Generator, Tensor
+from torch.utils.data import random_split
 
 from ..datasets import So2Sat
 from .geo import NonGeoDataModule
@@ -14,79 +16,186 @@
 class So2SatDataModule(NonGeoDataModule):
     """LightningDataModule implementation for the So2Sat dataset.
 
-    Uses the train/val/test splits from the dataset.
+    If using the version 2 dataset, we use the train/val/test splits from the dataset.
+    If using the version 3 datasets, we use a random 80/20 train/val split from the
+    "train" set and use the "test" set as the test set.
     """
 
-    # TODO: calculate mean/std dev of s1 bands
-    mean = torch.tensor(
-        [
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.12375696117681859,
-            0.1092774636368323,
-            0.1010855203267882,
-            0.1142398616114001,
-            0.1592656692023089,
-            0.18147236008771792,
-            0.1745740312291377,
-            0.19501607349635292,
-            0.15428468872076637,
-            0.10905050699570007,
-        ]
-    )
-    std = torch.tensor(
-        [
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.0,
-            0.03958795985905458,
-            0.047778262752410296,
-            0.06636616706371974,
-            0.06358874912497474,
-            0.07744387147984592,
-            0.09101635085921553,
-            0.09218466562387101,
-            0.10164581233948201,
-            0.09991773043519253,
-            0.08780632509122865,
-        ]
-    )
+    means_per_version: dict[str, Tensor] = {
+        "2": torch.tensor(
+            [
+                -0.00003591224260,
+                -0.00000765856128,
+                0.00005937385750,
+                0.00002516623150,
+                0.04420110660000,
+                0.25761027100000,
+                0.00075567433700,
+                0.00135034668000,
+                0.12375696117681,
+                0.10927746363683,
+                0.10108552032678,
+                0.11423986161140,
+                0.15926566920230,
+                0.18147236008771,
+                0.17457403122913,
+                0.19501607349635,
+                0.15428468872076,
+                0.10905050699570,
+            ]
+        ),
+        "3_random": torch.tensor(
+            [
+                -0.00005541164581,
+                -0.00001363245448,
+                0.00004558943283,
+                0.00002990907940,
+                0.04451951629749,
+                0.25862310103671,
+                0.00032720731137,
+                0.00123416595462,
+                0.12428656593186,
+                0.11001677362564,
+                0.10230652367417,
+                0.11532195526186,
+                0.15989486018315,
+                0.18204406482475,
+                0.17513562590622,
+                0.19565546643221,
+                0.15648722649020,
+                0.11122536338577,
+            ]
+        ),
+        "3_block": torch.tensor(
+            [
+                -0.00004632368791,
+                0.00001260869365,
+                0.00005305557337,
+                0.00003471369557,
+                0.04449937686171,
+                0.26046026815721,
+                0.00087815394475,
+                0.00086889627435,
+                0.12381869777901,
+                0.10944155483024,
+                0.10176911573221,
+                0.11465267892206,
+                0.15870528223797,
+                0.18053964470203,
+                0.17366821871719,
+                0.19390983961551,
+                0.15536490486611,
+                0.11057334452833,
+            ]
+        ),
+    }
+    means_per_version["3_culture_10"] = means_per_version["2"]
+
+    stds_per_version: dict[str, Tensor] = {
+        "2": torch.tensor(
+            [
+                0.17555201,
+                0.17556463,
+                0.45998793,
+                0.45598876,
+                2.85599092,
+                8.32480061,
+                2.44987574,
+                1.46473530,
+                0.03958795,
+                0.04777826,
+                0.06636616,
+                0.06358874,
+                0.07744387,
+                0.09101635,
+                0.09218466,
+                0.10164581,
+                0.09991773,
+                0.08780632,
+            ]
+        ),
+        "3_random": torch.tensor(
+            [
+                0.1756914,
+                0.1761190,
+                0.4600589,
+                0.4563601,
+                2.2492179,
+                7.9056503,
+                2.1917633,
+                1.3148480,
+                0.0392269,
+                0.0470917,
+                0.0653264,
+                0.0624057,
+                0.0758367,
+                0.0891717,
+                0.0905092,
+                0.0996856,
+                0.0990188,
+                0.0873386,
+            ]
+        ),
+        "3_block": torch.tensor(
+            [
+                0.1751797,
+                0.1754073,
+                0.4610124,
+                0.4572122,
+                0.8294254,
+                7.1771026,
+                0.9642598,
+                0.8770835,
+                0.0388311,
+                0.0464986,
+                0.0643833,
+                0.0616141,
+                0.0753004,
+                0.0886178,
+                0.0899500,
+                0.0991759,
+                0.0983276,
+                0.0865943,
+            ]
+        ),
+    }
+    stds_per_version["3_culture_10"] = stds_per_version["2"]
 
     def __init__(
         self,
         batch_size: int = 64,
         num_workers: int = 0,
         band_set: str = "all",
+        val_split_pct: float = 0.2,
         **kwargs: Any,
     ) -> None:
         """Initialize a new So2SatDataModule instance.
 
         Args:
             batch_size: Size of each mini-batch.
             num_workers: Number of workers for parallel data loading.
-            band_set: One of 'all', 's1', or 's2'.
+            band_set: One of 'all', 's1', 's2', or 'rgb'.
+            val_split_pct: Percentage of training data to use for validation in with
+                the version 3 datasets.
             **kwargs: Additional keyword arguments passed to
                 :class:`~torchgeo.datasets.So2Sat`.
+
+        .. versionadded:: 0.5
+           The *val_split_pct* parameter, and the 'rgb' argument to *band_set*.
         """
+        version = kwargs.get("version", "2")
         kwargs["bands"] = So2Sat.BAND_SETS[band_set]
+        self.val_split_pct = val_split_pct
 
         if band_set == "s1":
-            self.mean = self.mean[:8]
-            self.std = self.std[:8]
+            self.mean = self.means_per_version[version][:8]
+            self.std = self.stds_per_version[version][:8]
         elif band_set == "s2":
-            self.mean = self.mean[8:]
-            self.std = self.std[8:]
+            self.mean = self.means_per_version[version][8:]
+            self.std = self.stds_per_version[version][8:]
+        elif band_set == "rgb":
+            self.mean = self.means_per_version[version][[10, 9, 8]]
+            self.std = self.stds_per_version[version][[10, 9, 8]]
 
         super().__init__(So2Sat, batch_size, num_workers, **kwargs)
 
@@ -100,9 +209,22 @@ def setup(self, stage: str) -> None:
         Args:
             stage: Either 'fit', 'validate', 'test', or 'predict'.
         """
-        if stage in ["fit"]:
-            self.train_dataset = So2Sat(split="train", **self.kwargs)
-        if stage in ["fit", "validate"]:
-            self.val_dataset = So2Sat(split="validation", **self.kwargs)
-        if stage in ["test"]:
-            self.test_dataset = So2Sat(split="test", **self.kwargs)
+        if self.kwargs.get("version", "2") == "2":
+            if stage in ["fit"]:
+                self.train_dataset = So2Sat(split="train", **self.kwargs)
+            if stage in ["fit", "validate"]:
+                self.val_dataset = So2Sat(split="validation", **self.kwargs)
+            if stage in ["test"]:
+                self.test_dataset = So2Sat(split="test", **self.kwargs)
+        else:
+            if stage in ["fit", "validate"]:
+                dataset = So2Sat(split="train", **self.kwargs)
+                val_length = round(len(dataset) * self.val_split_pct)
+                train_length = len(dataset) - val_length
+                self.train_dataset, self.val_dataset = random_split(
+                    dataset,
+                    [train_length, val_length],
+                    generator=Generator().manual_seed(0),
+                )
+            if stage in ["test"]:
+                self.test_dataset = So2Sat(split="test", **self.kwargs)