Set keepdim=True

microsoft · Jul 14, 2024 · 5fb90fb · 5fb90fb
1 parent 36442fd
commit 5fb90fb
Show file tree

Hide file tree

Showing 27 changed files with 131 additions and 29 deletions.
diff --git a/torchgeo/datamodules/agrifieldnet.py b/torchgeo/datamodules/agrifieldnet.py
@@ -54,10 +54,13 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/chesapeake.py b/torchgeo/datamodules/chesapeake.py
@@ -122,8 +122,10 @@ def __init__(
             self.layers = ['naip-new', 'lc']
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.

diff --git a/torchgeo/datamodules/deepglobelandcover.py b/torchgeo/datamodules/deepglobelandcover.py
@@ -49,7 +49,10 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/fire_risk.py b/torchgeo/datamodules/fire_risk.py
@@ -38,7 +38,10 @@ def __init__(
             K.RandomErasing(p=0.1),
             K.ColorJitter(p=0.5, brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
             data_keys=None,
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/geo.py b/torchgeo/datamodules/geo.py
@@ -70,8 +70,10 @@ def __init__(
         # Data augmentation
         Transform = Callable[[dict[str, Tensor]], dict[str, Tensor]]
         self.aug: Transform = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
         self.train_aug: Transform | None = None
         self.val_aug: Transform | None = None
         self.test_aug: Transform | None = None

diff --git a/torchgeo/datamodules/gid15.py b/torchgeo/datamodules/gid15.py
@@ -51,13 +51,19 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
         self.predict_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.predict_aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets.
 

diff --git a/torchgeo/datamodules/inria.py b/torchgeo/datamodules/inria.py
@@ -49,18 +49,26 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
         self.aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
         self.predict_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+        self.predict_aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets.
 

diff --git a/torchgeo/datamodules/l7irish.py b/torchgeo/datamodules/l7irish.py
@@ -54,10 +54,13 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/l8biome.py b/torchgeo/datamodules/l8biome.py
@@ -54,10 +54,13 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/landcoverai.py b/torchgeo/datamodules/landcoverai.py
@@ -38,7 +38,12 @@ def __init__(
             K.RandomSharpness(p=0.5),
             K.ColorJitter(p=0.5, brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
             data_keys=None,
+            keepdim=True,
         )
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
+
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
diff --git a/torchgeo/datamodules/levircd.py b/torchgeo/datamodules/levircd.py
@@ -46,14 +46,20 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
         self.val_aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
         self.test_aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.val_aug.keepdim = True  # type: ignore[attr-defined]
+        self.test_aug.keepdim = True  # type: ignore[attr-defined]
+
 
 class LEVIRCDPlusDataModule(NonGeoDataModule):
     """LightningDataModule implementation for the LEVIR-CD+ dataset.
@@ -92,14 +98,20 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
         self.val_aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
         self.test_aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.val_aug.keepdim = True  # type: ignore[attr-defined]
+        self.test_aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets.
 

diff --git a/torchgeo/datamodules/naip.py b/torchgeo/datamodules/naip.py
@@ -57,8 +57,10 @@ def __init__(
         )
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.

diff --git a/torchgeo/datamodules/oscd.py b/torchgeo/datamodules/oscd.py
@@ -88,7 +88,10 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/potsdam.py b/torchgeo/datamodules/potsdam.py
@@ -51,7 +51,10 @@ def __init__(
             K.Normalize(mean=self.mean, std=self.std),
             _RandomNCrop(self.patch_size, batch_size),
             data_keys=None,
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/quakeset.py b/torchgeo/datamodules/quakeset.py
@@ -39,4 +39,7 @@ def __init__(
             K.RandomHorizontalFlip(p=0.5),
             K.RandomVerticalFlip(p=0.5),
             data_keys=['image'],
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
diff --git a/torchgeo/datamodules/resisc45.py b/torchgeo/datamodules/resisc45.py
@@ -44,4 +44,7 @@ def __init__(
             K.RandomErasing(p=0.1),
             K.ColorJitter(p=0.5, brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
             data_keys=None,
+            keepdim=True,
         )
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
diff --git a/torchgeo/datamodules/seco.py b/torchgeo/datamodules/seco.py
@@ -53,7 +53,10 @@ def __init__(
                 K.Normalize(mean=torch.tensor(0), std=1 / torch.tensor(255)),
                 K.Normalize(mean=_mean, std=_std),
                 data_keys=None,
+                keepdim=True,
             )
+            # https://github.com/kornia/kornia/issues/2848
+            self.aug.keepdim = True  # type: ignore[attr-defined]
 
     def setup(self, stage: str) -> None:
         """Set up datasets.

diff --git a/torchgeo/datamodules/sentinel2_cdl.py b/torchgeo/datamodules/sentinel2_cdl.py
@@ -68,15 +68,20 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.
 

diff --git a/torchgeo/datamodules/sentinel2_eurocrops.py b/torchgeo/datamodules/sentinel2_eurocrops.py
@@ -13,7 +13,6 @@
 from ..datasets import EuroCrops, Sentinel2, random_grid_cell_assignment
 from ..samplers import GridGeoSampler, RandomBatchGeoSampler
 from ..samplers.utils import _to_tuple
-from ..transforms import AugmentationSequential
 from .geo import GeoDataModule
 
 
@@ -64,21 +63,26 @@ def __init__(
             **self.eurocrops_kwargs,
         )
 
-        self.train_aug = AugmentationSequential(
+        self.train_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             K.RandomResizedCrop(_to_tuple(self.patch_size), scale=(0.6, 1.0)),
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
 
-        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
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.
 

diff --git a/torchgeo/datamodules/sentinel2_nccm.py b/torchgeo/datamodules/sentinel2_nccm.py
@@ -68,15 +68,20 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.
 

diff --git a/torchgeo/datamodules/sentinel2_south_america_soybean.py b/torchgeo/datamodules/sentinel2_south_america_soybean.py
@@ -67,15 +67,20 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std), data_keys=None
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets and samplers.
 

diff --git a/torchgeo/datamodules/southafricacroptype.py b/torchgeo/datamodules/southafricacroptype.py
@@ -12,7 +12,6 @@
 from ..datasets import SouthAfricaCropType, random_bbox_assignment
 from ..samplers import GridGeoSampler, RandomBatchGeoSampler
 from ..samplers.utils import _to_tuple
-from ..transforms import AugmentationSequential
 from .geo import GeoDataModule
 
 
@@ -49,21 +48,26 @@ def __init__(
             **kwargs,
         )
 
-        self.train_aug = AugmentationSequential(
+        self.train_aug = K.AugmentationSequential(
             K.Normalize(mean=self.mean, std=self.std),
             K.RandomResizedCrop(_to_tuple(self.patch_size), scale=(0.6, 1.0)),
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
-            data_keys=['image', 'mask'],
+            data_keys=None,
+            keepdim=True,
             extra_args={
                 DataKey.MASK: {'resample': Resample.NEAREST, 'align_corners': None}
             },
         )
 
-        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
         )
 
+        # https://github.com/kornia/kornia/issues/2848
+        self.train_aug.keepdim = True  # type: ignore[attr-defined]
+        self.aug.keepdim = True  # type: ignore[attr-defined]
+
     def setup(self, stage: str) -> None:
         """Set up datasets.