Update app.py

app.py

@@ -1,16 +1,6 @@
 # =========================================================
-# HUGGINGFACE DRIVER STATE DETECTION APP
-# =========================================================
-#
-# Upload driver image
-# → Predict fatigue state
-#
-# Classes:
-# - alert
-# - sleepy
-# - slowBlink
-# - yawning
-#
+# AI DRIVER SAFETY DETECTION SYSTEM
+# HuggingFace Gradio App
 # =========================================================
 
 import gradio as gr
@@ -26,6 +16,10 @@ from tensorflow.keras.preprocessing.image import img_to_array
 
 model = load_model("final_driver_state_model.h5")
 
+# =========================================================
+# CLASS NAMES
+# =========================================================
+
 CLASS_NAMES = [
     "alert",
     "sleepy",
@@ -34,7 +28,7 @@ CLASS_NAMES = [
 ]
 
 # =========================================================
-# RISK MAPPING
+# RISK LEVELS
 # =========================================================
 
 RISK_LEVELS = {
@@ -44,19 +38,38 @@ RISK_LEVELS = {
     "yawning": "LOW RISK"
 }
 
+# =========================================================
+# COLORS
+# =========================================================
+
+RISK_EMOJIS = {
+    "SAFE": "🟢",
+    "LOW RISK": "🟡",
+    "MEDIUM RISK": "🟠",
+    "HIGH RISK": "🔴"
+}
+
 # =========================================================
 # PREDICTION FUNCTION
 # =========================================================
 
 def predict_driver_state(image):
 
-    # =============================================
-    # IMAGE PREPROCESSING
-    # =============================================
+    # =====================================================
+    # VALIDATION
+    # =====================================================
 
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    if image is None:
+        return "Please upload an image.", None
 
-    image = cv2.resize(image, (224,224))
+    # =====================================================
+    # PREPROCESSING
+    # IMPORTANT:
+    # Gradio already gives RGB image
+    # DO NOT use cvtColor here
+    # =====================================================
+
+    image = cv2.resize(image, (224, 224))
 
     image = image.astype("float32") / 255.0
 
@@ -64,9 +77,9 @@ def predict_driver_state(image):
 
     image = np.expand_dims(image, axis=0)
 
-    # =============================================
-    # MODEL PREDICTION
-    # =============================================
+    # =====================================================
+    # PREDICTION
+    # =====================================================
 
     prediction = model.predict(image, verbose=0)
 
@@ -78,9 +91,11 @@ def predict_driver_state(image):
 
     risk_level = RISK_LEVELS[predicted_class]
 
-    # =============================================
-    # ALL CLASS PROBABILITIES
-    # =============================================
+    emoji = RISK_EMOJIS[risk_level]
+
+    # =====================================================
+    # CONFIDENCE SCORES
+    # =====================================================
 
     confidence_scores = {}
 
@@ -90,40 +105,50 @@ def predict_driver_state(image):
             prediction[0][i]
         )
 
-    # =============================================
+    # =====================================================
     # RESULT TEXT
-    # =============================================
+    # =====================================================
 
     result = f"""
 🚗 DRIVER STATE ANALYSIS
 
-Prediction: {predicted_class}
+Prediction:
+{predicted_class.upper()}
 
-Confidence: {confidence:.2f}
+Confidence:
+{confidence:.2f}
 
-Risk Level: {risk_level}
+Risk Level:
+{emoji} {risk_level}
 """
 
     return result, confidence_scores
 
 # =========================================================
-# GRADIO UI
+# EXAMPLES
+# =========================================================
+
+examples = []
+
+# =========================================================
+# UI
 # =========================================================
 
 title = "🚗 AI Driver Safety Detection System"
 
 description = """
-Upload a driver image to detect fatigue state using Deep Learning.
+Upload a driver image to analyze fatigue and attention state using Deep Learning.
 
-### Supported Driver States
-- Alert
-- Sleepy
-- Slow Blink
-- Yawning
+## Supported Driver States
+- 🟢 Alert
+- 🔴 Sleepy
+- 🟠 Slow Blink
+- 🟡 Yawning
 
-### AI Features
-✅ CNN-Based Classification  
+## AI Features
+✅ CNN-Based Driver State Classification  
 ✅ Fatigue Risk Analysis  
+✅ Deep Learning Inference  
 ✅ Real-Time Prediction Engine  
 """
 
@@ -141,12 +166,13 @@ interface = gr.Interface(
     ),
 
     outputs=[
+
         gr.Textbox(
             label="Prediction Result"
         ),
 
         gr.Label(
-            label="Class Confidence Scores"
+            label="Confidence Scores"
         )
     ],
 
@@ -154,7 +180,11 @@ interface = gr.Interface(
 
     description=description,
 
-    theme="soft"
+    examples=examples,
+
+    theme="soft",
+
+    allow_flagging="never"
 )
 
 # =========================================================