Repository: microsoft/torchgeo

Issue:
Add Transforms to compute popular remote sensing indices
Indices are commonly computed for various applications in remote sensing as linear combinations of existing bands in imagery e.g. [Normalized Difference Vegetation Index (NDVI)](https://www.indexdatabase.de/db/i-single.php?id=58) = `(NIR - RED) / (NIR + RED)`. We should look to add some `nn.Modules` under `torchgeo.transforms` for some popular ones.

- [Remote Sensing Index Database](https://www.indexdatabase.de/db/i.php)


Discussion:

Resolving patch:
diff --git a/torchgeo/transforms/__init__.py b/torchgeo/transforms/__init__.py
--- a/torchgeo/transforms/__init__.py
+++ b/torchgeo/transforms/__init__.py
@@ -3,9 +3,24 @@
 
 """TorchGeo transforms."""
 
-from .transforms import Identity, RandomHorizontalFlip, RandomVerticalFlip
+from .indices import AppendNDBI, AppendNDSI, AppendNDVI, AppendNDWI
+from .transforms import (
+    AugmentationSequential,
+    Identity,
+    RandomHorizontalFlip,
+    RandomVerticalFlip,
+)
 
-__all__ = ("Identity", "RandomHorizontalFlip", "RandomVerticalFlip")
+__all__ = (
+    "AppendNDBI",
+    "AppendNDSI",
+    "AppendNDVI",
+    "AppendNDWI",
+    "AugmentationSequential",
+    "Identity",
+    "RandomHorizontalFlip",
+    "RandomVerticalFlip",
+)
 
 # https://stackoverflow.com/questions/40018681
 for module in __all__:
diff --git a/torchgeo/transforms/indices.py b/torchgeo/transforms/indices.py
new file mode 100644
--- /dev/null
+++ b/torchgeo/transforms/indices.py
@@ -0,0 +1,242 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""TorchGeo indices transforms.
+
+For more information about indices see the following references:
+- https://www.indexdatabase.de/db/i.php
+- https://github.com/davemlz/awesome-spectral-indices
+"""
+
+from typing import Dict
+
+import torch
+from torch import Tensor
+from torch.nn import Module  # type: ignore[attr-defined]
+
+# https://github.com/pytorch/pytorch/issues/60979
+# https://github.com/pytorch/pytorch/pull/61045
+Module.__module__ = "torch.nn"
+
+
+_EPSILON = 1e-10
+
+
+def ndbi(swir: Tensor, nir: Tensor) -> Tensor:
+    """Compute Normalized Different Built-up Index (NDBI).
+
+    Args:
+        swir: tensor containing swir band
+        nir: tensor containing nir band
+
+    Returns:
+        tensor containing computed NDBI values
+    """
+    return (swir - nir) / ((swir + nir) + _EPSILON)
+
+
+def ndsi(green: Tensor, swir: Tensor) -> Tensor:
+    """Compute Normalized Different Snow Index (NDSI).
+
+    Args:
+        green: tensor containing green band
+        swir: tensor containing swir band
+
+    Returns:
+        tensor containing computed NDSI values
+    """
+    return (green - swir) / ((green + swir) + _EPSILON)
+
+
+def ndvi(red: Tensor, nir: Tensor) -> Tensor:
+    """Compute Normalized Different Vegetation Index (NDVI).
+
+    Args:
+        red: tensor containing red band
+        nir: tensor containing nir band
+
+    Returns:
+        tensor containing computed NDVI values
+    """
+    return (nir - red) / ((nir + red) + _EPSILON)
+
+
+def ndwi(green: Tensor, nir: Tensor) -> Tensor:
+    """Compute Normalized Different Water Index (NDWI).
+
+    Args:
+        green: tensor containing green band
+        nir: tensor containing nir band
+
+    Returns:
+        tensor containing computed NDWI values
+    """
+    return (green - nir) / ((green + nir) + _EPSILON)
+
+
+class AppendNDBI(Module):  # type: ignore[misc,name-defined]
+    """Normalized Difference Built-up Index (NDBI).
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.1080/01431160304987
+    """
+
+    def __init__(self, index_swir: int, index_nir: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_swir: index of the Short-wave Infrared (SWIR) band in the image
+            index_nir: index of the Near Infrared (NIR) band in the image
+        """
+        super().__init__()
+        self.dim = -3
+        self.index_nir = index_nir
+        self.index_swir = index_swir
+
+    def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Create a band for NDBI and append to image channels.
+
+        Args:
+            sample: a single data sample
+
+        Returns:
+            a sample where the image has an additional channel representing NDBI
+        """
+        if "image" in sample:
+            index = ndbi(
+                swir=sample["image"][:, self.index_swir],
+                nir=sample["image"][:, self.index_nir],
+            )
+            index = index.unsqueeze(self.dim)
+            sample["image"] = torch.cat(  # type: ignore[attr-defined]
+                [sample["image"], index], dim=self.dim
+            )
+
+        return sample
+
+
+class AppendNDSI(Module):  # type: ignore[misc,name-defined]
+    """Normalized Difference Snow Index (NDSI).
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.1109/IGARSS.1994.399618
+    """
+
+    def __init__(self, index_green: int, index_swir: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_green: index of the Green band in the image
+            index_swir: index of the Short-wave Infrared (SWIR) band in the image
+        """
+        super().__init__()
+        self.dim = -3
+        self.index_green = index_green
+        self.index_swir = index_swir
+
+    def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Create a band for NDSI and append to image channels.
+
+        Args:
+            sample: a single data sample
+
+        Returns:
+            a sample where the image has an additional channel representing NDSI
+        """
+        if "image" in sample:
+            index = ndsi(
+                green=sample["image"][:, self.index_green],
+                swir=sample["image"][:, self.index_swir],
+            )
+            index = index.unsqueeze(self.dim)
+            sample["image"] = torch.cat(  # type: ignore[attr-defined]
+                [sample["image"], index], dim=self.dim
+            )
+
+        return sample
+
+
+class AppendNDVI(Module):  # type: ignore[misc,name-defined]
+    """Normalized Difference Vegetation Index (NDVI).
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.1016/0034-4257(79)90013-0
+    """
+
+    def __init__(self, index_red: int, index_nir: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_red: index of the Red band in the image
+            index_nir: index of the Near Infrared (NIR) band in the image
+        """
+        super().__init__()
+        self.dim = -3
+        self.index_red = index_red
+        self.index_nir = index_nir
+
+    def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Create a band for NDVI and append to image channels.
+
+        Args:
+            sample: a single data sample
+
+        Returns:
+            a sample where the image has an additional channel representing NDVI
+        """
+        if "image" in sample:
+            index = ndvi(
+                red=sample["image"][:, self.index_red],
+                nir=sample["image"][:, self.index_nir],
+            )
+            index = index.unsqueeze(self.dim)
+            sample["image"] = torch.cat(  # type: ignore[attr-defined]
+                [sample["image"], index], dim=self.dim
+            )
+
+        return sample
+
+
+class AppendNDWI(Module):  # type: ignore[misc,name-defined]
+    """Normalized Difference Water Index (NDWI).
+
+    If you use this dataset in your research, please cite the following paper:
+
+    * https://doi.org/10.1080/01431169608948714
+    """
+
+    def __init__(self, index_green: int, index_nir: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_green: index of the Green band in the image
+            index_nir: index of the Near Infrared (NIR) band in the image
+        """
+        super().__init__()
+        self.dim = -3
+        self.index_green = index_green
+        self.index_nir = index_nir
+
+    def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Create a band for NDWI and append to image channels.
+
+        Args:
+            sample: a single data sample
+
+        Returns:
+            a sample where the image has an additional channel representing NDWI
+        """
+        if "image" in sample:
+            index = ndwi(
+                green=sample["image"][:, self.index_green],
+                nir=sample["image"][:, self.index_nir],
+            )
+            index = index.unsqueeze(self.dim)
+            sample["image"] = torch.cat(  # type: ignore[attr-defined]
+                [sample["image"], index], dim=self.dim
+            )
+
+        return sample
diff --git a/torchgeo/transforms/transforms.py b/torchgeo/transforms/transforms.py
--- a/torchgeo/transforms/transforms.py
+++ b/torchgeo/transforms/transforms.py
@@ -3,8 +3,9 @@
 
 """TorchGeo transforms."""
 
-from typing import Dict
+from typing import Dict, List, Union
 
+import kornia.augmentation as K
 import torch
 from torch import Tensor
 from torch.nn import Module  # type: ignore[attr-defined]
@@ -97,3 +98,65 @@ def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
             the unchanged input
         """
         return sample
+
+
+class AugmentationSequential(Module):  # type: ignore[misc]
+    """Wrapper around kornia AugmentationSequential to handle input dicts."""
+
+    def __init__(self, *args: Module, data_keys: List[str]) -> None:
+        """Initialize a new augmentation sequential instance.
+
+        Args:
+            *args: Sequence of kornia augmentations
+            data_keys: List of inputs to augment (e.g. ["image", "mask", "boxes"])
+        """
+        super().__init__()
+        self.data_keys = data_keys
+
+        keys = []
+        for key in data_keys:
+            if key == "image":
+                keys.append("input")
+            elif key == "boxes":
+                keys.append("bbox")
+            else:
+                keys.append(key)
+
+        self.augs = K.AugmentationSequential(*args, data_keys=keys)
+
+    def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Perform augmentations and update data dict.
+
+        Args:
+            sample: the input
+
+        Returns:
+            the augmented input
+        """
+        # Kornia augmentations require masks & boxes to be float
+        if "mask" in self.data_keys:
+            mask_dtype = sample["mask"].dtype
+            sample["mask"] = sample["mask"].to(torch.float)  # type:ignore[attr-defined]
+        if "boxes" in self.data_keys:
+            boxes_dtype = sample["boxes"].dtype
+            sample["boxes"] = sample["boxes"].to(
+                torch.float  # type:ignore[attr-defined]
+            )
+
+        inputs = [sample[k] for k in self.data_keys]
+        outputs_list: Union[Tensor, List[Tensor]] = self.augs(*inputs)
+        outputs_list = (
+            outputs_list if isinstance(outputs_list, list) else [outputs_list]
+        )
+        outputs: Dict[str, Tensor] = {
+            k: v for k, v in zip(self.data_keys, outputs_list)
+        }
+        sample.update(outputs)
+
+        # Convert masks & boxes to previous dtype
+        if "mask" in self.data_keys:
+            sample["mask"] = sample["mask"].to(mask_dtype)
+        if "boxes" in self.data_keys:
+            sample["boxes"] = sample["boxes"].to(boxes_dtype)
+
+        return sample


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -19,6 +19,8 @@ jobs:
         python-version: 3.9
     - name: Install pip dependencies
       run: |
+        # TODO: remove after kornia's next release (>0.5.11)
+        pip install git+https://github.com/kornia/kornia@master # required for mypy fixes
         pip install cython numpy  # needed for pycocotools
         pip install .[datasets,tests,train]
     - name: Run mypy checks
diff --git a/tests/transforms/test_indices.py b/tests/transforms/test_indices.py
new file mode 100644
--- /dev/null
+++ b/tests/transforms/test_indices.py
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from typing import Dict
+
+import pytest
+import torch
+from torch import Tensor
+
+from torchgeo.transforms import indices
+
+
+@pytest.fixture
+def sample() -> Dict[str, Tensor]:
+    return {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [[[1, 2, 3], [4, 5, 6], [7, 8, 9]]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[0, 0, 1], [0, 1, 1], [1, 1, 1]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+    }
+
+
+@pytest.fixture
+def batch() -> Dict[str, Tensor]:
+    return {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[0, 0, 1], [0, 1, 1], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+
+
+def test_ndbi(sample: Dict[str, Tensor]) -> None:
+    index = indices.ndbi(swir=sample["image"], nir=sample["image"])
+    assert index.ndim == 3
+    assert index.shape[-2:] == sample["image"].shape[-2:]
+
+
+def test_ndsi(sample: Dict[str, Tensor]) -> None:
+    index = indices.ndsi(green=sample["image"], swir=sample["image"])
+    assert index.ndim == 3
+    assert index.shape[-2:] == sample["image"].shape[-2:]
+
+
+def test_ndvi(sample: Dict[str, Tensor]) -> None:
+    index = indices.ndvi(red=sample["image"], nir=sample["image"])
+    assert index.ndim == 3
+    assert index.shape[-2:] == sample["image"].shape[-2:]
+
+
+def test_ndwi(sample: Dict[str, Tensor]) -> None:
+    index = indices.ndwi(green=sample["image"], nir=sample["image"])
+    assert index.ndim == 3
+    assert index.shape[-2:] == sample["image"].shape[-2:]
+
+
+def test_append_ndbi(batch: Dict[str, Tensor]) -> None:
+    b, c, h, w = batch["image"].shape
+    tr = indices.AppendNDBI(index_swir=0, index_nir=0)
+    output = tr(batch)
+    assert output["image"].shape == (b, c + 1, h, w)
+
+
+def test_append_ndsi(batch: Dict[str, Tensor]) -> None:
+    b, c, h, w = batch["image"].shape
+    tr = indices.AppendNDSI(index_green=0, index_swir=0)
+    output = tr(batch)
+    assert output["image"].shape == (b, c + 1, h, w)
+
+
+def test_append_ndvi(batch: Dict[str, Tensor]) -> None:
+    b, c, h, w = batch["image"].shape
+    tr = indices.AppendNDVI(index_red=0, index_nir=0)
+    output = tr(batch)
+    assert output["image"].shape == (b, c + 1, h, w)
+
+
+def test_append_ndwi(batch: Dict[str, Tensor]) -> None:
+    b, c, h, w = batch["image"].shape
+    tr = indices.AppendNDWI(index_green=0, index_nir=0)
+    output = tr(batch)
+    assert output["image"].shape == (b, c + 1, h, w)
diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py
--- a/tests/transforms/test_transforms.py
+++ b/tests/transforms/test_transforms.py
@@ -3,11 +3,13 @@
 
 from typing import Dict
 
+import kornia.augmentation as K
 import pytest
 import torch
+import torch.nn as nn
 from torch import Tensor
 
-from torchgeo.transforms import transforms
+from torchgeo.transforms import indices, transforms
 
 
 @pytest.fixture
@@ -25,9 +27,82 @@ def sample() -> Dict[str, Tensor]:
     }
 
 
+@pytest.fixture
+def batch_gray() -> Dict[str, Tensor]:
+    return {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[0, 0, 1], [0, 1, 1], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[0, 0, 0, 0], [1, 1, 2, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+
+
+@pytest.fixture
+def batch_rgb() -> Dict[str, Tensor]:
+    return {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                ]
+            ],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[0, 0, 1], [0, 1, 1], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[0, 0, 0, 0], [1, 1, 2, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+
+
+@pytest.fixture
+def batch_multispectral() -> Dict[str, Tensor]:
+    return {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                ]
+            ],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[0, 0, 1], [0, 1, 1], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[0, 0, 0, 0], [1, 1, 2, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+
+
 def assert_matching(output: Dict[str, Tensor], expected: Dict[str, Tensor]) -> None:
     for key in expected:
-        assert torch.allclose(output[key], expected[key])  # type: ignore[attr-defined]
+        err = f"output[{key}] != expected[{key}]"
+        equal = torch.allclose(output[key], expected[key])  # type: ignore[attr-defined]
+        assert equal, err
 
 
 def test_random_horizontal_flip(sample: Dict[str, Tensor]) -> None:
@@ -68,3 +143,169 @@ def test_identity(sample: Dict[str, Tensor]) -> None:
     tr = transforms.Identity()
     output = tr(sample)
     assert_matching(output, sample)
+
+
+def test_augmentation_sequential_gray(batch_gray: Dict[str, Tensor]) -> None:
+    expected = {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [[[[3, 2, 1], [6, 5, 4], [9, 8, 7]]]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 0, 0], [1, 1, 0], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[2, 0, 2, 0], [1, 1, 0, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+    augs = transforms.AugmentationSequential(
+        K.RandomHorizontalFlip(p=1.0),
+        data_keys=["image", "mask", "boxes"],
+    )
+    output = augs(batch_gray)
+    assert_matching(output, expected)
+
+
+def test_augmentation_sequential_rgb(batch_rgb: Dict[str, Tensor]) -> None:
+    expected = {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[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,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 0, 0], [1, 1, 0], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[2, 0, 2, 0], [1, 1, 0, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+    augs = transforms.AugmentationSequential(
+        K.RandomHorizontalFlip(p=1.0),
+        data_keys=["image", "mask", "boxes"],
+    )
+    output = augs(batch_rgb)
+    assert_matching(output, expected)
+
+
+def test_augmentation_sequential_multispectral(
+    batch_multispectral: Dict[str, Tensor]
+) -> None:
+    expected = {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[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,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 0, 0], [1, 1, 0], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[2, 0, 2, 0], [1, 1, 0, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+    augs = transforms.AugmentationSequential(
+        K.RandomHorizontalFlip(p=1.0),
+        data_keys=["image", "mask", "boxes"],
+    )
+    output = augs(batch_multispectral)
+    assert_matching(output, expected)
+
+
+def test_augmentation_sequential_image_only(
+    batch_multispectral: Dict[str, Tensor]
+) -> None:
+    expected = {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[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,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[0, 0, 1], [0, 1, 1], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[0, 0, 0, 0], [1, 1, 2, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+    augs = transforms.AugmentationSequential(
+        K.RandomHorizontalFlip(p=1.0),
+        data_keys=["image"],
+    )
+    output = augs(batch_multispectral)
+    assert_matching(output, expected)
+
+
+def test_sequential_transforms_augmentations(
+    batch_multispectral: Dict[str, Tensor]
+) -> None:
+    expected = {
+        "image": torch.tensor(  # type: ignore[attr-defined]
+            [
+                [
+                    [[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]],
+                    [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                    [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                    [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                    [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                ]
+            ],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "mask": torch.tensor(  # type: ignore[attr-defined]
+            [[[[1, 0, 0], [1, 1, 0], [1, 1, 1]]]],
+            dtype=torch.long,  # type: ignore[attr-defined]
+        ),
+        "boxes": torch.tensor(  # type: ignore[attr-defined]
+            [[2, 0, 2, 0], [1, 1, 0, 2]],
+            dtype=torch.float,  # type: ignore[attr-defined]
+        ),
+        "labels": torch.tensor([[0, 1]]),  # type: ignore[attr-defined]
+    }
+    train_transforms = nn.Sequential(  # type: ignore[attr-defined]
+        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),
+        transforms.AugmentationSequential(
+            K.RandomHorizontalFlip(p=1.0),
+            data_keys=["image", "mask", "boxes"],
+        ),
+    )
+    output = train_transforms(batch_multispectral)
+    assert_matching(output, expected)