Skip to content

Commit

Permalink
Added 3 Bands indices AppendRBNDVI , AppendGBNDVI and AppendGRNDVI (#450
Browse files Browse the repository at this point in the history
)

* 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 <ajstewart426@gmail.com>
  • Loading branch information
MATRIX4284 and adamjstewart authored Jun 27, 2022
1 parent 0783560 commit d75d3b8
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 3 deletions.
13 changes: 13 additions & 0 deletions tests/transforms/test_indices.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,17 @@

from torchgeo.transforms import (
AppendBNDVI,
AppendGBNDVI,
AppendGNDVI,
AppendGRNDVI,
AppendNBR,
AppendNDBI,
AppendNDRE,
AppendNDSI,
AppendNDVI,
AppendNDWI,
AppendNormalizedDifferenceIndex,
AppendRBNDVI,
AppendSWI,
AppendTriBandNormalizedDifferenceIndex,
)
Expand Down Expand Up @@ -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)
10 changes: 8 additions & 2 deletions torchgeo/transforms/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
86 changes: 85 additions & 1 deletion torchgeo/transforms/indices.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""
Expand Down Expand Up @@ -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)

0 comments on commit d75d3b8

Please sign in to comment.