color , status, probablity sum added

app.py

@@ -83,9 +83,6 @@ class ModelPredictor:
 
             # If model can do predict_proba, we interpret the "2" class as the second column
             if hasattr(model, "predict_proba"):
-                # Usually classes_ might be something like [0,1], but we want [1,2]. 
-                # You may need to check model.classes_ or do a small transformation.
-                # For demonstration, we'll assume the second column is the "2" probability:
                 y_prob_2 = model.predict_proba(user_input)[:, 1]  
                 probs.append(y_prob_2)
             else:
@@ -129,8 +126,7 @@ categories_dict = {
     ]
 }
 
-# NOTE: input_mapping below for capturing user choices => numeric codes. 
-# If the models expect [1,2] for "Yes"/"No", ensure those are correct here too.
+# NOTE: input_mapping below for capturing user choices => numeric codes.
 input_mapping = {
     'YMDESUD5ANYO': {
         "SUD only, no MDE": 1, 
@@ -330,21 +326,24 @@ def predict(
     count_ones = np.sum(all_preds == 1)
     
     # Evaluate severity using count_ones
-    severity_msg = predictor.evaluate_severity(count_ones)
-
-    # Sum up the probability of class=2 across all labels
-    # (We stored probs[i][0] = probability for the first row => class=2)
+    severity_base = predictor.evaluate_severity(count_ones)
+
+    # -------------------------------
+    # Sum of predicted probabilities
+    # -------------------------------
+    # 'probs' is a list of arrays; each array is the prob for class=2 from each model.
+    # If classes_=[1,2], then prob[0] is P(class=2). Probability of class=1 is (1 - P(class=2)).
+    # We'll sum them:
     sum_prob_2 = sum(prob[0] for prob in probs if not np.isnan(prob[0]))
-    # Append to severity message
-    severity_msg = f"{severity_msg} (Sum of Probabilities for 'Class 2' = {sum_prob_2:.2f})"
+    sum_prob_1 = sum((1 - prob[0]) for prob in probs if not np.isnan(prob[0]))
+    severity_msg = f"{severity_base} (Sum of Prob=1={sum_prob_1:.2f}, Prob=2={sum_prob_2:.2f})"
 
     # 4) Summarize predictions (with probabilities)
-    # Build label -> (pred_value, prob_value)
     label_prediction_info = {}
     for i, fname in enumerate(model_filenames):
         lbl_col = fname.split('.')[0]
         pred_val = preds[i][0]     # e.g. 1 or 2
-        prob_val = probs[i][0]     # probability for class=2 (?)
+        prob_val = probs[i][0]     # probability for class=2
         label_prediction_info[lbl_col] = (pred_val, prob_val)
 
     # Group them by domain
@@ -379,7 +378,7 @@ def predict(
                     icon = "❌"  # red cross
 
                 if not np.isnan(prob_val):
-                    text_prob = f"(Prob= {prob_val:.2f})"
+                    text_prob = f"(Prob= {prob_val:.2f} for class=2)"
                 else:
                     text_prob = "(No probability available)"
 
@@ -426,14 +425,22 @@ def predict(
             "Count": list(label_counts.values()),
             "Pred_Val": [label_prediction_info[lbl_col][0] for lbl_col in label_counts.keys()]
         })
-        # Assign colors based on predicted value (2 => green, 1 => red)
-        bar_lbl_df["Color"] = bar_lbl_df["Pred_Val"].apply(lambda x: "Green" if x == 2 else "Red")
+        # Assign legend text & color based on predicted value
+        #  - 2 => "Ok Mental Status" (green)
+        #  - 1 => "Bad Mental Status" (red)
+        bar_lbl_df["Mental Status"] = bar_lbl_df["Pred_Val"].apply(
+            lambda x: "Ok Mental Status" if x == 2 else "Bad Mental Status"
+        )
+
         fig_lbl = px.bar(
             bar_lbl_df, 
             x="Label", 
             y="Count", 
-            color="Color",
-            color_discrete_map={"Green": "green", "Red": "red"},
+            color="Mental Status",
+            color_discrete_map={
+                "Ok Mental Status": "green",
+                "Bad Mental Status": "red"
+            },
             title="Number of Patients with the Same Predicted Label"
         )
         fig_lbl.update_layout(width=1200, height=400)
@@ -443,7 +450,7 @@ def predict(
 
     return (
         final_str,         # 1) Prediction Results
-        severity_msg,      # 2) Mental Health Severity + sum of probabilities
+        severity_msg,      # 2) Mental Health Severity (with sums of prob=1 & prob=2)
         total_count_md,    # 3) Total Patient Count
         nn_md,             # 4) Nearest Neighbors Summary
         fig_in,            # 5) Bar Chart (input features)