Clocksp/face-emotion-recognition

1import gradio as gr2import cv23import numpy as np4import pickle5from functools import lru_cache6 7try:8    from util import get_face_landmarks9except Exception as e:10    raise ImportError(11        "Make sure util.py defines get_face_landmarks(image, draw=False)."12    ) from e13 14 15# ---- App Config ----16EMOTIONS = ["HAPPY", "SAD", "SURPRISED"]17MODEL_PATH = "model.pkl"18APP_TITLE = "Emotion Detector"19APP_DESC = (20    "Upload an image or use your webcam. Toggle 'Draw Landmarks' for visualization."21)22 23 24# ---- Model Loader (cached) ----25@lru_cache(maxsize=1)26def load_model():27    with open(MODEL_PATH, "rb") as f:28        return pickle.load(f)29 30 31# ---- Core Inference ----32def predict_emotion(image, draw_toggle):33 34    if image is None:35        return {"Status": 1.0}, None, "Please upload an image."36 37    draw = (draw_toggle == "ON")38    img_rgb = np.array(image)39 40    if img_rgb.ndim == 2:41        img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_GRAY2RGB)42 43    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)44 45    landmarks = get_face_landmarks(img_bgr, draw=draw)46 47    if landmarks is None or len(landmarks) == 0:48        return {"No face detected": 1.0}, img_rgb, "No face detected."49 50    model = load_model()51 52    # Prediction53    pred_idx = int(model.predict([landmarks])[0])54    pred_label = EMOTIONS[pred_idx] if 0 <= pred_idx < len(EMOTIONS) else str(pred_idx)55 56    # Confidence57    if hasattr(model, "predict_proba"):58        probs = model.predict_proba([landmarks])[0]59        confidence = {EMOTIONS[i]: float(probs[i]) for i in range(len(EMOTIONS))}60    else:61        confidence = {pred_label: 1.0}62 63    # Output image64    img_out = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) if draw else img_rgb65 66    status = f"Detected emotion: {pred_label}"67    return confidence, img_out, status68 69with gr.Blocks() as demo:70    gr.Markdown(f"# {APP_TITLE}\n{APP_DESC}")71 72    with gr.Row():73        with gr.Column():74            image_input = gr.Image(type="pil", sources=["upload", "webcam"])75            draw_toggle = gr.Radio(["OFF", "ON"], value="OFF", label="Draw Landmarks")76 77        with gr.Column():78            label_output = gr.Label(num_top_classes=3)79            image_output = gr.Image(type="numpy")80            status_output = gr.Textbox()81 82    gr.Examples(83        examples=[84            ["examples/happy.png", "OFF"],85            ["examples/sad.png", "OFF"],86            ["examples/surprised.png", "OFF"],87        ],88        inputs=[image_input, draw_toggle],89    )90 91    image_input.change(92        predict_emotion,93        [image_input, draw_toggle],94        [label_output, image_output, status_output],95        queue=False,96    )97 98    draw_toggle.change(99        predict_emotion,100        [image_input, draw_toggle],101        [label_output, image_output, status_output],102        queue=False,103        )104 105if __name__ == "__main__":106    demo.launch()