Add plot method and data.py to NAIP

tests/data/naip/data.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+
+import numpy as np
+import rasterio
+from rasterio.crs import CRS
+from rasterio.transform import Affine
+
+SIZE = 128
+
+np.random.seed(0)
+
+# from Chesapeak data.py
+wkt = """
+PROJCS["USA_Contiguous_Albers_Equal_Area_Conic_USGS_version",
+    GEOGCS["NAD83",
+        DATUM["North_American_Datum_1983",
+            SPHEROID["GRS 1980",6378137,298.257222101004,
+                AUTHORITY["EPSG","7019"]],
+            AUTHORITY["EPSG","6269"]],
+        PRIMEM["Greenwich",0],
+        UNIT["degree",0.0174532925199433,
+            AUTHORITY["EPSG","9122"]],
+        AUTHORITY["EPSG","4269"]],
+    PROJECTION["Albers_Conic_Equal_Area"],
+    PARAMETER["latitude_of_center",23],
+    PARAMETER["longitude_of_center",-96],
+    PARAMETER["standard_parallel_1",29.5],
+    PARAMETER["standard_parallel_2",45.5],
+    PARAMETER["false_easting",0],
+    PARAMETER["false_northing",0],
+    UNIT["metre",1,
+        AUTHORITY["EPSG","9001"]],
+    AXIS["Easting",EAST],
+    AXIS["Northing",NORTH]]
+"""
+
+
+def create_file(path: str, dtype: str, num_channels: int) -> None:
+    profile = {}
+    profile["driver"] = "GTiff"
+    profile["dtype"] = dtype
+    profile["count"] = num_channels
+    profile["crs"] = CRS.from_wkt(wkt)
+    profile["transform"] = Affine(
+        1.0, 0.0, 1303555.0000000005, 0.0, -1.0, 2535064.999999998
+    )
+    profile["height"] = SIZE
+    profile["width"] = SIZE
+    profile["compress"] = "lzw"
+    profile["predictor"] = 2
+
+    if "float" in profile["dtype"]:
+        Z = np.random.randn(SIZE, SIZE).astype(profile["dtype"])
+    else:
+        Z = np.random.randint(
+            np.iinfo(profile["dtype"]).max, size=(SIZE, SIZE), dtype=profile["dtype"]
+        )
+
+    src = rasterio.open(path, "w", **profile)
+    for i in range(1, profile["count"] + 1):
+        src.write(Z, i)
+
+
+if __name__ == "__main__":
+    filenames = ["m_3807511_ne_18_060_20181104.tif", "m_3807511_ne_18_060_20190605.tif"]
+
+    for f in filenames:
+        create_file(os.path.join(os.getcwd(), f), "uint8", 4)

tests/datasets/test_naip.py

@@ -4,18 +4,18 @@
 import os
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 import pytest
 import torch
 import torch.nn as nn
-from _pytest.monkeypatch import MonkeyPatch
 from rasterio.crs import CRS
 
 from torchgeo.datasets import NAIP, BoundingBox, IntersectionDataset, UnionDataset
 
 
 class TestNAIP:
     @pytest.fixture
-    def dataset(self, monkeypatch: MonkeyPatch) -> NAIP:
+    def dataset(self) -> NAIP:
         root = os.path.join("tests", "data", "naip")
         transforms = nn.Identity()  # type: ignore[no-untyped-call]
         return NAIP(root, transforms=transforms)
@@ -34,6 +34,12 @@ def test_or(self, dataset: NAIP) -> None:
         ds = dataset | dataset
         assert isinstance(ds, UnionDataset)
 
+    def test_plot(self, dataset: NAIP) -> None:
+        query = dataset.bounds
+        x = dataset[query]
+        dataset.plot(x, suptitle="Test")
+        plt.close()
+
     def test_no_data(self, tmp_path: Path) -> None:
         with pytest.raises(FileNotFoundError, match="No NAIP data was found in "):
             NAIP(str(tmp_path))

torchgeo/datasets/naip.py

@@ -3,6 +3,10 @@
 
 """National Agriculture Imagery Program (NAIP) dataset."""
 
+from typing import Any, Dict, Optional
+
+import matplotlib.pyplot as plt
+
 from .geo import RasterDataset
 
 
@@ -42,3 +46,37 @@ class NAIP(RasterDataset):
     # Plotting
     all_bands = ["R", "G", "B", "NIR"]
     rgb_bands = ["R", "G", "B"]
+
+    def plot(
+        self,
+        sample: Dict[str, Any],
+        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 string to use as a suptitle
+
+        Returns:
+            a matplotlib Figure with the rendered sample
+
+        .. versionchanged:: 0.3
+            Method now takes a sample dict, not a Tensor. Additionally, possible to
+            show subplot titles and/or use a custom suptitle.
+        """
+        image = sample["image"][0:3, :, :].permute(1, 2, 0)
+
+        fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(4, 4))
+
+        ax.imshow(image)
+        ax.axis("off")
+        if show_titles:
+            ax.set_title("Image")
+
+        if suptitle is not None:
+            plt.suptitle(suptitle)
+
+        return fig