Fix one print issue, enable multiclass predicting for classifiers

Author: nmaarnio

eis_toolkit/cli.py

@@ -2359,8 +2359,10 @@ def classifier_test_cli(
     typer.echo(f"Results: {json_str}")
 
     typer.echo(
-        f"Testing classifier model completed, writing rasters to \
-        {output_raster_probability} and {output_raster_classified}."
+        (
+            "Testing classifier model completed, writing rasters to "
+            f"{output_raster_probability} and {output_raster_classified}."
+        )
     )
 
 

eis_toolkit/prediction/machine_learning_predict.py

@@ -18,13 +18,12 @@ def predict_classifier(
     """
     Predict with a trained classifier model.
 
-    Only works for binary classification currently.
-
     Args:
         data: Data used to make predictions.
         model: Trained classifier or regressor. Can be any machine learning model trained with
             EIS Toolkit (Sklearn and Keras models).
-        classification_threshold: Threshold for classifying based on probabilities. Defaults to 0.5.
+        classification_threshold: Threshold for classifying based on probabilities. Only used for
+            binary classification. Defaults to 0.5.
         include_probabilities: If the probability array should be returned too. Defaults to True.
 
     Returns:
@@ -34,19 +33,27 @@ def predict_classifier(
         InvalidModelTypeException: Input model is not a classifier model.
     """
     if isinstance(model, keras.Model):
-        probabilities = model.predict(data).squeeze()
-        labels = probabilities >= classification_threshold
+        probabilities = model.predict(data).astype(np.float32)
+        if probabilities.shape[1] == 1:  # Binary classification
+            probabilities = probabilities.squeeze()
+            labels = (probabilities >= classification_threshold).astype(np.float32)
+        else:  # Multiclass classification
+            labels = probabilities.argmax(axis=-1).astype(np.float32)
         if include_probabilities:
-            return labels, probabilities.astype(np.float32)
+            return labels, probabilities
         else:
             return labels
     elif isinstance(model, BaseEstimator):
         if not is_classifier(model):
             raise InvalidModelTypeException(f"Expected a classifier model: {type(model)}.")
-        probabilities = model.predict_proba(data)[:, 1]
-        labels = (probabilities >= classification_threshold).astype(np.float32)
+        probabilities = model.predict_proba(data).astype(np.float32)
+        if probabilities.shape[1] == 2:  # Binary classification
+            probabilities = probabilities[:, 1]
+            labels = (probabilities >= classification_threshold).astype(np.float32)
+        else:  # Multiclass classification
+            labels = probabilities.argmax(axis=-1).astype(np.float32)
         if include_probabilities:
-            return labels, probabilities.astype(np.float32)
+            return labels, probabilities
         else:
             return labels
     else:

tests/prediction/machine_learning_general_test.py

@@ -118,7 +118,7 @@ def test_evaluate_model_sklearn():
         X_train, y_train, model=RF_MODEL, validation_method="none", metrics=CLF_METRICS, random_state=42
     )
 
-    predictions = predict_classifier(X_test, model, include_probabilities=False)
+    predictions = predict_classifier(X_test, model, classification_threshold=0.5, include_probabilities=False)
     accuracy = score_predictions(y_test, predictions, "accuracy")
     np.testing.assert_equal(accuracy, 1.0)
 
@@ -131,7 +131,7 @@ def test_predict_classifier_sklearn():
         X_train, y_train, model=RF_MODEL, validation_method="none", metrics=CLF_METRICS, random_state=42
     )
 
-    predicted_labels, predicted_probabilities = predict_classifier(X_test, model, True)
+    predicted_labels, predicted_probabilities = predict_classifier(X_test, model, include_probabilities=True)
     np.testing.assert_equal(len(predicted_labels), len(y_test))
     np.testing.assert_equal(len(predicted_probabilities), len(y_test))