batch_size as parameter

make batch_size a parameter of the feature extractor, to easily set it in the udf
WorldCereal · Jun 25, 2024 · de78e28 · de78e28
1 parent 98a182d
commit de78e28
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/worldcereal/openeo/feature_extractor.py b/src/worldcereal/openeo/feature_extractor.py
@@ -110,9 +110,11 @@ def execute(self, inarr: xr.DataArray) -> xr.DataArray:
             get_presto_features,
         )
 
+        batch_size = self._parameters.get("batch_size", 4096)
+
         self.logger.info("Extracting presto features")
         features = get_presto_features(
-            inarr, presto_model_url, self.epsg, batch_size=4096
+            inarr, presto_model_url, self.epsg, batch_size=batch_size
         )
         return features