add test for cultivated model

opendatacube · Aug 6, 2024 · 72b9fb4 · 72b9fb4
1 parent dbf76eb
commit 72b9fb4
Show file tree

Hide file tree

Showing 3 changed files with 478 additions and 9 deletions.
diff --git a/odc/stats/plugins/_worker.py b/odc/stats/plugins/_worker.py
@@ -44,6 +44,7 @@ def __init__(self, model_path):
     def setup(self, worker):
         worker.plugin_instance = self
         worker.predictors = {}
+        print(f"registered worker {worker}")
 
     def get_predictor(self):
         worker = get_worker()

diff --git a/odc/stats/plugins/lc_treelite_cultivated.py b/odc/stats/plugins/lc_treelite_cultivated.py
@@ -238,8 +238,8 @@ def predict(self, input_array):
             input_array,
             bands_indices,
             chunks=(
-                self.chunks["x"],
-                self.chunks["y"],
+                input_array.chunks[0],
+                input_array.chunks[1],
                 15 + len(bands_indices) - bands_indices["bcdev"] - 1,
             ),
             dtype="float32",
@@ -254,6 +254,7 @@ def predict(self, input_array):
             dtype="float32",
             name="cultivated_predict",
         )
+
         return cc
 
     def aggregate_results_from_group(self, predict_output):
@@ -263,10 +264,21 @@ def aggregate_results_from_group(self, predict_output):
         # for each pixel
         m_size = len(predict_output)
         if m_size > 1:
-            predict_output = da.stack(predict_output).sum(axis=0)
+            predict_output = da.stack(predict_output)
         else:
             predict_output = predict_output[0]
 
+        predict_output = expr_eval(
+            "where(a<nodata, 1-a, a)",
+            {"a": predict_output},
+            name="invert_output",
+            dtype="float32",
+            **{"nodata": NODATA},
+        )
+
+        if m_size > 1:
+            predict_output = predict_output.sum(axis=0)
+
         predict_output = expr_eval(
             "where((a/nodata)>=_l, nodata, a%nodata)",
             {"a": predict_output},
@@ -276,19 +288,19 @@ def aggregate_results_from_group(self, predict_output):
         )
 
         predict_output = expr_eval(
-            "where((a>=_l)&(a<nodata), _u, a)",
+            "where((a>0)&(a<nodata), _u, a)",
             {"a": predict_output},
-            name="output_classes_natural",
+            name="output_classes_cultivated",
             dtype="float32",
-            **{"_u": self.output_classes["natural"], "nodata": NODATA, "_l": m_size},
+            **{"_u": self.output_classes["cultivated"], "nodata": NODATA},
         )
 
         predict_output = expr_eval(
-            "where(a<_l, _nu, a)",
+            "where(a<=0, _nu, a)",
             {"a": predict_output},
-            name="output_classes_cultivated",
+            name="output_classes_natural",
             dtype="uint8",
-            **{"_nu": self.output_classes["cultivated"], "_l": m_size},
+            **{"_nu": self.output_classes["natural"]},
         )
 
         return predict_output.rechunk(-1, -1)