Fix for Blank Images in plot due to Float Tensor Ranges (#1712)

* fix: blank image when using plot * black check
microsoft · Nov 6, 2023 · d1e5a0e · d1e5a0e
1 parent 57b028c
commit d1e5a0e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/torchgeo/datasets/ucmerced.py b/torchgeo/datasets/ucmerced.py
@@ -214,6 +214,11 @@ def plot(
         .. versionadded:: 0.2
         """
         image = np.rollaxis(sample["image"].numpy(), 0, 3)
+
+        # Normalize the image if the max value is greater than 1
+        if image.max() > 1:
+            image = image.astype(np.float32) / 255.0  # Scale to [0, 1]
+
         label = cast(int, sample["label"].item())
         label_class = self.classes[label]