diff --git a/eis_toolkit/cli.py b/eis_toolkit/cli.py index 1ad52c24..3dce7288 100644 --- a/eis_toolkit/cli.py +++ b/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() diff --git a/eis_toolkit/utilities/nodata.py b/eis_toolkit/utilities/nodata.py index 3637e475..43456762 100644 --- a/eis_toolkit/utilities/nodata.py +++ b/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. diff --git a/tests/utilities/nodata_test.py b/tests/utilities/nodata_test.py index be5d5179..aaa88c25 100644 --- a/tests/utilities/nodata_test.py +++ b/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: