Added 3 Bands indices AppendRBNDVI , AppendGBNDVI and AppendGRNDVI (#450

) * added GBNDVI and GRNDVI * added RBNDVI * Updated for AppendGBNDVI AppendGRNDVI AppendRBNDVI * added AppendRBNDVI AppendGBNDVI AppendGRNDVI * Import and style fixes * Fix tests * Docstring fixes Co-authored-by: Adam J. Stewart <[email protected]>
microsoft · Jun 27, 2022 · d75d3b8 · d75d3b8
1 parent 0783560
commit d75d3b8
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 3 deletions.
diff --git a/tests/transforms/test_indices.py b/tests/transforms/test_indices.py
@@ -9,14 +9,17 @@
 
 from torchgeo.transforms import (
     AppendBNDVI,
+    AppendGBNDVI,
     AppendGNDVI,
+    AppendGRNDVI,
     AppendNBR,
     AppendNDBI,
     AppendNDRE,
     AppendNDSI,
     AppendNDVI,
     AppendNDWI,
     AppendNormalizedDifferenceIndex,
+    AppendRBNDVI,
     AppendSWI,
     AppendTriBandNormalizedDifferenceIndex,
 )
@@ -81,3 +84,13 @@ def test_append_normalized_difference_indices(
     tr = index(0, 0)
     output = tr(sample)
     assert output["image"].shape == (c + 1, h, w)
+
+
+@pytest.mark.parametrize("index", [AppendGBNDVI, AppendGRNDVI, AppendRBNDVI])
+def test_append_tri_band_normalized_difference_indices(
+    sample: Dict[str, Tensor], index: AppendTriBandNormalizedDifferenceIndex
+) -> None:
+    c, h, w = sample["image"].shape
+    tr = index(0, 0, 0)
+    output = tr(sample)
+    assert output["image"].shape == (c + 1, h, w)
diff --git a/torchgeo/transforms/__init__.py b/torchgeo/transforms/__init__.py
@@ -5,32 +5,38 @@
 
 from .indices import (
     AppendBNDVI,
+    AppendGBNDVI,
     AppendGNDVI,
+    AppendGRNDVI,
     AppendNBR,
     AppendNDBI,
     AppendNDRE,
     AppendNDSI,
     AppendNDVI,
     AppendNDWI,
     AppendNormalizedDifferenceIndex,
+    AppendRBNDVI,
     AppendSWI,
     AppendTriBandNormalizedDifferenceIndex,
 )
 from .transforms import AugmentationSequential
 
 __all__ = (
-    "AppendNormalizedDifferenceIndex",
     "AppendBNDVI",
+    "AppendGBNDVI",
     "AppendGNDVI",
+    "AppendGRNDVI",
     "AppendNBR",
     "AppendNDBI",
     "AppendNDRE",
     "AppendNDSI",
     "AppendNDVI",
     "AppendNDWI",
+    "AppendNormalizedDifferenceIndex",
+    "AppendRBNDVI",
     "AppendSWI",
-    "AugmentationSequential",
     "AppendTriBandNormalizedDifferenceIndex",
+    "AugmentationSequential",
 )
 
 # https://stackoverflow.com/questions/40018681

diff --git a/torchgeo/transforms/indices.py b/torchgeo/transforms/indices.py
@@ -308,7 +308,7 @@ class AppendTriBandNormalizedDifferenceIndex(Module):
 
     .. math::
 
-       \text{NDI} = \frac{A - {B + C}}{A + {B + C}}
+       \text{NDI} = \frac{A - (B + C)}{A + (B + C)}
 
     .. versionadded:: 0.3
     """
@@ -360,3 +360,87 @@ def forward(self, sample: Dict[str, Tensor]) -> Dict[str, Tensor]:
             sample["image"] = torch.cat([sample["image"], index], dim=self.dim)
 
         return sample
+
+
+class AppendGRNDVI(AppendTriBandNormalizedDifferenceIndex):
+    r"""Green-Red Normalized Difference Vegetation Index (GRNDVI).
+
+    Computes the following index:
+
+    .. math::
+
+       \text{GRNDVI} =
+           \frac{\text{NIR} - (\text{G} + \text{R})}{\text{NIR} + (\text{G} + \text{R})}
+
+    If you use this index in your research, please cite the following paper:
+
+    * https://doi.org/10.1016/S1672-6308(07)60027-4
+
+    .. versionadded:: 0.3
+    """
+
+    def __init__(self, index_nir: int, index_green: int, index_red: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_nir: index of the NIR band, e.g. B8 in Sentinel 2 imagery
+            index_green: index of the Green band, B3 in Sentinel 2 imagery
+            index_red: index of the Red band, B4 in Sentinel 2 imagery
+        """
+        super().__init__(index_a=index_nir, index_b=index_green, index_c=index_red)
+
+
+class AppendGBNDVI(AppendTriBandNormalizedDifferenceIndex):
+    r"""Green-Blue Normalized Difference Vegetation Index (GBNDVI).
+
+    Computes the following index:
+
+    .. math::
+
+       \text{GBNDVI} =
+           \frac{\text{NIR} - (\text{G} + \text{B})}{\text{NIR} + (\text{G} + \text{B})}
+
+    If you use this index in your research, please cite the following paper:
+
+    * https://doi.org/10.1016/S1672-6308(07)60027-4
+
+    .. versionadded:: 0.3
+    """
+
+    def __init__(self, index_nir: int, index_green: int, index_blue: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_nir: index of the NIR band, e.g. B8 in Sentinel 2 imagery
+            index_green: index of the Green band, B3 in Sentinel 2 imagery
+            index_blue: index of the Blue band, B2 in Sentinel 2 imagery
+        """
+        super().__init__(index_a=index_nir, index_b=index_green, index_c=index_blue)
+
+
+class AppendRBNDVI(AppendTriBandNormalizedDifferenceIndex):
+    r"""Red-Blue Normalized Difference Vegetation Index (RBNDVI).
+
+    Computes the following index:
+
+    .. math::
+
+       \text{RBNDVI} =
+           \frac{\text{NIR} - (\text{R} + \text{B})}{\text{NIR} + (\text{R} + \text{B})}
+
+    If you use this index in your research, please cite the following paper:
+
+    * https://doi.org/10.1016/S1672-6308(07)60027-4
+
+    .. versionadded:: 0.3
+    """
+
+    def __init__(self, index_nir: int, index_red: int, index_blue: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_nir: index of the NIR band, e.g. B8 in Sentinel 2 imagery
+            index_red: index of the Red band, B4 in Sentinel 2 imagery
+            index_blue: index of the Blue band, B2 in Sentinel 2 imagery
+        """
+        super().__init__(index_a=index_nir, index_b=index_red, index_c=index_blue)