Make RasterDataset not override custom cmaps (#422)

* Make RasterDataset not override custom cmaps Fixes #421 * Move if check outside of try for speed
microsoft · Mar 19, 2022 · 70578db · 70578db
1 parent bf848c1
commit 70578db
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
@@ -348,10 +348,11 @@ def __init__(
                 try:
                     with rasterio.open(filepath) as src:
                         # See if file has a color map
-                        try:
-                            self.cmap = src.colormap(1)
-                        except ValueError:
-                            pass
+                        if len(self.cmap) == 0:
+                            try:
+                                self.cmap = src.colormap(1)
+                            except ValueError:
+                                pass
 
                         if crs is None:
                             crs = src.crs