Add plotting method for CV4A Kenya Crop Type Dataset

tests/datasets/test_cv4a_kenya_crop_type.py

@@ -7,9 +7,11 @@
 from pathlib import Path
 from typing import Generator
 
+import matplotlib.pyplot as plt
 import pytest
 import torch
 import torch.nn as nn
+from _pytest.fixtures import SubRequest
 from _pytest.monkeypatch import MonkeyPatch
 from torch.utils.data import ConcatDataset
 
@@ -30,9 +32,12 @@ def fetch(dataset_id: str, **kwargs: str) -> Dataset:
 
 
 class TestCV4AKenyaCropType:
-    @pytest.fixture
+    @pytest.fixture(params=[tuple(["B01"]), tuple(CV4AKenyaCropType.band_names)])
     def dataset(
-        self, monkeypatch: Generator[MonkeyPatch, None, None], tmp_path: Path
+        self,
+        monkeypatch: Generator[MonkeyPatch, None, None],
+        tmp_path: Path,
+        request: SubRequest,
     ) -> CV4AKenyaCropType:
         radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
         monkeypatch.setattr(  # type: ignore[attr-defined]
@@ -56,6 +61,7 @@ def dataset(
         transforms = nn.Identity()  # type: ignore[attr-defined]
         return CV4AKenyaCropType(
             root,
+            bands=request.param,
             transforms=transforms,
             download=True,
             api_key="",
@@ -113,3 +119,17 @@ def test_invalid_bands(self) -> None:
 
         with pytest.raises(ValueError, match="is an invalid band name."):
             CV4AKenyaCropType(bands=("foo", "bar"))
+
+    def test_plot(self, dataset: CV4AKenyaCropType) -> None:
+        if not all(band in dataset.bands for band in dataset.RGB_BANDS):
+            with pytest.raises(ValueError, match="Dataset doesn't contain"):
+                x = dataset[0].copy()
+                dataset.plot(x, time_step=0, suptitle="Test")
+        else:
+            dataset.plot(dataset[0], time_step=0, suptitle="Test")
+            plt.close()
+
+            sample = dataset[0]
+            sample["prediction"] = sample["mask"].clone()
+            dataset.plot(sample, time_step=0, suptitle="Pred")
+            plt.close()

torchgeo/datasets/cv4a_kenya_crop_type.py

@@ -8,6 +8,7 @@
 from functools import lru_cache
 from typing import Callable, Dict, List, Optional, Tuple
 
+import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from PIL import Image
@@ -102,6 +103,8 @@ class CV4AKenyaCropType(VisionDataset):
         "CLD",
     )
 
+    RGB_BANDS = ["B04", "B03", "B02"]
+
     # Same for all tiles
     tile_height = 3035
     tile_width = 2016
@@ -400,3 +403,69 @@ def _download(self, api_key: Optional[str] = None) -> None:
         target_archive_path = os.path.join(self.root, self.target_meta["filename"])
         for fn in [image_archive_path, target_archive_path]:
             extract_archive(fn, self.root)
+
+    def plot(
+        self,
+        sample: Dict[str, Tensor],
+        time_step: int = 0,
+        suptitle: Optional[str] = None,
+    ) -> plt.Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample return by :meth:`__getitem__`
+            time_step: time step at which to access image, beginning with 0
+            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")
+
+        if "prediction" in sample:
+            prediction = sample["prediction"]
+            n_cols = 3
+        else:
+            n_cols = 2
+
+        image, mask = sample["image"], sample["mask"]
+
+        assert (
+            time_step <= image.shape[0] - 1
+        ), "The specified time step \
+            does not exist, image onyl contains {} time \
+                instances".format(
+            image.shape[0]
+        )
+
+        image = image[time_step, rgb_indices, :, :]
+        print(image.max())
+        print(image.min())
+        print(image.shape)
+        print(mask.shape)
+
+        fig, axs = plt.subplots(nrows=1, ncols=n_cols, figsize=(10, n_cols * 5))
+
+        axs[0].imshow(image.permute(1, 2, 0))
+        axs[0].axis("off")
+        axs[0].set_title("Image")
+        axs[1].imshow(mask)
+        axs[1].axis("off")
+        axs[1].set_title("Mask")
+
+        if "prediction" in sample:
+            axs[2].imshow(prediction)
+            axs[2].axis("off")
+            axs[2].set_title("Prediction")
+
+        if suptitle is not None:
+            plt.suptitle(suptitle)
+
+        return fig