Merge pull request #417 from GispoCoding/414-add-masking-parameter-to…

…-unify-raster-grids

414 add masking parameter to unify raster grids

eis_toolkit/cli.py

@@ -308,6 +308,14 @@ class ThresholdCriteria(str, Enum):
     outside = "outside"
 
 
+class MaskingMode(str, Enum):
+    """Masking modes for raster unification."""
+
+    extents = "extents"
+    full = "full"
+    none = "none"
+
+
 class KerasClassifierMetrics(str, Enum):
     """Metrics available for Keras classifier models."""
 
@@ -1405,7 +1413,7 @@ def unify_rasters_cli(
     base_raster: INPUT_FILE_OPTION,
     output_directory: OUTPUT_DIR_OPTION,
     resampling_method: Annotated[ResamplingMethods, typer.Option(case_sensitive=False)] = ResamplingMethods.nearest,
-    same_extent: bool = False,
+    masking: Annotated[MaskingMode, typer.Option(case_sensitive=False)] = MaskingMode.extents,
 ):
     """Unify rasters to match the base raster."""
     from eis_toolkit.raster_processing.unifying import unify_raster_grids
@@ -1416,11 +1424,12 @@ def unify_rasters_cli(
         to_unify = [rasterio.open(rstr) for rstr in rasters_to_unify]  # Open all rasters to be unified
         typer.echo("Progress: 25%")
 
+        masking_param = get_enum_values(masking)
         unified = unify_raster_grids(
             base_raster=raster,
             rasters_to_unify=to_unify,
             resampling_method=get_enum_values(resampling_method),
-            same_extent=same_extent,
+            masking=None if masking_param == "none" else masking_param,
         )
         [rstr.close() for rstr in to_unify]  # Close all rasters
     typer.echo("Progress: 75%")

eis_toolkit/raster_processing/masking.py

@@ -38,7 +38,7 @@ def mask_raster(
         raster_arr = raster.read()
     else:
         out_rasters = unify_raster_grids(
-            base_raster=base_raster, rasters_to_unify=[raster], resampling_method="nearest", same_extent=True
+            base_raster=base_raster, rasters_to_unify=[raster], resampling_method="nearest", masking="extents"
         )
         raster_arr = out_rasters[1][0]
 

eis_toolkit/raster_processing/unifying.py

@@ -1,7 +1,8 @@
 import numpy as np
 import rasterio
+import rasterio.profiles
 from beartype import beartype
-from beartype.typing import List, Literal, Sequence, Tuple
+from beartype.typing import List, Literal, Optional, Sequence, Tuple
 from rasterio import warp
 from rasterio.enums import Resampling
 from rasterio.profiles import Profile
@@ -43,11 +44,33 @@ def _calculate_snapped_grid(
     return dst_transform, dst_width, dst_height
 
 
+def _mask_nodata(
+    raster_array: np.ndarray,
+    nodata_value: Optional[float],
+    base_raster_array: np.ndarray,
+    base_raster_profile: rasterio.profiles.Profile,
+):
+    # Take only first band as nodata mask from base raster
+    if base_raster_array.ndim != 1:
+        base_raster_array = base_raster_array[0]
+
+    # Create the mask
+    base_nodata_value = base_raster_profile.get("nodata", np.nan)
+    mask = (base_raster_array == base_nodata_value) | np.isnan(base_raster_array)
+
+    # Apply to each band
+    if raster_array.ndim != 1:
+        for array in raster_array:
+            array[mask] = nodata_value
+    else:
+        raster_array[mask] = nodata_value
+
+
 def _unify_raster_grids(
     base_raster: rasterio.io.DatasetReader,
     rasters_to_unify: Sequence[rasterio.io.DatasetReader],
     resampling_method: Resampling,
-    same_extent: bool,
+    masking: Optional[Literal["extents", "extents_and_nodata"]],
 ) -> List[Tuple[np.ndarray, Profile]]:
 
     dst_crs = base_raster.crs
@@ -56,19 +79,25 @@ def _unify_raster_grids(
     dst_transform = base_raster.transform
     dst_resolution = (base_raster.transform.a, abs(base_raster.transform.e))
 
-    out_rasters = [(base_raster.read(), base_raster.profile.copy())]
+    base_raster_arr = base_raster.read()
+    base_raster_profile = base_raster.profile.copy()
+    out_rasters = [(base_raster_arr, base_raster_profile)]
 
     for raster in rasters_to_unify:
         out_profile = raster.profile.copy()
 
         # If we unify without clipping, things are more complicated and we need to
         # calculate corner coordinates, width and height, and snap the grid to nearest corner
-        if not same_extent:
+        if not masking:
             dst_transform, dst_width, dst_height = _calculate_snapped_grid(raster, dst_crs, dst_resolution)
 
-        # Initialize output raster arrary
         dst_array = np.empty((base_raster.count, dst_height, dst_width))
+        base_raster_nodata = base_raster_profile.get("nodata", np.nan)
         nodata = out_profile["nodata"]
+        if nodata is None:
+            nodata = base_raster_nodata
+            out_profile["nodata"] = base_raster_nodata
+
         dst_array.fill(nodata)
 
         src_array = raster.read()
@@ -85,6 +114,9 @@ def _unify_raster_grids(
             resampling=resampling_method,
         )[0]
 
+        if masking == "full":
+            _mask_nodata(out_image, nodata, base_raster_arr, base_raster_profile)
+
         out_profile.update({"transform": dst_transform, "width": dst_width, "height": dst_height, "crs": dst_crs})
 
         out_rasters.append((out_image, out_profile))
@@ -97,18 +129,26 @@ def unify_raster_grids(
     base_raster: rasterio.io.DatasetReader,
     rasters_to_unify: Sequence[rasterio.io.DatasetReader],
     resampling_method: Literal["nearest", "bilinear", "cubic", "average", "gauss", "max", "min"] = "nearest",
-    same_extent: bool = False,
+    masking: Optional[Literal["extents", "full"]] = "extents",
 ) -> List[Tuple[np.ndarray, Profile]]:
-    """Unifies (reprojects, resamples, aligns and optionally clips) given rasters relative to base raster.
+    """Unifies given rasters with the base raster.
+
+    Performs the following operations:
+    - Reprojecting
+    - Resampling
+    - Aligning / snapping
+    - Clipping / expanding extents (optional)
+    - Copying nodata cells from base raster (optional)
 
     Args:
         base_raster: The base raster to determine target raster grid properties.
         rasters_to_unify: Rasters to be unified with the base raster.
-        resampling_method: Resampling method. Most suitable
-            method depends on the dataset and context. Nearest, bilinear and cubic are some
-            common choices. This parameter defaults to nearest.
-        same_extent: If the unified rasters will be forced to have the same extent/bounds
-            as the base raster. Expands smaller rasters with nodata cells. Defaults to False.
+        resampling_method: Resampling method. Most suitable method depends on the dataset and context.
+            `nearest`, `bilinear` and `cubic` are some common choices. This parameter defaults to `nearest`.
+        masking: Controls if and how masking should be handled. If `extents`, the bounds of rasters to-be-unified
+            are matched with the base raster. Larger rasters are clipped and smaller rasters expanded (with nodata).
+            If `full`, copies nodata pixel locations from the base raster additionally. If None,
+            extents are not matched and nodata not copied. Defaults to `extents`.
 
     Returns:
         List of unified rasters' data and profiles. First element is the base raster.
@@ -120,5 +160,5 @@ def unify_raster_grids(
         raise InvalidParameterValueException("Rasters to unify is empty.")
 
     method = RESAMPLE_METHOD_MAP[resampling_method]
-    out_rasters = _unify_raster_grids(base_raster, rasters_to_unify, method, same_extent)
+    out_rasters = _unify_raster_grids(base_raster, rasters_to_unify, method, masking)
     return out_rasters

tests/raster_processing/masking_test.py

@@ -1,32 +1,27 @@
-import tempfile
+from pathlib import Path
 
-import geopandas as gpd
 import numpy as np
 import rasterio
 
-from eis_toolkit.raster_processing.clipping import clip_raster
 from eis_toolkit.raster_processing.masking import mask_raster
-from tests.raster_processing.clip_test import polygon_path as SMALL_VECTOR_PATH
 from tests.raster_processing.clip_test import raster_path as SMALL_RASTER_PATH
 
+test_dir = Path(__file__).parent.parent
+small_raster_clipped_path = test_dir.joinpath("data/remote/small_raster_clipped.tif")
+
 
 def test_mask_raster():
     """Test that masking raster works as intended."""
     with rasterio.open(SMALL_RASTER_PATH) as raster:
-        geodataframe = gpd.read_file(SMALL_VECTOR_PATH)
-        out_image, out_meta = clip_raster(raster, geodataframe)
-        with tempfile.NamedTemporaryFile() as tmpfile:
-            with rasterio.open(tmpfile.name, "w", **out_meta) as dest:
-                dest.write(out_image)
-            with rasterio.open(tmpfile.name) as base_raster:
-                old_nodata_count = np.count_nonzero(raster.read(1) == raster.nodata)
-                out_image, out_profile = mask_raster(raster, base_raster)
+        with rasterio.open(small_raster_clipped_path) as base_raster:
+            old_nodata_count = np.count_nonzero(raster.read(1) == raster.nodata)
+            out_image, out_profile = mask_raster(raster, base_raster)
 
-                new_nodata_count = np.count_nonzero(out_image == out_profile["nodata"])
+            new_nodata_count = np.count_nonzero(out_image == out_profile["nodata"])
 
-                # Check nodata count has increased
-                assert new_nodata_count > old_nodata_count
-                # Check that nodata exists now in identical locations in input raster and base raster
-                np.testing.assert_array_equal(
-                    base_raster.read(1) == base_raster.nodata, out_image[0] == out_profile["nodata"]
-                )
+            # Check nodata count has increased
+            assert new_nodata_count > old_nodata_count
+            # Check that nodata exists now in identical locations in input raster and base raster
+            np.testing.assert_array_equal(
+                base_raster.read(1) == base_raster.nodata, out_image[0] == out_profile["nodata"]
+            )

tests/raster_processing/unify_raster_grids_test.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
 import rasterio
 
 from eis_toolkit.exceptions import InvalidParameterValueException
 from eis_toolkit.raster_processing.unifying import unify_raster_grids
+from tests.raster_processing.masking_test import small_raster_clipped_path as base_raster_path_3
 
 test_dir = Path(__file__).parent.parent
 
@@ -18,11 +20,9 @@
 raster_to_unify_save_path_2 = test_dir.joinpath("data/local/results/unify_test_result_2.tif")
 
 
-def test_unify_raster_grids_case1():
+def test_unify_raster_grids():
     """
-    Tests unify function for case 1.
-
-    Nearest resampling with same_extent set to False (no clipping).
+    Tests unify raster grids without masking.
 
     The raster-to-unify differes from the base raster in...
     - crs
@@ -32,7 +32,7 @@ def test_unify_raster_grids_case1():
 
     with rasterio.open(raster_to_unify_path_1) as raster_to_unify:
         with rasterio.open(base_raster_path_1) as base_raster:
-            out_rasters = unify_raster_grids(base_raster, [raster_to_unify], "nearest", False)
+            out_rasters = unify_raster_grids(base_raster, [raster_to_unify], "nearest", masking=None)
             out_image, out_meta = out_rasters[1]
 
             assert len(out_rasters) == 2
@@ -48,11 +48,9 @@ def test_unify_raster_grids_case1():
         dst.write(out_image)
 
 
-def test_unify_raster_grids_case2():
+def test_unify_raster_grids_extent():
     """
-    Tests unify function for case 2.
-
-    Bilinear resampling with same_extent set to True (clipping).
+    Tests unify grids with extent masking.
 
     The raster-to-unify differs from the base raster in...
     - pixel size
@@ -62,7 +60,7 @@ def test_unify_raster_grids_case2():
 
     with rasterio.open(raster_to_unify_path_2) as raster_to_unify:
         with rasterio.open(base_raster_path_2) as base_raster:
-            out_rasters = unify_raster_grids(base_raster, [raster_to_unify], "bilinear", True)
+            out_rasters = unify_raster_grids(base_raster, [raster_to_unify], "bilinear", masking="extents")
             out_image, out_meta = out_rasters[1]
 
             assert len(out_rasters) == 2
@@ -80,6 +78,40 @@ def test_unify_raster_grids_case2():
         dst.write(out_image)
 
 
+def test_unify_raster_grids_full_masking():
+    """
+    Tests unify raster grids with full masking (extent and nodata).
+
+    The raster-to-unify differs from the base raster in...
+    - pixel size
+    - alignment
+    - bounds
+    """
+
+    with rasterio.open(raster_to_unify_path_1) as raster_to_unify:
+        with rasterio.open(base_raster_path_3) as base_raster:
+            out_rasters = unify_raster_grids(base_raster, [raster_to_unify], "bilinear", masking="full")
+            out_image, out_profile = out_rasters[1]
+
+            assert len(out_rasters) == 2
+            assert out_profile["crs"] == base_raster.crs
+            # Check pixel size
+            assert out_profile["transform"].a == base_raster.transform.a
+            assert out_profile["transform"].e == base_raster.transform.e
+            # Check grid alignment and bounds
+            assert out_profile["transform"].c == base_raster.transform.c
+            assert out_profile["transform"].f == base_raster.transform.f
+            assert out_profile["width"] == base_raster.width
+            assert out_profile["height"] == base_raster.height
+            # Check nodata locations
+            np.testing.assert_array_equal(
+                base_raster.read(1) == base_raster.nodata, out_image[0] == out_profile["nodata"]
+            )
+
+    with rasterio.open(raster_to_unify_save_path_2, "w", **out_profile) as dst:
+        dst.write(out_image)
+
+
 def test_unify_raster_grids_empty_raster_list():
     """Test that empty raster list raises correct exception."""
     with pytest.raises(InvalidParameterValueException):