Fix for new datasets

microsoft · Nov 5, 2024 · e34e120 · e34e120
1 parent 9c74986
commit e34e120
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 15 deletions.
diff --git a/torchgeo/datamodules/caffe.py b/torchgeo/datamodules/caffe.py
@@ -9,7 +9,6 @@
 import torch
 
 from ..datasets import CaFFe
-from ..transforms import AugmentationSequential
 from .geo import NonGeoDataModule
 
 
@@ -40,16 +39,18 @@ def __init__(
 
         self.size = size
 
-        self.train_aug = AugmentationSequential(
+        self.train_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             K.Resize(size),
             K.RandomHorizontalFlip(p=0.5),
             K.RandomVerticalFlip(p=0.5),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
         )
 
-        self.aug = AugmentationSequential(
+        self.aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             K.Resize(size),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
         )
diff --git a/torchgeo/datamodules/ftw.py b/torchgeo/datamodules/ftw.py
@@ -9,7 +9,6 @@
 import torch
 
 from ..datasets import FieldsOfTheWorld
-from ..transforms import AugmentationSequential
 from .geo import NonGeoDataModule
 
 
@@ -55,16 +54,17 @@ def __init__(
         self.val_countries = val_countries
         self.test_countries = test_countries
 
-        self.train_aug = AugmentationSequential(
+        self.train_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             K.RandomRotation(p=0.5, degrees=90),
             K.RandomHorizontalFlip(p=0.5),
             K.RandomVerticalFlip(p=0.5),
             K.RandomSharpness(p=0.5),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
         )
-        self.aug = AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=['image', 'mask']
+        self.aug = K.AugmentationSequential(
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
     def setup(self, stage: str) -> None:

diff --git a/torchgeo/datamodules/geonrw.py b/torchgeo/datamodules/geonrw.py
@@ -10,7 +10,6 @@
 from torch.utils.data import Subset
 
 from ..datasets import GeoNRW
-from ..transforms import AugmentationSequential
 from .geo import NonGeoDataModule
 from .utils import group_shuffle_split
 
@@ -38,14 +37,17 @@ def __init__(
         """
         super().__init__(GeoNRW, batch_size, num_workers, **kwargs)
 
-        self.train_aug = AugmentationSequential(
+        self.train_aug = K.AugmentationSequential(
             K.Resize(size),
             K.RandomHorizontalFlip(p=0.5),
             K.RandomVerticalFlip(p=0.5),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
         )
 
-        self.aug = AugmentationSequential(K.Resize(size), data_keys=['image', 'mask'])
+        self.aug = K.AugmentationSequential(
+            K.Resize(size), data_keys=None, keepdim=True
+        )
 
         self.size = size
 

diff --git a/torchgeo/datasets/geonrw.py b/torchgeo/datasets/geonrw.py
@@ -242,7 +242,7 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
 
         # rename to torchgeo standard keys
         sample['image'] = sample.pop('rgb').float()
-        sample['mask'] = sample.pop('seg').long()
+        sample['mask'] = sample.pop('seg').long().squeeze(0)
 
         if self.transforms:
             sample = self.transforms(sample)