clean up comments and spacing

GeoscienceAustralia · Sep 30, 2024 · a6e937f · a6e937f
1 parent c50ff73
commit a6e937f
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/Tools/dea_tools/classification.py b/Tools/dea_tools/classification.py
@@ -336,7 +336,7 @@ def _predict_func(model, input_xr, persist, proba, max_proba, clean, return_inpu
             print("   probabilities...")
             out_proba = model.predict_proba(input_data_flattened)
 
-            # convert to %
+            # return either one band with the max probability, or the whole probability array
             if max_proba == True:
                 print("  returning single probability band.")
                 out_proba = da.max(out_proba, axis=1) * 100.0
@@ -348,12 +348,10 @@ def _predict_func(model, input_xr, persist, proba, max_proba, clean, return_inpu
             else:
                 print("  returning class probability array.")
                 out_proba = out_proba * 100.0
-
                 class_names = model.classes_  # Get the unique class names from the fitted classifier
 
-                probabilities_dataset = xr.Dataset()
-
                 # Loop through each class (band)
+                probabilities_dataset = xr.Dataset()
                 for i, class_name in enumerate(class_names):
                     reshaped_band = out_proba[:, i].reshape(len(y), len(x))
                     reshaped_da = xr.DataArray(