Add plot method to base chesapeake (microsoft#417)

* plot method to base chesapeake

* fix typo

* Add plotting to Chesapeake datasets

* isort, yousort, weallsort

* putting back the globs

* putting back the geo.py change

Co-authored-by: Caleb Robinson <calebrob6@gmail.com>

tests/datasets/test_chesapeake.py

@@ -82,7 +82,13 @@ def test_not_downloaded(self, tmp_path: Path) -> None:
     def test_plot(self, dataset: Chesapeake13) -> None:
         query = dataset.bounds
         x = dataset[query]
-        dataset.plot(x["mask"])
+        dataset.plot(x, suptitle="Test")
+
+    def test_plot_prediction(self, dataset: Chesapeake13) -> None:
+        query = dataset.bounds
+        x = dataset[query]
+        x["prediction"] = x["mask"].clone()
+        dataset.plot(x, suptitle="Prediction")
 
     def test_url(self) -> None:
         ds = Chesapeake13(os.path.join("tests", "data", "chesapeake", "BAYWIDE"))

torchgeo/datasets/chesapeake.py

@@ -9,14 +9,17 @@
 from typing import Any, Callable, Dict, Optional, Sequence
 
 import fiona
+import matplotlib.pyplot as plt
 import numpy as np
 import pyproj
 import rasterio
 import rasterio.mask
 import shapely.geometry
 import shapely.ops
 import torch
+from matplotlib.colors import ListedColormap
 from rasterio.crs import CRS
+from torch import Tensor
 
 from .geo import GeoDataset, RasterDataset
 from .utils import BoundingBox, download_url, extract_archive
@@ -46,6 +49,24 @@ class Chesapeake(RasterDataset, abc.ABC):
 
     is_image = False
 
+    # subclasses use the 13 class cmap by default
+    cmap = {
+        0: (0, 0, 0, 0),
+        1: (0, 197, 255, 255),
+        2: (0, 168, 132, 255),
+        3: (38, 115, 0, 255),
+        4: (76, 230, 0, 255),
+        5: (163, 255, 115, 255),
+        6: (255, 170, 0, 255),
+        7: (255, 0, 0, 255),
+        8: (156, 156, 156, 255),
+        9: (0, 0, 0, 255),
+        10: (115, 115, 0, 255),
+        11: (230, 230, 0, 255),
+        12: (255, 255, 115, 255),
+        13: (197, 0, 255, 255),
+    }
+
     @property
     @abc.abstractmethod
     def base_folder(self) -> str:
@@ -107,6 +128,17 @@ def __init__(
 
         self._verify()
 
+        colors = []
+        for i in range(len(self.cmap)):
+            colors.append(
+                (
+                    self.cmap[i][0] / 255.0,
+                    self.cmap[i][1] / 255.0,
+                    self.cmap[i][2] / 255.0,
+                )
+            )
+        self._cmap = ListedColormap(colors)
+
         super().__init__(root, crs, res, transforms, cache)
 
     def _verify(self) -> None:
@@ -144,6 +176,72 @@ def _extract(self) -> None:
         """Extract the dataset."""
         extract_archive(os.path.join(self.root, self.zipfile))
 
+    def plot(  # type: ignore[override]
+        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 returned by :meth:`RasterDataset.__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.3
+        """
+        mask = sample["mask"].squeeze(0)
+        ncols = 1
+
+        showing_predictions = "prediction" in sample
+        if showing_predictions:
+            pred = sample["prediction"].squeeze(0)
+            ncols = 2
+
+        fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(4 * ncols, 4))
+
+        if showing_predictions:
+            axs[0].imshow(
+                mask,
+                vmin=0,
+                vmax=self._cmap.N - 1,
+                cmap=self._cmap,
+                interpolation="none",
+            )
+            axs[0].axis("off")
+            axs[1].imshow(
+                pred,
+                vmin=0,
+                vmax=self._cmap.N - 1,
+                cmap=self._cmap,
+                interpolation="none",
+            )
+            axs[1].axis("off")
+            if show_titles:
+                axs[0].set_title("Mask")
+                axs[1].set_title("Prediction")
+
+        else:
+            axs.imshow(
+                mask,
+                vmin=0,
+                vmax=self._cmap.N - 1,
+                cmap=self._cmap,
+                interpolation="none",
+            )
+            axs.axis("off")
+            if show_titles:
+                axs.set_title("Mask")
+
+        if suptitle is not None:
+            plt.suptitle(suptitle)
+
+        return fig
+
 
 class Chesapeake7(Chesapeake):
     """Complete 7-class dataset.
@@ -175,14 +273,6 @@ class Chesapeake7(Chesapeake):
         5: (156, 156, 156, 255),
         6: (0, 0, 0, 255),
         7: (197, 0, 255, 255),
-        8: (0, 0, 0, 0),
-        9: (0, 0, 0, 0),
-        10: (0, 0, 0, 0),
-        11: (0, 0, 0, 0),
-        12: (0, 0, 0, 0),
-        13: (0, 0, 0, 0),
-        14: (0, 0, 0, 0),
-        15: (0, 0, 0, 0),
     }