Update app.py

app.py

@@ -2,61 +2,164 @@ import gradio as gr
 import cv2
 import numpy as np
 import pickle
-from util import get_face_landmarks
-
-# Emotion labels
-emotions = ['HAPPY', 'SAD', 'SURPRISED']
-
-# Load the model once
-with open('model.pkl', 'rb') as f:
-    model = pickle.load(f)
-
-def predict_emotion(image):
-    # Convert PIL image to OpenCV format
-    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-
-    # Extract landmarks
-    face_landmarks = get_face_landmarks(img, draw=False, static_image_mode=True)
-
-    if face_landmarks is None or len(face_landmarks) == 0:
-        # Draw "No Face Detected" on the image
-        cv2.putText(img, "No Face Detected", (10, img.shape[0] - 10),
-                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
-        return {"No face detected": 1.0}, cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "⚠ No face detected in the image."
-
-    # Predict emotion and probabilities
-    output = model.predict([face_landmarks])
-    predicted_emotion = emotions[int(output[0])]
-
-    if hasattr(model, "predict_proba"):
-        probs = model.predict_proba([face_landmarks])[0]
-        confidence_dict = {emotions[i]: float(probs[i]) for i in range(len(emotions))}
-    else:
-        confidence_dict = {predicted_emotion: 1.0}
-
-    # Annotate image with predicted emotion
-    cv2.putText(img, predicted_emotion, (10, img.shape[0] - 10),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
-    return confidence_dict, cv2.cvtColor(img, cv2.COLOR_BGR2RGB), f" Detected emotion: {predicted_emotion}"
-
-# Example images
-examples = [["examples/happy.png"], ["examples/sad.png"], ["examples/surprised.png"]]
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=predict_emotion,
-    inputs=gr.Image(type="pil", label="Upload Image or Use Webcam", sources=["upload", "webcam"]),
-    outputs=[
-        gr.Label(num_top_classes=3, label="Predicted Emotion & Confidence"),
-        gr.Image(type="numpy", label="Annotated Image"),
-        gr.Textbox(label="Status", interactive=False)
-    ],
-    title="Emotion Detector",
-    description="Upload an image or use webcam to detect emotions (HAPPY, SAD, SURPRISED).",
-    examples=examples,
-    theme="default"
+from functools import lru_cache
+
+# ---- Import your landmarks util ----
+try:
+    from util import get_face_landmarks
+except Exception as e:
+    raise ImportError(
+        "Could not import 'get_face_landmarks' from util.py. "
+        "Make sure util.py exists and defines get_face_landmarks(img, draw: bool, static_image_mode: bool)."
+    ) from e
+
+
+# ---- App Config ----
+EMOTIONS = ["HAPPY", "SAD", "SURPRISED"]
+MODEL_PATH = "model.pkl"
+APP_TITLE = "Emotion Detector"
+APP_DESC = (
+    "Upload an image or use your webcam. Toggle 'Draw Landmarks' for visualization. "
+    "If no face is detected, the original image is shown and status explains why."
 )
 
+
+# ---- Model Loader (cached) ----
+@lru_cache(maxsize=1)
+def load_model():
+    with open(MODEL_PATH, "rb") as f:
+        model = pickle.load(f)
+    return model
+
+
+# ---- Core Inference ----
+def predict_emotion(image, draw_toggle):
+    """
+    image: PIL.Image (from gr.Image with type='pil')
+    draw_toggle: 'OFF' or 'ON'
+    """
+    # Input validation
+    if image is None:
+        return {"Status": 1.0}, None, "Please upload an image."
+
+    draw = (draw_toggle == "ON")
+
+    # Convert PIL -> OpenCV BGR
+    try:
+        img_rgb = np.array(image)  # PIL -> RGB ndarray
+        if img_rgb.ndim == 2:  # grayscale to 3-ch
+            img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_GRAY2RGB)
+        img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+    except Exception:
+        return {"Status": 1.0}, None, "⚠ Could not read the image. Please try a different one."
+
+    # Extract landmarks (your util may also draw on img internally when draw=True)
+    try:
+        landmarks = get_face_landmarks(img_bgr, draw=draw, static_image_mode=True)
+    except Exception as e:
+        return {"Status": 1.0}, img_rgb, f"⚠ Landmark extraction failed: {e}"
+
+    # Handle no-face case (do NOT annotate; return original image)
+    if landmarks is None or (hasattr(landmarks, "__len__") and len(landmarks) == 0):
+        return {"No face detected": 1.0}, img_rgb, "⚠ No face detected in the image."
+
+    # Load model
+    try:
+        model = load_model()
+    except FileNotFoundError:
+        return {"Status": 1.0}, img_rgb, "⚠ model.pkl not found in repo root."
+    except Exception as e:
+        return {"Status": 1.0}, img_rgb, f"⚠ Failed to load model: {e}"
+
+    # Predict
+    try:
+        output = model.predict([landmarks])
+        pred_idx = int(output[0])
+        pred_label = EMOTIONS[pred_idx] if 0 <= pred_idx < len(EMOTIONS) else str(pred_idx)
+
+        # Confidence/probabilities if available
+        if hasattr(model, "predict_proba"):
+            probs = model.predict_proba([landmarks])[0]
+            confidence = {EMOTIONS[i]: float(probs[i]) for i in range(len(EMOTIONS))}
+        else:
+            confidence = {pred_label: 1.0}
+
+        # Always draw predicted text on the copy we return (but ONLY if a face exists)
+        img_annot = img_bgr.copy()
+        try:
+            cv2.putText(
+                img_annot,
+                pred_label,
+                (10, img_annot.shape[0] - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1.0,
+                (0, 255, 0),
+                2,
+                cv2.LINE_AA,
+            )
+        except Exception:
+            # If drawing fails, just return the original image
+            img_annot = img_bgr
+
+        img_out = cv2.cvtColor(img_annot, cv2.COLOR_BGR2RGB)
+        status = f"✅ Detected emotion: {pred_label}"
+        return confidence, img_out, status
+
+    except Exception as e:
+        return {"Status": 1.0}, img_rgb, f"⚠ Inference failed: {e}"
+
+
+# ---- Gradio UI ----
+with gr.Blocks(theme="default") as demo:
+    gr.Markdown(f"# {APP_TITLE}\n{APP_DESC}")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(
+                type="pil",
+                label="Upload Image or Use Webcam",
+                sources=["upload", "webcam"],
+                interactive=True,
+            )
+            draw_toggle = gr.Radio(
+                choices=["OFF", "ON"],
+                value="OFF",
+                label="Draw Landmarks",
+                interactive=True,
+            )
+            gr.Markdown(
+                "Tip: Switch **Draw Landmarks** ON to visualize key points (if your `util.get_face_landmarks` draws them)."
+            )
+
+        with gr.Column(scale=1):
+            label_output = gr.Label(num_top_classes=3, label="Predicted Emotion & Confidence")
+            image_output = gr.Image(type="numpy", label="Annotated Image")
+            status_output = gr.Textbox(label="Status", interactive=False)
+
+    # Examples (ensure these files exist in /examples)
+    gr.Examples(
+        examples=[
+            ["examples/happy.png", "OFF"],
+            ["examples/sad.png", "OFF"],
+            ["examples/surprised.png", "OFF"],
+        ],
+        inputs=[image_input, draw_toggle],
+        label="Try examples",
+    )
+
+    # Real-time: changing either the image or the toggle re-runs inference automatically
+    image_input.change(
+        fn=predict_emotion,
+        inputs=[image_input, draw_toggle],
+        outputs=[label_output, image_output, status_output],
+        queue=False,
+    )
+    draw_toggle.change(
+        fn=predict_emotion,
+        inputs=[image_input, draw_toggle],
+        outputs=[label_output, image_output, status_output],
+        queue=False,
+    )
+
 if __name__ == "__main__":
     demo.launch()