Remove AugPipe

tests/datamodules/test_fair1m.py

@@ -35,7 +35,7 @@ def test_plot(self, datamodule: FAIR1MDataModule) -> None:
         batch = next(iter(datamodule.val_dataloader()))
         sample = {
             'image': batch['image'][0],
-            'boxes': batch['boxes'][0],
+            'bbox_xyxy': batch['bbox_xyxy'][0],
             'label': batch['label'][0],
         }
         datamodule.plot(sample)

tests/datasets/test_fair1m.py

@@ -70,9 +70,9 @@ def test_getitem(self, dataset: FAIR1M) -> None:
         assert x['image'].shape[0] == 3
 
         if dataset.split != 'test':
-            assert isinstance(x['boxes'], torch.Tensor)
+            assert isinstance(x['bbox_xyxy'], torch.Tensor)
             assert isinstance(x['label'], torch.Tensor)
-            assert x['boxes'].shape[-2:] == (5, 2)
+            assert x['bbox_xyxy'].shape[-2:] == (5, 2)
             assert x['label'].ndim == 1
 
     def test_len(self, dataset: FAIR1M) -> None:
@@ -124,6 +124,6 @@ def test_plot(self, dataset: FAIR1M) -> None:
         plt.close()
 
         if dataset.split != 'test':
-            x['prediction_boxes'] = x['boxes'].clone()
+            x['prediction_bbox_xyxy'] = x['bbox_xyxy'].clone()
             dataset.plot(x)
             plt.close()

tests/datasets/test_forestdamage.py

@@ -36,7 +36,7 @@ def test_getitem(self, dataset: ForestDamage) -> None:
         assert isinstance(x, dict)
         assert isinstance(x['image'], torch.Tensor)
         assert isinstance(x['label'], torch.Tensor)
-        assert isinstance(x['boxes'], torch.Tensor)
+        assert isinstance(x['bbox_xyxy'], torch.Tensor)
         assert x['image'].shape[0] == 3
         assert x['image'].ndim == 3
 
@@ -67,6 +67,6 @@ def test_plot(self, dataset: ForestDamage) -> None:
 
     def test_plot_prediction(self, dataset: ForestDamage) -> None:
         x = dataset[0].copy()
-        x['prediction_boxes'] = x['boxes'].clone()
+        x['prediction_bbox_xyxy'] = x['bbox_xyxy'].clone()
         dataset.plot(x, suptitle='Prediction')
         plt.close()

tests/datasets/test_idtrees.py

@@ -57,11 +57,11 @@ def test_getitem(self, dataset: IDTReeS) -> None:
 
         if 'label' in x:
             assert isinstance(x['label'], torch.Tensor)
-        if 'boxes' in x:
-            assert isinstance(x['boxes'], torch.Tensor)
-            if x['boxes'].ndim != 1:
-                assert x['boxes'].ndim == 2
-                assert x['boxes'].shape[-1] == 4
+        if 'bbox_xyxy' in x:
+            assert isinstance(x['bbox_xyxy'], torch.Tensor)
+            if x['bbox_xyxy'].ndim != 1:
+                assert x['bbox_xyxy'].ndim == 2
+                assert x['bbox_xyxy'].shape[-1] == 4
 
     def test_len(self, dataset: IDTReeS) -> None:
         assert len(dataset) == 3
@@ -87,8 +87,8 @@ def test_plot(self, dataset: IDTReeS) -> None:
         dataset.plot(x, show_titles=False)
         plt.close()
 
-        if 'boxes' in x:
-            x['prediction_boxes'] = x['boxes']
+        if 'bbox_xyxy' in x:
+            x['prediction_bbox_xyxy'] = x['bbox_xyxy']
             dataset.plot(x, show_titles=True)
             plt.close()
         if 'label' in x:

tests/datasets/test_nasa_marine_debris.py

@@ -28,9 +28,9 @@ def test_getitem(self, dataset: NASAMarineDebris) -> None:
         x = dataset[0]
         assert isinstance(x, dict)
         assert isinstance(x['image'], torch.Tensor)
-        assert isinstance(x['boxes'], torch.Tensor)
+        assert isinstance(x['bbox_xyxy'], torch.Tensor)
         assert x['image'].shape[0] == 3
-        assert x['boxes'].shape[-1] == 4
+        assert x['bbox_xyxy'].shape[-1] == 4
 
     def test_len(self, dataset: NASAMarineDebris) -> None:
         assert len(dataset) == 5
@@ -50,6 +50,6 @@ def test_plot(self, dataset: NASAMarineDebris) -> None:
         plt.close()
         dataset.plot(x, show_titles=False)
         plt.close()
-        x['prediction_boxes'] = x['boxes'].clone()
+        x['prediction_bbox_xyxy'] = x['bbox_xyxy'].clone()
         dataset.plot(x)
         plt.close()

tests/datasets/test_pastis.py

@@ -52,7 +52,7 @@ def test_getitem_instance(self, dataset: PASTIS) -> None:
         assert isinstance(x, dict)
         assert isinstance(x['image'], torch.Tensor)
         assert isinstance(x['mask'], torch.Tensor)
-        assert isinstance(x['boxes'], torch.Tensor)
+        assert isinstance(x['bbox_xyxy'], torch.Tensor)
         assert isinstance(x['label'], torch.Tensor)
 
     def test_len(self, dataset: PASTIS) -> None:

tests/datasets/test_reforestree.py

@@ -36,11 +36,11 @@ def test_getitem(self, dataset: ReforesTree) -> None:
         assert isinstance(x, dict)
         assert isinstance(x['image'], torch.Tensor)
         assert isinstance(x['label'], torch.Tensor)
-        assert isinstance(x['boxes'], torch.Tensor)
+        assert isinstance(x['bbox_xyxy'], torch.Tensor)
         assert isinstance(x['agb'], torch.Tensor)
         assert x['image'].shape[0] == 3
         assert x['image'].ndim == 3
-        assert len(x['boxes']) == 2
+        assert len(x['bbox_xyxy']) == 2
 
     def test_len(self, dataset: ReforesTree) -> None:
         assert len(dataset) == 2
@@ -67,6 +67,6 @@ def test_plot(self, dataset: ReforesTree) -> None:
 
     def test_plot_prediction(self, dataset: ReforesTree) -> None:
         x = dataset[0].copy()
-        x['prediction_boxes'] = x['boxes'].clone()
+        x['prediction_bbox_xyxy'] = x['bbox_xyxy'].clone()
         dataset.plot(x, suptitle='Prediction')
         plt.close()

tests/datasets/test_vhr10.py

@@ -41,10 +41,10 @@ def test_getitem(self, dataset: VHR10) -> None:
             assert isinstance(x, dict)
             assert isinstance(x['image'], torch.Tensor)
             if dataset.split == 'positive':
-                assert isinstance(x['labels'], torch.Tensor)
-                assert isinstance(x['boxes'], torch.Tensor)
-                if 'masks' in x:
-                    assert isinstance(x['masks'], torch.Tensor)
+                assert isinstance(x['label'], torch.Tensor)
+                assert isinstance(x['bbox_xyxy'], torch.Tensor)
+                if 'mask' in x:
+                    assert isinstance(x['mask'], torch.Tensor)
 
     def test_len(self, dataset: VHR10) -> None:
         if dataset.split == 'positive':
@@ -82,10 +82,10 @@ def test_plot(self, dataset: VHR10) -> None:
             scores = [0.7, 0.3, 0.7]
             for i in range(3):
                 x = dataset[i]
-                x['prediction_labels'] = x['labels']
-                x['prediction_boxes'] = x['boxes']
+                x['prediction_labels'] = x['label']
+                x['prediction_bbox_xyxy'] = x['bbox_xyxy']
                 x['prediction_scores'] = torch.Tensor([scores[i]])
-                if 'masks' in x:
-                    x['prediction_masks'] = x['masks']
+                if 'mask' in x:
+                    x['prediction_masks'] = x['mask']
                     dataset.plot(x, show_feats='masks')
                     plt.close()

tests/datasets/test_zuericrop.py

@@ -35,7 +35,7 @@ def test_getitem(self, dataset: ZueriCrop) -> None:
         assert isinstance(x, dict)
         assert isinstance(x['image'], torch.Tensor)
         assert isinstance(x['mask'], torch.Tensor)
-        assert isinstance(x['boxes'], torch.Tensor)
+        assert isinstance(x['bbox_xyxy'], torch.Tensor)
         assert isinstance(x['label'], torch.Tensor)
 
         # Image tests
@@ -46,8 +46,8 @@ def test_getitem(self, dataset: ZueriCrop) -> None:
         assert x['mask'].shape[-2:] == x['image'].shape[-2:]
 
         # Bboxes tests
-        assert x['boxes'].ndim == 2
-        assert x['boxes'].shape[1] == 4
+        assert x['bbox_xyxy'].ndim == 2
+        assert x['bbox_xyxy'].shape[1] == 4
 
         # Labels tests
         assert x['label'].ndim == 1

tests/transforms/test_transforms.py

@@ -23,7 +23,7 @@ def batch_gray() -> dict[str, Tensor]:
     return {
         'image': torch.tensor([[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]], dtype=torch.float),
         'mask': torch.tensor([[[0, 0, 1], [0, 1, 1], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
 
@@ -42,7 +42,7 @@ def batch_rgb() -> dict[str, Tensor]:
             dtype=torch.float,
         ),
         'mask': torch.tensor([[[0, 0, 1], [0, 1, 1], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[0.0, 1.0, 1.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
 
@@ -63,7 +63,7 @@ def batch_multispectral() -> dict[str, Tensor]:
             dtype=torch.float,
         ),
         'mask': torch.tensor([[[0, 0, 1], [0, 1, 1], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[0.0, 1.0, 1.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
 
@@ -79,12 +79,10 @@ def test_augmentation_sequential_gray(batch_gray: dict[str, Tensor]) -> None:
     expected = {
         'image': torch.tensor([[[[3, 2, 1], [6, 5, 4], [9, 8, 7]]]], dtype=torch.float),
         'mask': torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[1.0, 0.0, 3.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
-    augs = transforms.AugmentationSequential(
-        K.RandomHorizontalFlip(p=1.0), data_keys=['image', 'mask', 'boxes']
-    )
+    augs = K.AugmentationSequential(K.RandomHorizontalFlip(p=1.0), data_keys=None)
     output = augs(batch_gray)
     assert_matching(output, expected)
 
@@ -102,12 +100,10 @@ def test_augmentation_sequential_rgb(batch_rgb: dict[str, Tensor]) -> None:
             dtype=torch.float,
         ),
         'mask': torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[1.0, 0.0, 3.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[1.0, 1.0, 2.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
-    augs = transforms.AugmentationSequential(
-        K.RandomHorizontalFlip(p=1.0), data_keys=['image', 'mask', 'boxes']
-    )
+    augs = K.AugmentationSequential(K.RandomHorizontalFlip(p=1.0), data_keys=None)
     output = augs(batch_rgb)
     assert_matching(output, expected)
 
@@ -119,51 +115,43 @@ def test_augmentation_sequential_multispectral(
         'image': torch.tensor(
             [
                 [
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+                    [[7, 8, 9], [4, 5, 6], [1, 2, 3]],
+                    [[7, 8, 9], [4, 5, 6], [1, 2, 3]],
+                    [[7, 8, 9], [4, 5, 6], [1, 2, 3]],
+                    [[7, 8, 9], [4, 5, 6], [1, 2, 3]],
+                    [[7, 8, 9], [4, 5, 6], [1, 2, 3]],
                 ]
             ],
             dtype=torch.float,
         ),
-        'mask': torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[1.0, 0.0, 3.0, 2.0]], dtype=torch.float),
+        'mask': torch.tensor([[[1, 1, 1], [0, 1, 1], [0, 0, 1]]], dtype=torch.long),
+        'bbox_xyxy': torch.tensor([[0.0, 0.0, 1.0, 1.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
-    augs = transforms.AugmentationSequential(
-        K.RandomHorizontalFlip(p=1.0), data_keys=['image', 'mask', 'boxes']
-    )
+    augs = K.AugmentationSequential(K.RandomVerticalFlip(p=1.0), data_keys=None)
     output = augs(batch_multispectral)
     assert_matching(output, expected)
 
 
 def test_augmentation_sequential_image_only(
     batch_multispectral: dict[str, Tensor],
 ) -> None:
-    expected = {
-        'image': torch.tensor(
+    expected_image = torch.tensor(
+        [
             [
-                [
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                    [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
-                ]
-            ],
-            dtype=torch.float,
-        ),
-        'mask': torch.tensor([[[0, 0, 1], [0, 1, 1], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[0.0, 0.0, 2.0, 2.0]], dtype=torch.float),
-        'labels': torch.tensor([[0, 1]]),
-    }
-    augs = transforms.AugmentationSequential(
-        K.RandomHorizontalFlip(p=1.0), data_keys=['image']
+                [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+                [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+                [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+                [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+                [[3, 2, 1], [6, 5, 4], [9, 8, 7]],
+            ]
+        ],
+        dtype=torch.float,
     )
-    output = augs(batch_multispectral)
-    assert_matching(output, expected)
+
+    augs = K.AugmentationSequential(K.RandomHorizontalFlip(p=1.0), data_keys=['image'])
+    aug_image = augs(batch_multispectral['image'])
+    assert torch.allclose(aug_image, expected_image)
 
 
 def test_sequential_transforms_augmentations(
@@ -188,17 +176,17 @@ def test_sequential_transforms_augmentations(
             dtype=torch.float,
         ),
         'mask': torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.long),
-        'boxes': torch.tensor([[1.0, 0.0, 3.0, 2.0]], dtype=torch.float),
+        'bbox_xyxy': torch.tensor([[1.0, 1.0, 2.0, 2.0]], dtype=torch.float),
         'labels': torch.tensor([[0, 1]]),
     }
-    train_transforms = transforms.AugmentationSequential(
+    train_transforms = K.AugmentationSequential(
         indices.AppendNBR(index_nir=0, index_swir=0),
         indices.AppendNDBI(index_swir=0, index_nir=0),
         indices.AppendNDSI(index_green=0, index_swir=0),
         indices.AppendNDVI(index_red=0, index_nir=0),
         indices.AppendNDWI(index_green=0, index_nir=0),
         K.RandomHorizontalFlip(p=1.0),
-        data_keys=['image', 'mask', 'boxes'],
+        data_keys=None,
     )
     output = train_transforms(batch_multispectral)
     assert_matching(output, expected)

torchgeo/datamodules/fair1m.py

@@ -26,8 +26,8 @@ def collate_fn(batch: list[dict[str, Tensor]]) -> dict[str, Any]:
     output: dict[str, Any] = {}
     output['image'] = torch.stack([sample['image'] for sample in batch])
 
-    if 'boxes' in batch[0]:
-        output['boxes'] = [sample['boxes'] for sample in batch]
+    if 'bbox_xyxy' in batch[0]:
+        output['bbox_xyxy'] = [sample['bbox_xyxy'] for sample in batch]
     if 'label' in batch[0]:
         output['label'] = [sample['label'] for sample in batch]
 

torchgeo/datamodules/nasa_marine_debris.py

@@ -10,9 +10,8 @@
 from torch.utils.data import random_split
 
 from ..datasets import NASAMarineDebris
-from ..transforms import AugmentationSequential
 from .geo import NonGeoDataModule
-from .utils import AugPipe, collate_fn_detection
+from .utils import collate_fn_detection
 
 
 class NASAMarineDebrisDataModule(NonGeoDataModule):
@@ -46,11 +45,8 @@ def __init__(
         self.val_split_pct = val_split_pct
         self.test_split_pct = test_split_pct
 
-        self.aug = AugPipe(
-            AugmentationSequential(
-                K.Normalize(mean=self.mean, std=self.std), data_keys=['image', 'boxes']
-            ),
-            batch_size,
+        self.aug = K.AugmentationSequential(
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
         self.collate_fn = collate_fn_detection