Use datamodule plot instead of dataset plot

torchgeo/datamodules/bigearthnet.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -176,3 +177,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.BigEarthNet.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/cowc.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 from torch import Generator  # type: ignore[attr-defined]
 from torch.utils.data import DataLoader, random_split
@@ -121,3 +122,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.COWC.plot`."""
+        return self.val_dataset.dataset.plot(*args, **kwargs)

torchgeo/datamodules/etci2021.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch import Generator  # type: ignore[attr-defined]
@@ -149,3 +150,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.ETCI2021.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/eurosat.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -146,3 +147,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.EuroSAT.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/landcoverai.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, Optional
 
 import kornia.augmentation as K
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 from torch.utils.data import DataLoader
 
@@ -158,3 +159,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.LandCoverAI.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/nasa_marine_debris.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, List, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch import Tensor
@@ -138,3 +139,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             shuffle=False,
             collate_fn=collate_fn,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.NASAMarineDebris.plot`."""
+        return self.val_dataset.dataset.plot(*args, **kwargs)

torchgeo/datamodules/resisc45.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, Optional
 
 import kornia.augmentation as K
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -156,3 +157,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.RESISC45.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/so2sat.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional, cast
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -185,18 +186,6 @@ def setup(self, stage: Optional[str] = None) -> None:
                 So2Sat, temp_train + self.val_dataset + self.test_dataset
             )
 
-        # So2Sat dataset doesn't know how to plot any band set other than "all"
-        # TODO: move band selection to the Dataset level so that plot knows about it
-        if self.bands == "rgb":
-            # delattr doesn't work for some reason
-            # https://stackoverflow.com/a/1684219/5828163
-            def noattr() -> None:
-                raise AttributeError
-
-            self.val_dataset.plot = (  # type: ignore[assignment]
-                lambda *args, **kwargs: noattr()
-            )
-
     def train_dataloader(self) -> DataLoader[Any]:
         """Return a DataLoader for training.
 
@@ -235,3 +224,12 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.NASAMarineDebris.plot`."""
+        # So2Sat dataset doesn't know how to plot any band set other than "all"
+        # TODO: move band selection to the Dataset level so that plot knows about it
+        if self.bands == "rgb":
+            raise AttributeError
+
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/datamodules/ucmerced.py

@@ -5,6 +5,7 @@
 
 from typing import Any, Dict, Optional
 
+import matplotlib.pyplot as plt
 import pytorch_lightning as pl
 import torch
 import torchvision
@@ -123,3 +124,7 @@ def test_dataloader(self) -> DataLoader[Any]:
             num_workers=self.num_workers,
             shuffle=False,
         )
+
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Run :meth:`torchgeo.datasets.UCMerced.plot`."""
+        return self.val_dataset.plot(*args, **kwargs)

torchgeo/trainers/classification.py

@@ -184,7 +184,7 @@ def validation_step(  # type: ignore[override]
             try:
                 datamodule = self.trainer.datamodule  # type: ignore[attr-defined]
                 sample = unbind_samples(batch)[0]
-                fig = datamodule.val_dataset.plot(sample)
+                fig = datamodule.plot(sample)
                 summary_writer = self.logger.experiment
                 summary_writer.add_figure(
                     f"image/{batch_idx}", fig, global_step=self.global_step
@@ -349,7 +349,7 @@ def validation_step(  # type: ignore[override]
             try:
                 datamodule = self.trainer.datamodule  # type: ignore[attr-defined]
                 sample = unbind_samples(batch)[0]
-                fig = datamodule.val_dataset.plot(sample)
+                fig = datamodule.plot(sample)
                 summary_writer = self.logger.experiment
                 summary_writer.add_figure(
                     f"image/{batch_idx}", fig, global_step=self.global_step

torchgeo/trainers/regression.py

@@ -113,7 +113,7 @@ def validation_step(  # type: ignore[override]
             try:
                 datamodule = self.trainer.datamodule  # type: ignore[attr-defined]
                 sample = unbind_samples(batch)[0]
-                fig = datamodule.val_dataset.plot(sample)
+                fig = datamodule.plot(sample)
                 summary_writer = self.logger.experiment
                 summary_writer.add_figure(
                     f"image/{batch_idx}", fig, global_step=self.global_step

torchgeo/trainers/segmentation.py

@@ -178,7 +178,7 @@ def validation_step(  # type: ignore[override]
             try:
                 datamodule = self.trainer.datamodule  # type: ignore[attr-defined]
                 sample = unbind_samples(batch)[0]
-                fig = datamodule.val_dataset.plot(sample)
+                fig = datamodule.plot(sample)
                 summary_writer = self.logger.experiment
                 summary_writer.add_figure(
                     f"image/{batch_idx}", fig, global_step=self.global_step