diff --git a/Source/SVMWrapper.cpp b/Source/SVMWrapper.cpp
index 5424619..314bae2 100644
--- a/Source/SVMWrapper.cpp
+++ b/Source/SVMWrapper.cpp
@@ -15,6 +15,7 @@ extern "C" void __cdecl print_log(const char *str) {
 struct SVMWrapper {
     struct svm_model *model;
     int cols;
+    int classes;
     double *means;
     double *sdevs;
     struct svm_parameter param;
@@ -86,7 +87,7 @@ void destroy_problem(struct svm_problem *problem) {
 }
 
 extern "C" {
-    __declspec(dllexport) void *__cdecl trainModel(float features[], float labels[], int rows, int cols) {
+    __declspec(dllexport) void *__cdecl trainModel(float features[], float labels[], int rows, int cols, int classes, int probability, float C) {
         if (!(rows > 0 && cols > 0))
             return NULL;
 
@@ -94,6 +95,7 @@ extern "C" {
 
         struct SVMWrapper *sw = new SVMWrapper();
         sw->cols = cols;
+        sw->classes = classes;
         sw->means = new double[cols];
         sw->sdevs = new double[cols];
 
@@ -104,11 +106,11 @@ extern "C" {
         sw->param.coef0 = 0;
         sw->param.nu = 0.5;
         sw->param.cache_size = 100;
-        sw->param.C = 4;
+        sw->param.C = C;
         sw->param.eps = 0.001;
         sw->param.p = 0.1;
         sw->param.shrinking = 1;
-        sw->param.probability = 0;
+        sw->param.probability = probability;
         sw->param.nr_weight = 0;
         sw->param.weight_label = NULL;
         sw->param.weight = NULL;
@@ -120,22 +122,23 @@ extern "C" {
         return sw;
     }
 
-    __declspec(dllexport) void __cdecl predict(void *ptr, float features[], float predictions[], int rows) {
+    __declspec(dllexport) void __cdecl predict(void *ptr, float features[], float predictions[], double probabilities[], int rows) {
         struct SVMWrapper *sw = (struct SVMWrapper *)ptr;
         int cols = sw->cols;
 
         svmnode **x = make_svmnode_array(features, rows, cols, sw->means, sw->sdevs, 0);
         for (int r = 0; r < rows; r++)
-            predictions[r] = (float)svm_predict(sw->model, x[r]);
+            predictions[r] = (float)svm_predict_probability(sw->model, x[r], &(probabilities[r * sw->classes]));
         destroy_svmnode_array(x);
 
         return;
     }
 
-    __declspec(dllexport) void *__cdecl loadModel(char *filename, int cols, double means[], double sdevs[]) {
+    __declspec(dllexport) void *__cdecl loadModel(char *filename, int cols, int classes, double means[], double sdevs[]) {
         struct SVMWrapper *sw = new SVMWrapper();
         sw->cols = cols;
         sw->model = svm_load_model(filename);
+        sw->classes = classes;
         sw->means = new double[cols];
         sw->sdevs = new double[cols];
         for (int c = 0; c < cols; c++) {
@@ -145,10 +148,11 @@ extern "C" {
         return sw;
     }
 
-    __declspec(dllexport) void *__cdecl loadModelString(char *modelString, int cols, double means[], double sdevs[]) {
+    __declspec(dllexport) void *__cdecl loadModelString(char *modelString, int cols, int classes, double means[], double sdevs[]) {
         struct SVMWrapper *sw = new SVMWrapper();
         sw->cols = cols;
         sw->model = svm_load_model_string(modelString);
+        sw->classes = classes;
         sw->means = new double[cols];
         sw->sdevs = new double[cols];
         for (int c = 0; c < cols; c++) {
@@ -168,7 +172,7 @@ extern "C" {
         }
     }
 
-    __declspec(dllexport) char *__cdecl getModel(void *ptr) {
+    __declspec(dllexport) char *__cdecl saveModelString (void *ptr) {
         struct SVMWrapper *sw = (struct SVMWrapper *)ptr;
         return svm_save_model_string(sw->model);
     }
diff --git a/Unity/OpenSeeExpression.cs b/Unity/OpenSeeExpression.cs
index 4143d62..ca251c8 100644
--- a/Unity/OpenSeeExpression.cs
+++ b/Unity/OpenSeeExpression.cs
@@ -25,6 +25,8 @@ public class OpenSeeExpression : MonoBehaviour
     public int overRecordingSkip = 5;
     [Tooltip("This is the filename used for loading and saving expression data using the load and save flags.")]
     public string filename = "";
+    [Tooltip("When enabled, training will be done in such a way that the model can output probabilities. This is a lot slower!")]
+    public bool enableProbabilityTraining = false;
     [Header("Toggles")]
     [Tooltip("When enabled, calibration data will be collected from the given OpenSee component.")]
     public bool recording = false;
@@ -54,6 +56,8 @@ public class OpenSeeExpression : MonoBehaviour
     public float expressionTime = 0f;
     [Tooltip("This is currently detected expression.")]
     public string expression = "";
+    [Tooltip("These are the probabilities of all expressions. The labels in order can be retrieved through the GetClassLabels methods.")]
+    public double[] probabilities;
     [Header("Training statistics")]
     [Tooltip("This shows the accuracy result of the last training run.")]
     public float accuracy = 0.0f;
@@ -362,7 +366,10 @@ public class OpenSeeExpression : MonoBehaviour
                 i_test++;
             }
         }
-        model.TrainModel(X_train, y_train, rows_train, cols);
+        int probability = 0;
+        if (enableProbabilityTraining)
+            probability = 1;
+        model.TrainModel(X_train, y_train, rows_train, cols, probability);
         confusionMatrix = model.ConfusionMatrix(X_test, y_test, i_test, out accuracy);
         confusionMatrixString = SVMModel.FormatMatrix(confusionMatrix, classLabels);
         for (int label = 0; label < classes; label++) {
@@ -399,7 +406,7 @@ public class OpenSeeExpression : MonoBehaviour
         float[] predictionData = new float[cols];
         for (int i = 0; i < cols; i++)
             predictionData[i] = faceData[indices[i]];
-        float[] prediction = model.Predict(predictionData, 1);
+        float[] prediction = model.Predict(predictionData, out probabilities, 1);
         int predictedExpression = (int)Mathf.Round(prediction[0]);
         if (predictedExpression == lastPrediction)
             lastPredictionCount++;
@@ -464,6 +471,13 @@ public class OpenSeeExpression : MonoBehaviour
             return false;
         }
     }
+    
+    public string[] GetClassLabels() {
+        List<string> labels = new List<string>();
+        foreach (string l in classLabels)
+            labels.Add(String.Copy(l));
+        return labels.ToArray();
+    }
 
     public string[] GetExpressions() {
         List<string> keys = new List<string>();
diff --git a/Unity/Plugins/SVMModel.cs b/Unity/Plugins/SVMModel.cs
index 9beb8d1..ef5a0ad 100644
--- a/Unity/Plugins/SVMModel.cs
+++ b/Unity/Plugins/SVMModel.cs
@@ -12,16 +12,16 @@ namespace OpenSee {
 public class SVMModel {
 	#region DllImport
 	[DllImport("SVMModel", CallingConvention = CallingConvention.Cdecl)]
-	private static extern System.IntPtr trainModel(float[] features, float[] labels, int rows, int cols, float C);
+	private static extern System.IntPtr trainModel(float[] features, float[] labels, int rows, int cols, int classes, int probability, float C);
 
 	[DllImport("SVMModel", CallingConvention = CallingConvention.Cdecl)]
-	private static extern void predict(System.IntPtr model, float[] features, [Out] float[] predictions, int rows);
+	private static extern void predict(System.IntPtr model, float[] features, [Out] float[] predictions, [Out] double[] probabilities, int rows);
 
 	[DllImport("SVMModel", CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
-	private static extern System.IntPtr loadModel([MarshalAs(UnmanagedType.LPStr)]string filename, int cols, double[] means, double[] sdevs);
+	private static extern System.IntPtr loadModel([MarshalAs(UnmanagedType.LPStr)]string filename, int cols, int classes, double[] means, double[] sdevs);
 
 	[DllImport("SVMModel", CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
-	private static extern System.IntPtr loadModelString([MarshalAs(UnmanagedType.LPStr)]string modelString, int cols, double[] means, double[] sdevs);
+	private static extern System.IntPtr loadModelString([MarshalAs(UnmanagedType.LPStr)]string modelString, int cols, int classes, double[] means, double[] sdevs);
 
 	[DllImport("SVMModel", CallingConvention = CallingConvention.Cdecl)]
 	private static extern void getScales(System.IntPtr model, [Out] double[] means, [Out] double[] sdevs);
@@ -63,7 +63,7 @@ public class SVMModel {
             }
             cols = smr.cols;
             maxClasses = smr.maxClasses;
-            System.IntPtr model = loadModelString(smr.modelString, smr.cols, smr.means, smr.sdevs);
+            System.IntPtr model = loadModelString(smr.modelString, smr.cols, smr.maxClasses, smr.means, smr.sdevs);
             return model;
         }
 
@@ -118,7 +118,7 @@ public class SVMModel {
     }
 
     // features is a rows*cols long float array with rows written one after another. labels is a rows long float array with the corresponding classes as integral numbers from 0 to num_classes-1, where all classes have to exist. Training on more than 10000 rows is disabled.
-    public bool TrainModel(float[] features, float[] labels, int rows, int cols) {
+    public bool TrainModel(float[] features, float[] labels, int rows, int cols, int probability) {
         DestroyModel();
         if (rows > 10000)
             return false;
@@ -143,7 +143,7 @@ public class SVMModel {
 
         maxClasses = max + 1;
         this.cols = cols;
-        model = trainModel(features, labels, rows, cols, 12f);
+        model = trainModel(features, labels, rows, cols, maxClasses, probability, 2f);
         haveModel = true;
 
         return true;
@@ -163,7 +163,8 @@ public class SVMModel {
         }
         int[,] confusionMatrix = new int[maxClasses, maxClasses];
         float[] predictions = new float[rows];
-        predict(model, features, predictions, rows);
+        double[] probabilities = new double[rows * maxClasses];
+        predict(model, features, predictions, probabilities, rows);
         for (int i = 0; i < rows; i++) {
             predictions[i] = (float)Math.Round(predictions[i]);
             if (predictions[i] >= 0 && predictions[i] < maxClasses) {
@@ -219,15 +220,22 @@ public class SVMModel {
     }
 
     // features is a rows*cols long float array with rows written one after another.
-    public float[] Predict(float[] features, int rows) {
+    public float[] Predict(float[] features, out double[] probabilities, int rows) {
+        probabilities = null;
         if (!haveModel)
             return null;
         if (rows < 1 || features.Length != rows * cols)
             return null;
         float[] predictions = new float[rows];
-        predict(model, features, predictions, rows);
+        probabilities = new double[rows * maxClasses];
+        predict(model, features, predictions, probabilities, rows);
         return predictions;
     }
+    
+    public float[] Predict(float[] features, int rows) {
+        double[] probabilities;
+        return Predict(features, out probabilities, rows);
+    }
 
     public bool Ready() {
         return haveModel;
diff --git a/Unity/Plugins/SVMModel.dll b/Unity/Plugins/SVMModel.dll
index 831c85c..a273c8c 100644
Binary files a/Unity/Plugins/SVMModel.dll and b/Unity/Plugins/SVMModel.dll differ