Added GNDVI and SWI Indices (microsoft#371)

* Update indices.py Added MNDWI to support modification of normalised difference water index (NDWI) to enhance open water features in remotely sensed imagery. * Update indices.py * Update indices.py * Update indices.py Added a new Index Soil Water Index * Update indices.py Style error fixed * Update indices.py * SWI addition * Update indices.py * Update indices.py * Update indices.py * Update indices.py Changed the comment to fit in the code accordingly * Update indices.py modified * Update indices.py corrected the citation * Update indices.py * Ubuntu CUDA 11.3 docker for torchgeo * Update Dockerfile_ubuntu_1804_cuda_11_3_cudnn_8_torchgeo * Delete Dockerfile_ubuntu_1804_cuda_11_3_cudnn_8_torchgeo * Added Dockerfile Added Dockerfile * deleted docker due to wrong issue number * Added GNDVI Index Added GNDVI Index * corrected docstring * updated refernce index to doi * Added Triband Normalized Indexes * Update indices.py Saved * Updated GNDVI * Delete transforms.ipynb deleting * Added the notebook changes * Delete transforms.ipynb * Added the notebook changes * Added the spaces * Update indices.py * formatting corrected * Running black * Reset transforms * Clean up docstrings * Adding tests and adding to list of transforms exported by the module * Consistency Co-authored-by: Caleb Robinson <calebrob6@gmail.com>
yichiac · Feb 1, 2022 · a576867 · a576867
1 parent 7c65ce6
commit a576867
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 3 deletions.
diff --git a/tests/transforms/test_indices.py b/tests/transforms/test_indices.py
@@ -7,7 +7,16 @@
 import torch
 from torch import Tensor
 
-from torchgeo.transforms import indices
+from torchgeo.transforms import (
+    AppendGNDVI,
+    AppendNBR,
+    AppendNDBI,
+    AppendNDSI,
+    AppendNDVI,
+    AppendNDWI,
+    AppendNormalizedDifferenceIndex,
+    AppendSWI,
+)
 
 
 @pytest.fixture
@@ -41,13 +50,26 @@ def batch() -> Dict[str, Tensor]:
 
 def test_append_index_sample(sample: Dict[str, Tensor]) -> None:
     c, h, w = sample["image"].shape
-    tr = indices.AppendNormalizedDifferenceIndex(index_a=0, index_b=0)
+    tr = AppendNormalizedDifferenceIndex(index_a=0, index_b=0)
     output = tr(sample)
     assert output["image"].shape == (c + 1, h, w)
 
 
 def test_append_index_batch(batch: Dict[str, Tensor]) -> None:
     b, c, h, w = batch["image"].shape
-    tr = indices.AppendNormalizedDifferenceIndex(index_a=0, index_b=0)
+    tr = AppendNormalizedDifferenceIndex(index_a=0, index_b=0)
     output = tr(batch)
     assert output["image"].shape == (b, c + 1, h, w)
+
+
+@pytest.mark.parametrize(
+    "index",
+    [AppendNBR, AppendNDBI, AppendNDSI, AppendNDVI, AppendNDWI, AppendSWI, AppendGNDVI],
+)
+def test_append_normalized_difference_indices(
+    sample: Dict[str, Tensor], index: AppendNormalizedDifferenceIndex
+) -> None:
+    c, h, w = sample["image"].shape
+    tr = index(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
@@ -4,22 +4,26 @@
 """TorchGeo transforms."""
 
 from .indices import (
+    AppendGNDVI,
     AppendNBR,
     AppendNDBI,
     AppendNDSI,
     AppendNDVI,
     AppendNDWI,
     AppendNormalizedDifferenceIndex,
+    AppendSWI,
 )
 from .transforms import AugmentationSequential
 
 __all__ = (
     "AppendNormalizedDifferenceIndex",
+    "AppendGNDVI",
     "AppendNBR",
     "AppendNDBI",
     "AppendNDSI",
     "AppendNDVI",
     "AppendNDWI",
+    "AppendSWI",
     "AugmentationSequential",
 )
 

diff --git a/torchgeo/transforms/indices.py b/torchgeo/transforms/indices.py
@@ -165,3 +165,39 @@ def __init__(self, index_green: int, index_nir: int) -> None:
             index_nir: index of the Near Infrared (NIR) band in the image
         """
         super().__init__(index_a=index_green, index_b=index_nir)
+
+
+class AppendSWI(AppendNormalizedDifferenceIndex):
+    """Standardized Water-Level Index (SWI).
+
+    If you use this index in your research, please cite the following paper:
+
+    * https://doi.org/10.3390/w13121647
+    """
+
+    def __init__(self, index_red: int, index_swir: int) -> None:
+        """Initialize a new transform instance.
+
+        Args:
+            index_red: index of the VRE1 band, e.g. B5 in Sentinel 2 imagery
+            index_swir: index of the SWIR2 band, e.g. B11 in Sentinel 2 imagery
+        """
+        super().__init__(index_a=index_red, index_b=index_swir)
+
+
+class AppendGNDVI(AppendNormalizedDifferenceIndex):
+    """Green Normalized Difference Vegetation Index (GNDVI).
+
+    If you use this index in your research, please cite the following paper:
+
+    * https://doi.org/10.2134/agronj2001.933583x
+    """
+
+    def __init__(self, index_nir: int, index_green: 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, e.g. B3 in Sentinel 2 imagery
+        """
+        super().__init__(index_a=index_nir, index_b=index_green)