Added inspection in udf and corrected type annotations

Author: Hneuschmidt

efast/distance_transform_udf.py

@@ -1,10 +1,13 @@
+from openeo.udf.debug import inspect
 from scipy.ndimage import distance_transform_edt
 import numpy as np
 import xarray as xr
 from openeo.udf import XarrayDataCube
 
 def apply_datacube(cube: XarrayDataCube, context: dict) -> XarrayDataCube:
     array = cube.get_array()
+    inspect(array.dims, message="array dims")
+    inspect(array.sizes, message="array sizes")
     distance = distance_transform_edt(array)
     clipped = np.clip(distance, 0, 255)
-    return XarrayDataCube(xr.DataArray(clipped, dims=["bands", "t", "y", "x"]))
+    return XarrayDataCube(xr.DataArray(clipped, dims=["t", "bands", "y", "x"]))

efast/openeo/preprocessing/general.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import openeo
+from openeo import processes
 import scipy
 
 
@@ -44,7 +45,7 @@ def get_s2_cube(self, connection):
             bands=self.s2_bands,
         )
 
-    def get_s3_cube(self, connection):
+    def get_s3_cube(self, connection) -> openeo.DataCube:
         return connection.load_collection(
             "SENTINEL3_SYN_L2_SYN",
             spatial_extent=self.bbox,
@@ -53,9 +54,10 @@ def get_s3_cube(self, connection):
         )
 
 
+# TODO ratio and max_distance arguments unused
 def distance_to_clouds(
-    cube: openeo.DataCube, tolerance_percentage=0.05, ratio=30, max_distance=255
-):
+    cube: processes.ProcessBuilder | openeo.DataCube, tolerance_percentage=0.05, ratio=30, max_distance=255
+) -> processes.ProcessBuilder:
     return _distance_to_clouds_udf(
         cube,
         tolerance_percentage=tolerance_percentage,
@@ -79,8 +81,8 @@ def _distance_to_clouds_kernel(
 # TODO implement max_distance as a parameter to the UDF
 # TODO replace hard coded tile size (366)
 def _distance_to_clouds_udf(
-    cube: openeo.DataCube, tolerance_percentage=0.05, ratio=30, max_distance=255
-):
+    cube: processes.ProcessBuilder | openeo.DataCube, tolerance_percentage=0.05, ratio=30, max_distance=255
+) -> processes.ProcessBuilder:
     udf = openeo.UDF.from_file("efast/distance_transform_udf.py")
     dtc = cube.apply_neighborhood(
         udf,