Repository: microsoft/torchgeo

Issue:
Getting RGB images from EuroSAT
Hi!

In some datasets, it seems the bands can be selected when constructing but this not the case for EuroSAT. How can one get the normalized RGB images (similar to [`eurosat/rgb`](https://www.tensorflow.org/datasets/catalog/eurosat#eurosatrgb_default_config) version in tensorflow_datasets) ?


Discussion:
From the plotting code, I can do:

```python
img = eurosat[0]["image"]
(img[[3,2,1],:,:] / 3000).clip(0., 1.)
```
But this can still be nice to have the RGB version 
Related to #289, @nilsleh might be interested in this

Resolving patch:
diff --git a/torchgeo/datasets/eurosat.py b/torchgeo/datasets/eurosat.py
--- a/torchgeo/datasets/eurosat.py
+++ b/torchgeo/datasets/eurosat.py
@@ -4,10 +4,11 @@
 """EuroSAT dataset."""
 
 import os
-from typing import Callable, Dict, Optional, cast
+from typing import Callable, Dict, Optional, Sequence, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
+import torch
 from torch import Tensor
 
 from .geo import VisionClassificationDataset
@@ -83,10 +84,31 @@ class EuroSAT(VisionClassificationDataset):
         "Forest",
     ]
 
+    all_band_names = (
+        "B01",
+        "B02",
+        "B03",
+        "B04",
+        "B05",
+        "B06",
+        "B07",
+        "B08",
+        "B08A",
+        "B09",
+        "B10",
+        "B11",
+        "B12",
+    )
+
+    RGB_BANDS = ("B04", "B03", "B02")
+
+    BAND_SETS = {"all": all_band_names, "rgb": RGB_BANDS}
+
     def __init__(
         self,
         root: str = "data",
         split: str = "train",
+        bands: Sequence[str] = BAND_SETS["all"],
         transforms: Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]] = None,
         download: bool = False,
         checksum: bool = False,
@@ -96,19 +118,31 @@ def __init__(
         Args:
             root: root directory where dataset can be found
             split: one of "train", "val", or "test"
+            bands: a sequence of band names to load
             transforms: a function/transform that takes input sample and its target as
                 entry and returns a transformed version
             download: if True, download dataset and store it in the root directory
             checksum: if True, check the MD5 of the downloaded files (may be slow)
 
         Raises:
+            AssertionError: if ``split`` argument is invalid
             RuntimeError: if ``download=False`` and data is not found, or checksums
                 don't match
+
         """
         self.root = root
         self.transforms = transforms
         self.download = download
         self.checksum = checksum
+
+        assert split in ["train", "val", "test"]
+
+        self._validate_bands(bands)
+        self.bands = bands
+        self.band_indices = Tensor(
+            [self.all_band_names.index(b) for b in bands if b in self.all_band_names]
+        ).long()
+
         self._verify()
 
         valid_fns = set()
@@ -124,6 +158,26 @@ def __init__(
             is_valid_file=is_in_split,
         )
 
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        """Return an index within the dataset.
+
+        Args:
+            index: index to return
+        Returns:
+            data and label at that index
+        """
+        image, label = self._load_image(index)
+
+        image = torch.index_select(  # type: ignore[attr-defined]
+            image, dim=0, index=self.band_indices
+        )
+        sample = {"image": image, "label": label}
+
+        if self.transforms is not None:
+            sample = self.transforms(sample)
+
+        return sample
+
     def _check_integrity(self) -> bool:
         """Check integrity of dataset.
 
@@ -184,6 +238,23 @@ def _extract(self) -> None:
         filepath = os.path.join(self.root, self.filename)
         extract_archive(filepath)
 
+    def _validate_bands(self, bands: Sequence[str]) -> None:
+        """Validate list of bands.
+
+        Args:
+            bands: user-provided sequence of bands to load
+
+        Raises:
+            AssertionError: if ``bands`` is not a sequence
+            ValueError: if an invalid band name is provided
+
+        .. versionadded:: 0.3
+        """
+        assert isinstance(bands, Sequence), "'bands' must be a sequence"
+        for band in bands:
+            if band not in self.all_band_names:
+                raise ValueError(f"'{band}' is an invalid band name.")
+
     def plot(
         self,
         sample: Dict[str, Tensor],
@@ -200,9 +271,20 @@ def plot(
         Returns:
             a matplotlib Figure with the rendered sample
 
+        Raises:
+            ValueError: if RGB bands are not found in dataset
+
         .. versionadded:: 0.2
         """
-        image = np.rollaxis(sample["image"][[3, 2, 1]].numpy(), 0, 3).copy()
+        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")
+
+        image = np.take(sample["image"].numpy(), indices=rgb_indices, axis=0)
+        image = np.rollaxis(image, 0, 3)
         image = np.clip(image / 3000, 0, 1)
 
         label = cast(int, sample["label"].item())


Test patch:
diff --git a/tests/datasets/test_eurosat.py b/tests/datasets/test_eurosat.py
--- a/tests/datasets/test_eurosat.py
+++ b/tests/datasets/test_eurosat.py
@@ -58,7 +58,9 @@ def dataset(
         root = str(tmp_path)
         split = request.param
         transforms = nn.Identity()  # type: ignore[attr-defined]
-        return EuroSAT(root, split, transforms, download=True, checksum=True)
+        return EuroSAT(
+            root=root, split=split, transforms=transforms, download=True, checksum=True
+        )
 
     def test_getitem(self, dataset: EuroSAT) -> None:
         x = dataset[0]
@@ -66,6 +68,14 @@ def test_getitem(self, dataset: EuroSAT) -> None:
         assert isinstance(x["image"], torch.Tensor)
         assert isinstance(x["label"], torch.Tensor)
 
+    def test_invalid_split(self) -> None:
+        with pytest.raises(AssertionError):
+            EuroSAT(split="foo")
+
+    def test_invalid_bands(self) -> None:
+        with pytest.raises(ValueError):
+            EuroSAT(bands=("OK", "BK"))
+
     def test_len(self, dataset: EuroSAT) -> None:
         assert len(dataset) == 2
 
@@ -100,3 +110,8 @@ def test_plot(self, dataset: EuroSAT) -> None:
         x["prediction"] = x["label"].clone()
         dataset.plot(x)
         plt.close()
+
+    def test_plot_rgb(self, dataset: EuroSAT, tmp_path: Path) -> None:
+        dataset = EuroSAT(root=str(tmp_path), bands=("B03",))
+        with pytest.raises(ValueError, match="doesn't contain some of the RGB bands"):
+            dataset.plot(dataset[0], suptitle="Single Band")