Updated ChesapeakeCVPR to use NLCD colormap

torchgeo/datasets/chesapeake.py

@@ -22,7 +22,7 @@
 from matplotlib.figure import Figure
 from rasterio.crs import CRS
 from torch import Tensor
-
+from .nlcd import NLCD
 from .geo import GeoDataset, RasterDataset
 from .utils import BoundingBox, download_url, extract_archive
 
@@ -574,15 +574,15 @@ def __init__(
         self.cache = cache
         self.download = download
         self.checksum = checksum
-
+        self.cmap = NLCD.cmap
         self._verify()
 
         super().__init__(transforms)
 
-        lc_colors = np.zeros((max(self.lc_cmap.keys()) + 1, 4))
-        lc_colors[list(self.lc_cmap.keys())] = list(self.lc_cmap.values())
-        lc_colors = lc_colors[:, :3] / 255
-        self._lc_cmap = ListedColormap(lc_colors)
+        nlcd_colors = np.zeros((max(self.nlcd_cmap.keys()) + 1, 4))
+        nlcd_colors[list(self.nlcd_cmap.keys())] = list(self.lc_cmap.values())
+        nlcd_colors = nlcd_colors[:, :3] / 255
+        self._nlcd_cmap = ListedColormap(nlcd_colors)
 
         nlcd_colors = np.zeros((max(self.nlcd_cmap.keys()) + 1, 4))
         nlcd_colors[list(self.nlcd_cmap.keys())] = list(self.nlcd_cmap.values())
@@ -618,7 +618,9 @@ def __init__(
                             "prior_from_cooccurrences_101_31_no_osm_no_buildings": prior_fn,  # noqa: E501
                         },
                     )
-
+    @property
+    def _nlcd_cmap(self):
+        return self.cmap
     def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         """Retrieve image/mask and metadata indexed by query.