Use gray cmap

microsoft · Dec 17, 2024 · f10d198 · f10d198
1 parent 460556c
commit f10d198
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchgeo/datasets/gbm.py b/torchgeo/datasets/gbm.py
@@ -52,13 +52,13 @@ def plot(
         ncols = 2 if 'prediction' in sample else 1
         fig, axs = plt.subplots(ncols=ncols, squeeze=False)
 
-        axs[0, 0].imshow(sample['mask'])
+        axs[0, 0].imshow(sample['mask'], cmap='gray')
         axs[0, 0].axis('off')
         if show_titles:
             axs[0, 0].set_title('Mask')
 
         if 'prediction' in sample:
-            axs[0, 1].imshow(sample['prediction'])
+            axs[0, 1].imshow(sample['prediction'], cmap='gray')
             axs[0, 1].axis('off')
             if show_titles:
                 axs[0, 1].set_title('Prediction')