Add plot method and new band names for Sen12MS dataset

tests/datasets/test_sen12ms.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Generator
 
+import matplotlib.pyplot as plt
 import pytest
 import torch
 import torch.nn as nn
@@ -83,6 +84,19 @@ def test_band_subsets(self) -> None:
             x = ds[0]["image"]
             assert x.shape[0] == len(bands)
 
+    def test_invalid_bands(self) -> None:
+        with pytest.raises(ValueError):
+            SEN12MS(bands=tuple(["OK", "BK"]))
+
+    def test_plot(self, dataset: SEN12MS) -> None:
+        dataset.plot(dataset[0], suptitle="Test")
+        plt.close()
+
+        sample = dataset[0]
+        sample["prediction"] = sample["mask"].clone()
+        dataset.plot(sample, suptitle="prediction")
+        plt.close()
+
 
 class TestSEN12MSDataModule:
     @pytest.fixture(scope="class", params=["all", "s1", "s2-all", "s2-reduced"])

torchgeo/datasets/sen12ms.py

@@ -4,8 +4,9 @@
 """SEN12MS dataset."""
 
 import os
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, Optional, Tuple
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pytorch_lightning as pl
 import rasterio
@@ -69,50 +70,113 @@ class SEN12MS(VisionDataset):
        This download will likely take several hours.
     """  # noqa: E501
 
-    BAND_SETS: Dict[str, List[int]] = {
-        "all": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
-        "s1": [0, 1],
-        "s2-all": [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
-        "s2-reduced": [3, 4, 5, 9, 12, 13],
+    # BAND_SETS: Dict[str, List[int]] = {
+    #     "all": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+    #     "s1": [0, 1],
+    #     "s2-all": [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+    #     "s2-reduced": [3, 4, 5, 9, 12, 13],
+    # }
+
+    BAND_SETS: Dict[str, Tuple[str, ...]] = {
+        "all": tuple(
+            [
+                "VV",
+                "VH",
+                "B01",
+                "B02",
+                "B03",
+                "B04",
+                "B05",
+                "B06",
+                "B07",
+                "B08",
+                "B8A",
+                "B09",
+                "B10",
+                "B11",
+                "B12",
+            ]
+        ),
+        "s1": tuple(["VV", "VH"]),
+        "s2-all": tuple(
+            [
+                "B01",
+                "B02",
+                "B03",
+                "B04",
+                "B05",
+                "B06",
+                "B07",
+                "B08",
+                "B8A",
+                "B09",
+                "B10",
+                "B11",
+                "B12",
+            ]
+        ),
+        "s2-reduced": tuple(["B02", "B03", "B04", "B08", "B10", "B11"]),
     }
 
+    band_names = tuple(
+        [
+            "VV",
+            "VH",
+            "B01",
+            "B02",
+            "B03",
+            "B04",
+            "B05",
+            "B06",
+            "B07",
+            "B08",
+            "B8A",
+            "B09",
+            "B10",
+            "B11",
+            "B12",
+        ]
+    )
+
+    RGB_BANDS = ["B04", "B03", "B02"]
+
     filenames = [
         "ROIs1158_spring_lc.tar.gz",
         "ROIs1158_spring_s1.tar.gz",
         "ROIs1158_spring_s2.tar.gz",
-        "ROIs1868_summer_lc.tar.gz",
-        "ROIs1868_summer_s1.tar.gz",
-        "ROIs1868_summer_s2.tar.gz",
-        "ROIs1970_fall_lc.tar.gz",
-        "ROIs1970_fall_s1.tar.gz",
-        "ROIs1970_fall_s2.tar.gz",
-        "ROIs2017_winter_lc.tar.gz",
-        "ROIs2017_winter_s1.tar.gz",
-        "ROIs2017_winter_s2.tar.gz",
+        # "ROIs1868_summer_lc.tar.gz",
+        # "ROIs1868_summer_s1.tar.gz",
+        # "ROIs1868_summer_s2.tar.gz",
+        # "ROIs1970_fall_lc.tar.gz",
+        # "ROIs1970_fall_s1.tar.gz",
+        # "ROIs1970_fall_s2.tar.gz",
+        # "ROIs2017_winter_lc.tar.gz",
+        # "ROIs2017_winter_s1.tar.gz",
+        # "ROIs2017_winter_s2.tar.gz",
         "train_list.txt",
         "test_list.txt",
     ]
     light_filenames = [
         "ROIs1158_spring",
-        "ROIs1868_summer",
-        "ROIs1970_fall",
-        "ROIs2017_winter",
+        # "ROIs1868_summer",
+        # "ROIs1970_fall",
+        # "ROIs2017_winter",
         "train_list.txt",
         "test_list.txt",
     ]
     md5s = [
         "6e2e8fa8b8cba77ddab49fd20ff5c37b",
         "fba019bb27a08c1db96b31f718c34d79",
         "d58af2c15a16f376eb3308dc9b685af2",
-        "2c5bd80244440b6f9d54957c6b1f23d4",
-        "01044b7f58d33570c6b57fec28a3d449",
-        "4dbaf72ecb704a4794036fe691427ff3",
-        "9b126a68b0e3af260071b3139cb57cee",
-        "19132e0aab9d4d6862fd42e8e6760847",
-        "b8f117818878da86b5f5e06400eb1866",
-        "0fa0420ef7bcfe4387c7e6fe226dc728",
-        "bb8cbfc16b95a4f054a3d5380e0130ed",
-        "3807545661288dcca312c9c538537b63",
+        # "2c5bd80244440b6f9d54957c6b1f23d4",
+        # "01044b7f58d33570c6b57fec28a3d449",
+        # "4dbaf72ecb704a4794036fe691427ff3",
+        # "9b126a68b0e3af260071b3139cb57cee",
+        # "19132e0aab9d4d6862fd42e8e6760847",
+        # "b8f117818878da86b5f5e06400eb1866",
+        # "0fa0420ef7bcfe4387c7e6fe226dc728",
+        # "bb8cbfc16b95a4f054a3d5380e0130ed",
+        # "3807545661288dcca312c9c538537b63",
         "0a68d4e1eb24f128fccdb930000b2546",
         "c7faad064001e646445c4c634169484d",
     ]
@@ -121,7 +185,7 @@ def __init__(
         self,
         root: str = "data",
         split: str = "train",
-        bands: List[int] = BAND_SETS["all"],
+        bands: Tuple[str, ...] = BAND_SETS["all"],
         transforms: Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]] = None,
         checksum: bool = False,
     ) -> None:
@@ -147,9 +211,14 @@ def __init__(
         """
         assert split in ["train", "test"]
 
+        self._validate_bands(bands)
+        self.band_indices = torch.tensor(  # type: ignore[attr-defined]
+            [self.band_names.index(b) for b in bands]
+        ).long()
+        self.bands = bands
+
         self.root = root
         self.split = split
-        self.bands = torch.tensor(bands).long()  # type: ignore[attr-defined]
         self.transforms = transforms
         self.checksum = checksum
 
@@ -180,7 +249,7 @@ def __getitem__(self, index: int) -> Dict[str, Tensor]:
 
         image = torch.cat(tensors=[s1, s2], dim=0)  # type: ignore[attr-defined]
         image = torch.index_select(  # type: ignore[attr-defined]
-            image, dim=0, index=self.bands
+            image, dim=0, index=self.band_indices
         )
 
         sample: Dict[str, Tensor] = {"image": image, "mask": lc}
@@ -223,12 +292,30 @@ def _load_raster(self, filename: str, source: str) -> Tensor:
             tensor: Tensor = torch.from_numpy(array)  # type: ignore[attr-defined]
             return tensor
 
+    def _validate_bands(self, bands: Tuple[str, ...]) -> None:
+        """Validate list of bands.
+
+        Args:
+            bands: user-provided tuple of bands to load
+
+        Raises:
+            AssertionError: if ``bands`` is not a tuple
+            ValueError: if an invalid band name is provided
+        """
+        assert isinstance(bands, tuple), "The list of bands must be a tuple"
+        for band in bands:
+            if band not in self.band_names:
+                raise ValueError(f"'{band}' is an invalid band name.")
+
     def _check_integrity_light(self) -> bool:
         """Checks the integrity of the dataset structure.
 
         Returns:
             True if the dataset directories and split files are found, else False
         """
+        import pdb
+
+        pdb.set_trace()
         for filename in self.light_filenames:
             filepath = os.path.join(self.root, filename)
             if not os.path.exists(filepath):
@@ -247,6 +334,63 @@ def _check_integrity(self) -> bool:
                 return False
         return True
 
+    def plot(
+        self,
+        sample: Dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: Optional[str] = None,
+    ) -> plt.Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample return by :meth:`__getitem__`
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional suptitle to use for figure
+
+        Returns;
+            a matplotlib Figure with the rendered sample
+
+        .. versionadded:: 0.2
+        """
+        rgb_indices = []
+        for band in self.RGB_BANDS:
+            if band in self.bands:
+                rgb_indices.append(self.bands.index(band))
+            else:
+                raise ValueError("Dataset doesn't contain some of the RGB bands")
+        import pdb
+
+        pdb.set_trace()
+        image, _ = sample["image"][rgb_indices, ...], sample["mask"]
+        ncols = 2
+
+        showing_predictions = "prediction" in sample
+        if showing_predictions:
+            _ = sample["prediction"]
+            ncols += 1
+
+        fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(10, ncols * 5))
+
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[0].axis("off")
+        # axs[1].imshow(mask)
+        # axs[1].axis("off")
+
+        if showing_predictions:
+            # axs[2].imshow(prediction)
+            axs[2].axis("off")
+
+        if show_titles:
+            axs[0].set_title("Image")
+            axs[1].set_title("Mask")
+            if showing_predictions:
+                axs[2].set_title("Prediction")
+
+        if suptitle is not None:
+            plt.suptitle(suptitle)
+
+        return fig
+
 
 class SEN12MSDataModule(pl.LightningDataModule):
     """LightningDataModule implementation for the SEN12MS dataset.