Remove AugPipe

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_boxes'] = x['bbox_xyxy'].clone()
         dataset.plot(x)
         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['class'], 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['class']
+                x['prediction_boxes'] = 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()

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

torchgeo/datamodules/utils.py

@@ -4,75 +4,19 @@
 """Common datamodule utilities."""
 
 import math
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from typing import Any
 
 import numpy as np
 import torch
-from einops import rearrange
 from torch import Tensor
-from torch.nn import Module
 
 
 # Based on lightning_lite.utilities.exceptions
 class MisconfigurationException(Exception):
     """Exception used to inform users of misuse with Lightning."""
 
 
-class AugPipe(Module):
-    """Pipeline for applying augmentations sequentially on select data keys.
-
-    .. versionadded:: 0.6
-    """
-
-    def __init__(
-        self, augs: Callable[[dict[str, Any]], dict[str, Any]], batch_size: int
-    ) -> None:
-        """Initialize a new AugPipe instance.
-
-        Args:
-            augs: Augmentations to apply.
-            batch_size: Batch size
-        """
-        super().__init__()
-        self.augs = augs
-        self.batch_size = batch_size
-
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        """Apply the augmentation.
-
-        Args:
-            batch: Input batch.
-
-        Returns:
-            Augmented batch.
-        """
-        batch_len = len(batch['image'])
-        for bs in range(batch_len):
-            batch_dict = {
-                'image': batch['image'][bs],
-                'labels': batch['labels'][bs],
-                'boxes': batch['boxes'][bs],
-            }
-
-            if 'masks' in batch:
-                batch_dict['masks'] = batch['masks'][bs]
-
-            batch_dict = self.augs(batch_dict)
-
-            batch['image'][bs] = batch_dict['image']
-            batch['labels'][bs] = batch_dict['labels']
-            batch['boxes'][bs] = batch_dict['boxes']
-
-            if 'masks' in batch:
-                batch['masks'][bs] = batch_dict['masks']
-
-        # Stack images
-        batch['image'] = rearrange(batch['image'], 'b () c h w -> b c h w')
-
-        return batch
-
-
 def collate_fn_detection(batch: list[dict[str, Tensor]]) -> dict[str, Any]:
     """Custom collate fn for object detection and instance segmentation.
 
@@ -85,17 +29,23 @@ def collate_fn_detection(batch: list[dict[str, Tensor]]) -> dict[str, Any]:
     .. versionadded:: 0.6
     """
     output: dict[str, Any] = {}
-    output['image'] = [sample['image'] for sample in batch]
-    output['boxes'] = [sample['boxes'].float() for sample in batch]
-    if 'labels' in batch[0]:
-        output['labels'] = [sample['labels'] for sample in batch]
+    output['image'] = torch.stack([sample['image'] for sample in batch])
+    # Get bbox key as it can be one of {"bbox", "bbox_xyxy", "bbox_xywh"}
+    bbox_key = 'boxes'
+    for key in batch[0].keys():
+        if key in {'bbox', 'bbox_xyxy', 'bbox_xywh'}:
+            bbox_key = key
+
+    output[bbox_key] = [sample[bbox_key].float() for sample in batch]
+    if 'class' in batch[0].keys():
+        output['class'] = [sample['class'] for sample in batch]
     else:
-        output['labels'] = [
-            torch.tensor([1] * len(sample['boxes'])) for sample in batch
+        output['class'] = [
+            torch.tensor([1] * len(sample[bbox_key])) for sample in batch
         ]
 
-    if 'masks' in batch[0]:
-        output['masks'] = [sample['masks'] for sample in batch]
+    if 'mask' in batch[0]:
+        output['mask'] = [sample['mask'] for sample in batch]
     return output
 
 

torchgeo/datamodules/vhr10.py

@@ -11,9 +11,8 @@
 
 from ..datasets import VHR10
 from ..samplers.utils import _to_tuple
-from ..transforms import AugmentationSequential
 from .geo import NonGeoDataModule
-from .utils import AugPipe, collate_fn_detection
+from .utils import collate_fn_detection
 
 
 class VHR10DataModule(NonGeoDataModule):
@@ -52,24 +51,17 @@ def __init__(
 
         self.collate_fn = collate_fn_detection
 
-        self.train_aug = AugPipe(
-            AugmentationSequential(
-                K.Normalize(mean=self.mean, std=self.std),
-                K.Resize(self.patch_size),
-                K.RandomHorizontalFlip(),
-                K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.7),
-                K.RandomVerticalFlip(),
-                data_keys=['image', 'boxes', 'masks'],
-            ),
-            batch_size,
+        self.train_aug = K.AugmentationSequential(
+            K.Normalize(mean=self.mean, std=self.std),
+            K.RandomHorizontalFlip(),
+            K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=0.7),
+            K.RandomVerticalFlip(),
+            data_keys=None,
+            keepdim=True,
         )
-        self.aug = AugPipe(
-            AugmentationSequential(
-                K.Normalize(mean=self.mean, std=self.std),
-                K.Resize(self.patch_size),
-                data_keys=['image', 'boxes', 'masks'],
-            ),
-            batch_size,
+
+        self.aug = K.AugmentationSequential(
+            K.Normalize(mean=self.mean, std=self.std), data_keys=None, keepdim=True
         )
 
     def setup(self, stage: str) -> None:
@@ -78,6 +70,9 @@ def setup(self, stage: str) -> None:
         Args:
             stage: Either 'fit', 'validate', 'test', or 'predict'.
         """
+        self.kwargs['transforms'] = K.AugmentationSequential(
+            K.Resize(self.patch_size), data_keys=None, keepdim=True
+        )
         self.dataset = VHR10(**self.kwargs)
         generator = torch.Generator().manual_seed(0)
         self.train_dataset, self.val_dataset, self.test_dataset = random_split(

torchgeo/datasets/nasa_marine_debris.py

@@ -105,7 +105,7 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
         indices = w_check & h_check
         boxes = boxes[indices]
 
-        sample = {'image': image, 'boxes': boxes}
+        sample = {'image': image, 'bbox_xyxy': boxes}
 
         if self.transforms is not None:
             sample = self.transforms(sample)
@@ -161,8 +161,10 @@ def plot(
 
         sample['image'] = sample['image'].byte()
         image = sample['image']
-        if 'boxes' in sample and len(sample['boxes']):
-            image = draw_bounding_boxes(image=sample['image'], boxes=sample['boxes'])
+        if 'bbox_xyxy' in sample and len(sample['bbox_xyxy']):
+            image = draw_bounding_boxes(
+                image=sample['image'], boxes=sample['bbox_xyxy']
+            )
         image_arr = image.permute((1, 2, 0)).numpy()
 
         if 'prediction_boxes' in sample and len(sample['prediction_boxes']):

torchgeo/datasets/vhr10.py

@@ -248,9 +248,9 @@ def __getitem__(self, index: int) -> dict[str, Any]:
 
         if sample['label']['annotations']:
             sample = self.coco_convert(sample)
-            sample['labels'] = sample['label']['labels']
-            sample['boxes'] = sample['label']['boxes']
-            sample['masks'] = sample['label']['masks']
+            sample['class'] = sample['label']['labels']
+            sample['bbox_xyxy'] = sample['label']['boxes']
+            sample['mask'] = sample['label']['masks'].float()
             del sample['label']
 
         if self.transforms is not None:
@@ -400,11 +400,11 @@ def plot(
         if show_feats != 'boxes':
             skimage = lazy_import('skimage')
 
-        boxes = sample['boxes'].cpu().numpy()
-        labels = sample['labels'].cpu().numpy()
-
-        if 'masks' in sample:
-            masks = [mask.squeeze().cpu().numpy() for mask in sample['masks']]
+        image = sample['image'].permute(1, 2, 0).numpy()
+        boxes = sample['bbox_xyxy'].cpu().numpy()
+        labels = sample['class'].cpu().numpy()
+        if 'mask' in sample:
+            masks = [mask.squeeze().cpu().numpy() for mask in sample['mask']]
 
         n_gt = len(boxes)
 
@@ -459,7 +459,7 @@ def plot(
             )
 
             # Add masks
-            if show_feats in {'masks', 'both'} and 'masks' in sample:
+            if show_feats in {'masks', 'both'} and 'mask' in sample:
                 mask = masks[i]
                 contours = skimage.measure.find_contours(mask, 0.5)
                 for verts in contours:

torchgeo/trainers/detection.py

@@ -238,8 +238,12 @@ def training_step(
         """
         x = batch['image']
         batch_size = x.shape[0]
+        # Get bbox key as it can be one of {"bbox", "bbox_xyxy", "bbox_xywh"}
+        for key in batch.keys():
+            if key in {'bbox', 'bbox_xyxy', 'bbox_xywh'}:
+                bbox_key = key
         y = [
-            {'boxes': batch['boxes'][i], 'labels': batch['labels'][i]}
+            {'boxes': batch[bbox_key][i], 'labels': batch['class'][i]}
             for i in range(batch_size)
         ]
         loss_dict = self(x, y)
@@ -259,8 +263,12 @@ def validation_step(
         """
         x = batch['image']
         batch_size = x.shape[0]
+        # Get bbox key as it can be one of {"bbox", "bbox_xyxy", "bbox_xywh"}
+        for key in batch.keys():
+            if key in {'bbox', 'bbox_xyxy', 'bbox_xywh'}:
+                bbox_key = key
         y = [
-            {'boxes': batch['boxes'][i], 'labels': batch['labels'][i]}
+            {'boxes': batch[bbox_key][i], 'labels': batch['class'][i]}
             for i in range(batch_size)
         ]
         y_hat = self(x)
@@ -313,8 +321,12 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         """
         x = batch['image']
         batch_size = x.shape[0]
+        # Get bbox key as it can be one of {"bbox", "bbox_xyxy", "bbox_xywh"}
+        for key in batch.keys():
+            if key in {'bbox', 'bbox_xyxy', 'bbox_xywh'}:
+                bbox_key = key
         y = [
-            {'boxes': batch['boxes'][i], 'labels': batch['labels'][i]}
+            {'boxes': batch[bbox_key][i], 'labels': batch['class'][i]}
             for i in range(batch_size)
         ]
         y_hat = self(x)