Onnx example

python/OnnxUdf/example.py

@@ -0,0 +1,35 @@
+import openeo
+
+# This will be required to import onnxruntime in the UDF.
+dependencies_url = "https://artifactory.vgt.vito.be:443/auxdata-public/openeo/onnx_dependencies.zip"
+# You can upload your own model and paste the url here.
+model_url = "https://artifactory.vgt.vito.be:443/auxdata-public/openeo/test_onnx_model.zip"
+
+spatial_extent = dict(zip(["west", "south", "east", "north"], [8.908, 53.791, 8.96, 54.016]))
+temporal_extent = ["2022-10-01", "2022-12-01"]
+
+connection = openeo.connect("openeo.vito.be").authenticate_oidc()
+
+s2_cube = connection.load_collection(
+    "TERRASCOPE_S2_TOC_V2",
+    temporal_extent=temporal_extent,
+    spatial_extent=spatial_extent,
+    bands=["B04"],
+    max_cloud_cover=20,
+)
+
+udf = openeo.UDF.from_file("onnx_udf.py")
+s2_cube = s2_cube.apply_neighborhood(
+    udf,
+    size=[{"dimension": "x", "value": 256, "unit": "px"}, {"dimension": "y", "value": 256, "unit": "px"}],
+    overlap=[{"dimension": "x", "value": 0, "unit": "px"}, {"dimension": "y", "value": 0, "unit": "px"}],
+)
+
+# We pass the model and dependencies as zip files to the UDF through the job options.
+job_options = {
+    "udf-dependency-archives": [
+        f"{dependencies_url}#tmp/extra_venv",
+        f"{model_url}#tmp/extra_files",
+    ],
+}
+s2_cube.execute_batch("output.nc", job_options=job_options)

python/OnnxUdf/onnx_udf.py

@@ -0,0 +1,22 @@
+import sys
+from openeo.udf import XarrayDataCube
+from typing import Dict
+import xarray as xr
+from openeo.udf.debug import inspect
+
+
+def apply_datacube(cube: XarrayDataCube, context: Dict) -> XarrayDataCube:
+    sys.path.insert(0, "tmp/extra_venv")
+    import onnxruntime as ort
+
+    input_data = cube.get_array().isel(t=0).values  # Only perform inference for the first date.
+    input_data = input_data[None, ...]  # Neural network expects shape (1, 1, 256, 256)
+    inspect(input_data, "input data")
+    ort_session = ort.InferenceSession("tmp/extra_files/test_model.onnx")
+    ort_inputs = {ort_session.get_inputs()[0].name: input_data}
+    ort_outputs = ort_session.run(None, ort_inputs)
+    output_data = xr.DataArray(ort_outputs[0])
+    output_data = output_data.rename({"dim_0": "t", "dim_1": "bands", "dim_2": "y", "dim_3": "x"})
+    inspect(output_data, "output data")
+
+    return XarrayDataCube(output_data)