Refactored Sentinel-2 preprocessing (openeo) to be useful for S3

Author: Hneuschmidt

efast/openeo/__init__.py

@@ -0,0 +1,5 @@
+from .general import connect, TestArea
+from .general import extract_mask, distance_to_clouds
+from . import s2, s3
+
+__all__ = ["TestArea", "connect", "extract_mask", "s2", "s3", "distance_to_clouds"]

efast/openeo/preprocessing/__init__.py

@@ -1,14 +1,16 @@
-from typing import Mapping, Sequence
+import operator
+from typing import Tuple
 import openeo
 import numpy as np
 import scipy
+from collections.abc import Callable, Mapping, Iterable
 
 
 def connect():
     return openeo.connect("https://openeo.dataspace.copernicus.eu/")
 
 
-# TODO move somewhere else, not specific to sentinel-2
+
 class TestArea:
     # aoi_wkt = "POINT (-15.432283 15.402828)"
     directions = ["west", "south", "east", "north"]
@@ -19,14 +21,14 @@ def __init__(
         self,
         *,
         bbox: Mapping[str, float] = bbox,
-        s2_bands: Sequence[str] = ["B02", "B03", "B04", "B8A", "SCL"],
-        s3_bands: Sequence[str] = [
+        s2_bands: Iterable[str] = ["B02", "B03", "B04", "B8A", "SCL"],
+        s3_bands: Iterable[str] = [
             "Syn_Oa04_reflectance",
             "Syn_Oa06_reflectance",
             "Syn_Oa08_reflectance",
             "Syn_Oa17_reflectance",
             ],
-        temporal_extent: Sequence[str] = ["2022-06-01", "2022-06-30"],
+        temporal_extent: Iterable[str] = ["2022-06-01", "2022-06-30"],
     ) -> None:
         self.bbox = bbox
         self.s2_bands = s2_bands
@@ -50,39 +52,6 @@ def get_s3_cube(self, connection):
         )
 
 
-def extract_cloud_mask(cube: openeo.DataCube):
-    scl = cube.filter_bands(["SCL"])
-    # 0: No data
-    # 3: Cloud shadow
-    # 8-10: Clouds
-    # 11: Snow or ice
-
-    mask = (scl == 0) | (scl == 3) | (scl > 7)
-    return mask
-
-
-def calculate_large_grid_cloud_mask(cube: openeo.DataCube, tolerance_percentage: float = 0.05, grid_side_length: int=300):
-    cloud_mask = extract_cloud_mask(cube)
-    # FIXME check also if there is negative or zero data, otherwise results will differ
-    
-    # TODO this could better be resample_cube_spatial, because we are matching to a sentinel-3 cube
-    cloud_mask = cloud_mask * 1.0 # convert to float
-    cloud_mask_resampled = cloud_mask.resample_spatial(grid_side_length, method="average") # resample to sentinel-3 size
-
-    # TODO extract UDF to file
-    # UDF to apply an element-wise less than operation. Normal "<" does not properly work on openEO datacubes
-    udf = openeo.UDF(f"""
-import numpy as np
-import xarray as xr
-def apply_datacube(cube: XarrayDataCube, context: dict) -> XarrayDataCube:
-    array = cube.get_array()
-    array = array < {tolerance_percentage}
-    #return XarrayDataCube(xr.DataArray(array, dims=["t", "x", "y", "bands"]))
-    return XarrayDataCube(xr.DataArray(array, dims=["bands", "x", "y"]))
-            """)
-    return cloud_mask_resampled.apply(process=udf)
-
-
 def distance_to_clouds(
     cube: openeo.DataCube, tolerance_percentage=0.05, ratio=30, max_distance=255
 ):
@@ -116,3 +85,66 @@ def _distance_to_clouds_udf(
         overlap=[],
     )
     return dtc
+
+
+def extract_mask(
+        cube: openeo.DataCube,
+        mask_values: Mapping[str, Iterable[int]],
+        *,
+        operations: Mapping[Tuple[str, int], Callable],
+        ) -> openeo.DataCube:
+    """
+    Generic method to extract a mask from a data cube.
+    Generate a mask that has a value of ``True`` whereever the band specified
+    as a key in ``mask_values`` is equal to one of the values speicified
+    as a value. Operations other than equality comparison can be specified
+    via ``operations``.
+
+        import operator
+        extract_mask(
+                cube,
+                mask_values = {
+                    "SCL": (0, 3, 7),
+                    },
+                operations = {
+                    # use ``>`` instead of ``==``
+                    # to compare band ``"SCL"`` to value ``7``
+                    ("SCL", 7): operator.gt,
+                    }
+                )
+
+        # scl = cube.band("SCL")
+        # mask = (scl == 0) | (scl == 3) | (scl > 7)
+
+
+    :param cube: The data cube containing at least the band(s) used for masking
+    :param mask_values: Mapping of band names to the values that should be masked.
+    :param operations: Used to specify a comparison operation different to ``==``
+        when comparing bands and values. Operations are applied as ``op(band, value)``
+    
+    """
+
+    assert(len(mask_values) > 0), "'mask_values' cannot be empty."
+    def reduce_single_band(band_name, values_to_mask, mask=None):
+        band = cube.band(band_name)
+        vm_iter = iter(values_to_mask)
+
+        if mask is None:
+            vm_first = next(vm_iter)
+            mask = band == vm_first
+
+        for vm in vm_iter:
+            op = operations.get((band_name, vm), operator.eq)
+            mask |= op(band, vm)
+
+        return mask
+
+    first_band_name, *bands = mask_values.keys()
+    first_vm, *values_to_mask = mask_values.values()
+
+    mask = reduce_single_band(first_band_name, first_vm, mask=None)
+
+    for band_name, values_to_mask in zip(bands, values_to_mask):
+        mask |= reduce_single_band(band_name, values_to_mask, mask)
+
+    return mask

efast/s2_processing_openeo.py efast/openeo/preprocessing/general.pyefast/s2_processing_openeo.py renamed to efast/openeo/preprocessing/general.py

@@ -0,0 +1,24 @@
+import operator
+
+import openeo
+
+from . import extract_mask
+
+
+def extract_cloud_mask_s2(cube: openeo.DataCube) -> openeo.DataCube:
+    return extract_mask(
+        cube,
+        {"SCL": [0, 3, 7]},
+        operations={
+            ("SCL", 7): operator.gt,
+        },
+    )
+
+def calculate_large_grid_cloud_mask(cube: openeo.DataCube, tolerance_percentage: float = 0.05, grid_side_length: int=300):
+    cloud_mask = extract_cloud_mask_s2(cube)
+    # FIXME check also if there is negative or zero data, otherwise results will differ
+    
+    # TODO this could better be resample_cube_spatial, because we are matching to a sentinel-3 cube
+    cloud_mask = cloud_mask * 1.0 # convert to float
+    cloud_mask_resampled = cloud_mask.resample_spatial(grid_side_length, method="average") # resample to sentinel-3 size
+    return cloud_mask_resampled < tolerance_percentage

efast/openeo/preprocessing/s2.py

@@ -1,6 +1,7 @@
 from collections.abc import Generator
 import shutil
 import time
+import pytest
 
 from contextlib import contextmanager
 from pathlib import Path
@@ -17,7 +18,11 @@
 
 from openeo.udf import execute_local_udf
 
-from efast import s2_processing, s2_processing_openeo, s3_processing
+from efast.openeo import preprocessing
+# TODO this should live somewhere else
+from efast.openeo.preprocessing import connect
+from efast import s2_processing, s3_processing
+
 
 TEST_DATA_ROOT = Path(__file__).parent.parent / "test_data"
 TEST_DATA_S2 = TEST_DATA_ROOT / "S2"
@@ -33,7 +38,8 @@
 TEST_DATE = "20220618"
 TEST_DATE_DASH = "2022-06-18"
 
-SKIP_LOCAL = False
+SKIP_LOCAL = True
+DOWNLOAD_INTERMEDIATE_RESULTS = True
 SMALL_AREA = True
 
 
@@ -57,6 +63,7 @@ def create_temp_dir_and_copy_files(
 ############################## S2 Processing ##############################
 
 
+@pytest.mark.skip
 def test_distance_to_cloud():
     with create_temp_dir_and_copy_files(TEST_DATA_S2, sub="raw/", pattern=None) as tmp:
         tmp_path = Path(tmp.name)
@@ -86,39 +93,41 @@ def test_distance_to_cloud():
 
         # openeo
 
-        conn = s2_processing_openeo.connect()
+        conn = connect()
         conn.authenticate_oidc()
 
-        test_area = s2_processing_openeo.TestArea(
+        test_area = preprocessing.TestArea(
             bbox=bounds,
             s2_bands=["SCL"],
             temporal_extent=(TEST_DATE_DASH, TEST_DATE_DASH),
         )
         cube = test_area.get_s2_cube(conn)
 
-        dtc_input = s2_processing_openeo.calculate_large_grid_cloud_mask(
+        dtc_input = preprocessing.s2.calculate_large_grid_cloud_mask(
             cube, tolerance_percentage=0.05, grid_side_length=300
         )
-        dtc = s2_processing_openeo.distance_to_clouds(
+        dtc = preprocessing.distance_to_clouds(
             dtc_input, tolerance_percentage=0.05, ratio=30
         )
         download_path = tmp_path / "test_distance_to_cloud.tif"
 
         print("openEO execution")
-
-        # intermediate results for debugging
-        BASE_DIR = Path("openeo_results")
-        # BASE_DIR.mkdir(exist_ok=True)
-        # cloud_mask.download(BASE_DIR / "cloud_mask.tif")
-        # cloud_mask_resampled.download(BASE_DIR / "cloud_mask_resampled.tif")
-        # dtc_input.download(BASE_DIR / "dtc_input.tif")
-
         before = time.perf_counter()
-        dtc.download(download_path)
-        shutil.copy(download_path, BASE_DIR)
+        #dtc.download(download_path)
         elapsed = time.perf_counter() - before
         print(f"executed and downloaded in {elapsed:.2f}s")
 
+        if DOWNLOAD_INTERMEDIATE_RESULTS:
+            BASE_DIR = Path("openeo_results")
+            BASE_DIR.mkdir(exist_ok=True)
+            print("downloading input")  # TMP
+            (dtc_input * 1.0).download(BASE_DIR / "dtc_input.tif")
+            #shutil.copy(download_path, BASE_DIR)
+        print("downloading result") # TMP
+        dtc.download(download_path) # TMP
+        shutil.copy(download_path, BASE_DIR) # TMP
+
+
         with rasterio.open(download_path, "r") as ds:
             dtc_openeo = ds.read(1)
 
@@ -140,7 +149,8 @@ def test_distance_to_cloud_synthetic_cube():
         .mean(1)
     ) < tolerance
     cube = xr.DataArray(cube, dims=["x", "y"])
-    cube_resampled = xr.DataArray(cube_resampled, dims=["x", "y"])
+    #cube = cube.add_dimension(name="bands", label="mask", type="bands")
+    cube_resampled = xr.DataArray(cube_resampled[np.newaxis, np.newaxis], dims=["bands", "t", "x", "y"])
 
     udf = openeo.UDF.from_file("efast/distance_transform_udf.py")
     dtc_local_udf = (
@@ -245,7 +255,7 @@ def inner_data_acquisition_s3(tmpdir):
 
     # openeo
 
-    conn = s2_processing_openeo.connect()
+    conn = connect()
     conn.authenticate_oidc()
 
     bands = [
@@ -256,7 +266,7 @@ def inner_data_acquisition_s3(tmpdir):
         "CLOUD_flags",
         "OLC_flags",
     ]
-    test_area = s2_processing_openeo.TestArea(
+    test_area = preprocessing.TestArea(
         bbox=bounds, s3_bands=bands, temporal_extent=(TEST_DATE_DASH, TEST_DATE_DASH)
     )
     cube = test_area.get_s3_cube(conn)