Add LoveDADataModule to the trainer tests

tests/conf/loveda.yaml

@@ -0,0 +1,19 @@
+experiment:
+  task: "loveda"
+  module:
+    loss: "ce"
+    model: "unet"
+    backbone: "resnet18"
+    weights: null
+    learning_rate: 1e-3
+    learning_rate_schedule_patience: 6
+    verbose: false
+    in_channels: 3
+    num_classes: 8
+    num_filters: 1
+    ignore_index: null
+  datamodule:
+    root: "tests/data/loveda"
+    download: true
+    batch_size: 1
+    num_workers: 0

tests/trainers/test_segmentation.py

@@ -17,6 +17,7 @@
     ETCI2021DataModule,
     InriaAerialImageLabelingDataModule,
     LandCoverAIDataModule,
+    LoveDADataModule,
     NAIPChesapeakeDataModule,
     OSCDDataModule,
     SEN12MSDataModule,
@@ -44,6 +45,7 @@ class TestSemanticSegmentationTask:
             ("inria_val", InriaAerialImageLabelingDataModule),
             ("inria_test", InriaAerialImageLabelingDataModule),
             ("landcoverai", LandCoverAIDataModule),
+            ("loveda", LoveDADataModule),
             ("naipchesapeake", NAIPChesapeakeDataModule),
             ("oscd_all", OSCDDataModule),
             ("oscd_rgb", OSCDDataModule),
@@ -81,7 +83,9 @@ def test_trainer(
         # Instantiate trainer
         trainer = Trainer(fast_dev_run=True, log_every_n_steps=1, max_epochs=1)
         trainer.fit(model=model, datamodule=datamodule)
-        trainer.test(model=model, datamodule=datamodule)
+
+        if hasattr(datamodule, "test_dataset") or hasattr(datamodule, "test_sampler"):
+            trainer.test(model=model, datamodule=datamodule)
 
         if hasattr(datamodule, "predict_dataset"):
             trainer.predict(model=model, datamodule=datamodule)

torchgeo/datamodules/loveda.py

@@ -67,18 +67,19 @@ def setup(self, stage: Optional[str] = None) -> None:
             stage: stage to set up
         """
         train_transforms = self.preprocess
-        val_test_transforms = self.preprocess
+        val_predict_transforms = self.preprocess
 
         self.train_dataset = LoveDA(
             split="train", transforms=train_transforms, **self.kwargs
         )
 
         self.val_dataset = LoveDA(
-            split="val", transforms=val_test_transforms, **self.kwargs
+            split="val", transforms=val_predict_transforms, **self.kwargs
         )
 
-        self.test_dataset = LoveDA(
-            split="test", transforms=val_test_transforms, **self.kwargs
+        # Test set masks are not public, use for prediction instead
+        self.predict_dataset = LoveDA(
+            split="test", transforms=val_predict_transforms, **self.kwargs
         )
 
     def train_dataloader(self) -> DataLoader[Any]:
@@ -107,14 +108,14 @@ def val_dataloader(self) -> DataLoader[Any]:
             shuffle=False,
         )
 
-    def test_dataloader(self) -> DataLoader[Any]:
-        """Return a DataLoader for testing.
+    def predict_dataloader(self) -> DataLoader[Any]:
+        """Return a DataLoader for prediction.
 
         Returns:
-            testing data loader
+            predict data loader
         """
         return DataLoader(
-            self.test_dataset,
+            self.predict_dataset,
             batch_size=self.batch_size,
             num_workers=self.num_workers,
             shuffle=False,

train.py

@@ -21,6 +21,7 @@
     EuroSATDataModule,
     InriaAerialImageLabelingDataModule,
     LandCoverAIDataModule,
+    LoveDADataModule,
     NAIPChesapeakeDataModule,
     NASAMarineDebrisDataModule,
     OSCDDataModule,
@@ -52,6 +53,7 @@
     "etci2021": (SemanticSegmentationTask, ETCI2021DataModule),
     "inria": (SemanticSegmentationTask, InriaAerialImageLabelingDataModule),
     "landcoverai": (SemanticSegmentationTask, LandCoverAIDataModule),
+    "loveda": (SemanticSegmentationTask, LoveDADataModule),
     "naipchesapeake": (SemanticSegmentationTask, NAIPChesapeakeDataModule),
     "nasa_marine_debris": (ObjectDetectionTask, NASAMarineDebrisDataModule),
     "oscd": (SemanticSegmentationTask, OSCDDataModule),