Update app.py

app.py

@@ -17,7 +17,9 @@ from tensorflow.keras.preprocessing.image import img_to_array
 model = load_model("final_driver_state_model.h5")
 
 # =========================================================
-# CLASS NAMES
+# CLASS LABELS
+# IMPORTANT:
+# Must match training class order exactly
 # =========================================================
 
 CLASS_NAMES = [
@@ -39,7 +41,7 @@ RISK_LEVELS = {
 }
 
 # =========================================================
-# COLORS
+# EMOJIS
 # =========================================================
 
 RISK_EMOJIS = {
@@ -50,24 +52,17 @@ RISK_EMOJIS = {
 }
 
 # =========================================================
-# PREDICTION FUNCTION
+# IMAGE PREPROCESSING
+# IMPORTANT:
+# Match training preprocessing
 # =========================================================
 
-def predict_driver_state(image):
-
-    # =====================================================
-    # VALIDATION
-    # =====================================================
+def preprocess_image(image):
 
-    if image is None:
-        return "Please upload an image.", None
-
-    # =====================================================
-    # PREPROCESSING
-    # IMPORTANT:
-    # Gradio already gives RGB image
-    # DO NOT use cvtColor here
-    # =====================================================
+    # -----------------------------------------------------
+    # Gradio already provides RGB image
+    # DO NOT use cvtColor
+    # -----------------------------------------------------
 
     image = cv2.resize(image, (224, 224))
 
@@ -77,13 +72,41 @@ def predict_driver_state(image):
 
     image = np.expand_dims(image, axis=0)
 
+    return image
+
+# =========================================================
+# PREDICTION FUNCTION
+# =========================================================
+
+def predict_driver_state(image):
+
+    if image is None:
+
+        return (
+            "Please upload an image.",
+            {}
+        )
+
+    # =====================================================
+    # PREPROCESS
+    # =====================================================
+
+    processed_image = preprocess_image(image)
+
     # =====================================================
     # PREDICTION
     # =====================================================
 
-    prediction = model.predict(image, verbose=0)
+    prediction = model.predict(
+        processed_image,
+        verbose=0
+    )
 
-    class_index = np.argmax(prediction)
+    # =====================================================
+    # RESULTS
+    # =====================================================
+
+    class_index = int(np.argmax(prediction))
 
     predicted_class = CLASS_NAMES[class_index]
 
@@ -125,13 +148,7 @@ Risk Level:
     return result, confidence_scores
 
 # =========================================================
-# EXAMPLES
-# =========================================================
-
-examples = []
-
-# =========================================================
-# UI
+# TITLE & DESCRIPTION
 # =========================================================
 
 title = "🚗 AI Driver Safety Detection System"
@@ -153,7 +170,7 @@ Upload a driver image to analyze fatigue and attention state using Deep Learning
 """
 
 # =========================================================
-# INTERFACE
+# GRADIO INTERFACE
 # =========================================================
 
 interface = gr.Interface(
@@ -178,13 +195,7 @@ interface = gr.Interface(
 
     title=title,
 
-    description=description,
-
-    examples=examples,
-
-    theme="soft",
-
-    allow_flagging="never"
+    description=description
 )
 
 # =========================================================