feat: Add tool to replace pixel values with nodata

Author: mipeso

eis_toolkit/cli.py

@@ -344,6 +344,16 @@ class WeightsType(str, Enum):
     descending = "descending"
 
 
+class ReplaceCondition(str, Enum):
+    """Replace conditions for replace with nodata."""
+
+    equal = "equal"
+    less_than = "less_than"
+    greater_than = "greater_than"
+    less_than_or_equal = "less_than_or_equal"
+    greater_than_or_equal = "greater_than_or_equal"
+
+
 INPUT_FILE_OPTION = Annotated[
     Path,
     typer.Option(
@@ -4077,6 +4087,31 @@ def convert_raster_nodata_cli(
     typer.echo(f"Converting nodata completed, writing raster to {output_raster}.")
 
 
+@app.command()
+def replace_with_nodata_cli(
+    input_raster: INPUT_FILE_OPTION,
+    output_raster: OUTPUT_FILE_OPTION,
+    target_value: Annotated[float, typer.Option()],
+    nodata_value: float = None,
+    replace_condition: Annotated[ReplaceCondition, typer.Option(case_sensitive=False)] = ReplaceCondition.equal,
+):
+    """Replace raster pixel values with nodata."""
+    from eis_toolkit.utilities.nodata import replace_with_nodata
+
+    typer.echo("Progress: 10%")
+
+    with rasterio.open(input_raster) as raster:
+        typer.echo("Progress: 25%")
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+    typer.echo("Progress: 70%")
+
+    with rasterio.open(output_raster, "w", **out_meta) as dst:
+        dst.write(out_image)
+    typer.echo("Progres: 100%")
+
+    typer.echo(f"Raster pixel values replaced with nodata, writing raster to {output_raster}.")
+
+
 @app.command()
 def set_raster_nodata_cli(
     input_raster: INPUT_FILE_OPTION, output_raster: OUTPUT_FILE_OPTION, new_nodata: float = typer.Option()

eis_toolkit/utilities/nodata.py

@@ -4,7 +4,7 @@
 import numpy as np
 import rasterio
 from beartype import beartype
-from beartype.typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
+from beartype.typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple, Union
 from rasterio import profiles
 
 from eis_toolkit.exceptions import InvalidParameterValueException
@@ -101,6 +101,60 @@ def convert_raster_nodata(
     return out_image, out_meta
 
 
+@beartype
+def replace_with_nodata(
+    input_raster: rasterio.io.DatasetReader,
+    target_value: Number,
+    nodata_value: Optional[Number] = None,
+    replace_condition: Literal[
+        "equal", "less_than", "greater_than", "less_than_or_equal", "greater_than_or_equal"
+    ] = "equal",
+) -> Tuple[np.ndarray, dict]:
+    """
+    Replace pixels values with nodata for a raster.
+
+    Can be used either for replacing all pixels with certain value with nodata, or for replacing all pixels with
+    values less than, greater than, less than or equal to, or greater than or equal to the target value with nodata.
+
+    Args:
+        input_raster: Input raster dataset.
+        target_value: Value to be replaced with nodata.
+        nodata_value: Value that will be used as nodata. If not provided, nodata is determined from input raster.
+        replace_condition: Whether to replace pixels with certain value or values less than, greater than, less than or
+            equal, or greater than or equal to target_value with nodata.
+
+    Returns:
+        The input raster data with specified pixels replaced with nodata, and updated metadata.
+
+    Raises:
+        InvalidParameterValueException: Nodata is provided and not found in the input raster.
+    """
+
+    if nodata_value is None:
+        nodata_value = input_raster.nodata
+        if nodata_value is None:
+            raise InvalidParameterValueException("Nodata not provided and not found in the input raster.")
+
+    raster_arr = input_raster.read()
+
+    if replace_condition == "equal":
+        values_to_replace = target_value
+    elif replace_condition == "less_than":
+        values_to_replace = raster_arr[raster_arr < target_value].tolist()
+    elif replace_condition == "greater_than":
+        values_to_replace = raster_arr[raster_arr > target_value].tolist()
+    elif replace_condition == "less_than_or_equal":
+        values_to_replace = raster_arr[raster_arr <= target_value].tolist()
+    elif replace_condition == "greater_than_or_equal":
+        values_to_replace = raster_arr[raster_arr >= target_value].tolist()
+
+    out_image = replace_values(raster_arr, values_to_replace, nodata_value)
+    out_meta = input_raster.meta.copy()
+    out_meta["nodata"] = nodata_value
+
+    return out_image, out_meta
+
+
 @beartype
 def nodata_to_nan(data: np.ndarray, nodata_value: Number) -> np.ndarray:
     """Convert specified nodata_value to np.nan.

tests/utilities/nodata_test.py

@@ -8,6 +8,7 @@
     handle_nodata_as_nan,
     nan_to_nodata,
     nodata_to_nan,
+    replace_with_nodata,
     set_raster_nodata,
     unify_raster_nodata,
 )
@@ -38,6 +39,76 @@ def test_convert_raster_nodata():
         assert out_meta["nodata"] == -999
 
 
+def test_replace_with_nodata():
+    """Test that replacing raster pixel values with nodata works as expected."""
+    target_value = 2.705
+    nodata_value = -999
+
+    with rasterio.open(SMALL_RASTER_PATH) as raster:
+        raster_data = raster.read()
+        nr_of_pixels = np.count_nonzero(raster_data == target_value)
+        assert nr_of_pixels > 0
+        assert np.count_nonzero(raster_data == nodata_value) == 0
+
+        replace_condition = "equal"
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+        assert np.count_nonzero(out_image == nodata_value) > 0
+        assert np.count_nonzero(out_image == target_value) == 0
+        assert out_meta["nodata"] == -999
+
+    with rasterio.open(SMALL_RASTER_PATH) as raster:
+        raster_data = raster.read()
+        # Ensure some pixels exist that are less than the target value
+        nr_of_pixels_less_than_target = np.count_nonzero(raster_data < target_value)
+        assert nr_of_pixels_less_than_target > 0
+
+        replace_condition = "less_than"
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+
+        assert np.count_nonzero(out_image == nodata_value) == nr_of_pixels_less_than_target
+        assert np.count_nonzero((out_image < target_value) & (out_image != nodata_value)) == 0
+        assert out_meta["nodata"] == -999
+
+    with rasterio.open(SMALL_RASTER_PATH) as raster:
+        raster_data = raster.read()
+        # Ensure some pixels exist that are greater than the target value
+        nr_of_pixels_greater_than_target = np.count_nonzero(raster_data > target_value)
+        assert nr_of_pixels_greater_than_target > 0
+
+        replace_condition = "greater_than"
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+
+        assert np.count_nonzero(out_image == nodata_value) == nr_of_pixels_greater_than_target
+        assert np.count_nonzero(out_image > target_value) == 0
+        assert out_meta["nodata"] == -999
+
+    with rasterio.open(SMALL_RASTER_PATH) as raster:
+        raster_data = raster.read()
+        # Ensure some pixels exist that are less than or equal to the target value
+        nr_of_pixels_less_than_or_equal = np.count_nonzero(raster_data <= target_value)
+        assert nr_of_pixels_less_than_or_equal > 0
+
+        replace_condition = "less_than_or_equal"
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+
+        assert np.count_nonzero(out_image == nodata_value) == nr_of_pixels_less_than_or_equal
+        assert np.count_nonzero((out_image <= target_value) & (out_image != nodata_value)) == 0
+        assert out_meta["nodata"] == -999
+
+    with rasterio.open(SMALL_RASTER_PATH) as raster:
+        raster_data = raster.read()
+        # Ensure some pixels exist that are greater than or equal to the target value
+        nr_of_pixels_greater_than_or_equal = np.count_nonzero(raster_data >= target_value)
+        assert nr_of_pixels_greater_than_or_equal > 0
+
+        replace_condition = "greater_than_or_equal"
+        out_image, out_meta = replace_with_nodata(raster, target_value, nodata_value, replace_condition)
+
+        assert np.count_nonzero(out_image == nodata_value) == nr_of_pixels_greater_than_or_equal
+        assert np.count_nonzero(out_image >= target_value) == 0
+        assert out_meta["nodata"] == -999
+
+
 def test_unify_raster_nodata():
     """Test that unifying raster nodata for multiple rasters works as expected."""
     with rasterio.open(SMALL_RASTER_PATH) as raster: