Update app.py

app.py

@@ -57,11 +57,20 @@ def predict(amount, user_txn, driver_txn, pair_count, hour, day):
 
     outputs = sess.run(None, {sess.get_inputs()[0].name: X})
 
-    # Case 1: probability output exists (most common)
-    if len(outputs[0].shape) == 2:
-        fraud_prob = outputs[0][0][1]   # class-1 probability
+    scores = outputs[0]
+
+    # Case 1: class probabilities [prob_0, prob_1]
+    if scores.ndim == 2 and scores.shape[1] == 2:
+        fraud_prob = scores[0][1]
+
+    # Case 2: single probability
+    elif scores.ndim == 2 and scores.shape[1] == 1:
+        fraud_prob = scores[0][0]
+
+    # Case 3: raw score → sigmoid
     else:
-        # Case 2: raw score → apply sigmoid
-        fraud_prob = 1 / (1 + np.exp(-outputs[0][0]))
+        raw_score = scores[0]
+        fraud_prob = 1 / (1 + np.exp(-raw_score))
 
     return float(fraud_prob)
+