Update app.py

app.py

@@ -1,48 +1,85 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForImageClassification, AutoImageProcessor
-from PIL import Image
 import numpy as np
+from PIL import Image
 import os
+import warnings
+
+warnings.filterwarnings("ignore")
+
+# --- Optional torch and tf loading ---
+try:
+    import torch
+    IS_TORCH_AVAILABLE = True
+except ImportError:
+    IS_TORCH_AVAILABLE = False
 
-# Check for TensorFlow
 try:
     import tensorflow as tf
     IS_TF_AVAILABLE = True
 except ImportError:
     IS_TF_AVAILABLE = False
 
-# --- Load Model ---
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+
+# --- Model loading ---
 model = None
 processor = None
 is_pytorch_model = True
-
 model_name_or_path = "ravi86/mood_detector"
-local_h5_path = "./my_model.h5"
-
-try:
-    model = AutoModelForImageClassification.from_pretrained(model_name_or_path)
-    processor = AutoImageProcessor.from_pretrained(model_name_or_path)
-except:
-    if IS_TF_AVAILABLE:
+local_model_dir = "./model"
+local_h5_path = os.path.join(local_model_dir, "my_model.h5")
+
+# Try Hugging Face PyTorch
+if IS_TORCH_AVAILABLE:
+    try:
+        model = AutoModelForImageClassification.from_pretrained(model_name_or_path)
+        processor = AutoImageProcessor.from_pretrained(model_name_or_path)
+        is_pytorch_model = True
+        print("Loaded PyTorch model from Hugging Face.")
+    except:
+        pass
+
+# Try Hugging Face TensorFlow
+if model is None and IS_TF_AVAILABLE:
+    try:
+        model = AutoModelForImageClassification.from_pretrained(model_name_or_path, from_tf=True)
+        processor = AutoImageProcessor.from_pretrained(model_name_or_path)
+        is_pytorch_model = False
+        print("Loaded TensorFlow model from Hugging Face.")
+    except:
+        pass
+
+# Try local Transformers model
+if model is None:
+    try:
+        model = AutoModelForImageClassification.from_pretrained(local_model_dir)
+        processor = AutoImageProcessor.from_pretrained(local_model_dir)
+        is_pytorch_model = hasattr(model, 'parameters')
+        print("Loaded local Transformers model.")
+    except:
+        pass
+
+# Try raw .h5
+if model is None and IS_TF_AVAILABLE and os.path.exists(local_h5_path):
+    try:
+        model = tf.keras.models.load_model(local_h5_path)
         try:
-            model = tf.keras.models.load_model(local_h5_path)
-            processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
-            is_pytorch_model = False
+            processor = AutoImageProcessor.from_pretrained(local_model_dir)
         except:
-            raise RuntimeError("Could not load .h5 model.")
-    else:
-        raise RuntimeError("Model loading failed. No valid model found.")
+            processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+        is_pytorch_model = False
+        print("Loaded local Keras .h5 model.")
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {e}")
 
 if model is None or processor is None:
-    raise RuntimeError("Model or processor not loaded.")
+    raise RuntimeError("Failed to load model and processor.")
 
-if is_pytorch_model:
+if is_pytorch_model and IS_TORCH_AVAILABLE:
     model.eval()
 
-# --- Labels and Music Mapping ---
+# --- Emotion & Spotify Map ---
 emotions = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
-
 spotify_playlist_mapping = {
     "Angry": "https://open.spotify.com/playlist/37i9dQZF1DX2LTjeP1y0aR",
     "Disgust": "https://open.spotify.com/playlist/37i9dQZF1DXcK3k3gJ6usM",
@@ -53,56 +90,55 @@ spotify_playlist_mapping = {
     "Neutral": "https://open.spotify.com/playlist/37i9dQZF1DXasMvN3R0sVw"
 }
 
-# --- Predict Function ---
-def classify_expression_and_suggest_music(image_input):
+# --- Inference Function ---
+def classify_expression_and_suggest_music(image_input: np.ndarray):
     if image_input is None:
-        return "No webcam input detected.", ""
+        return "No image detected. Please enable your webcam.", ""
 
     image = Image.fromarray(image_input).convert("L").resize((48, 48))
+    inputs = processor(images=image, return_tensors="pt" if is_pytorch_model else "tf")
 
-    inputs = processor(images=image, return_tensors="pt")
-    inputs_for_model = inputs['pixel_values']
-
-    if not is_pytorch_model and IS_TF_AVAILABLE:
-        pixel_values_np = inputs['pixel_values'].squeeze(0).numpy()
-        inputs_for_model = tf.expand_dims(tf.convert_to_tensor(pixel_values_np), 0)
-
-    with torch.no_grad():
-        if is_pytorch_model:
-            outputs = model(inputs_for_model)
-            logits = outputs.logits
-        else:
-            outputs = model(inputs_for_model)
-            logits = outputs
-
+    if not is_pytorch_model:
+        pixel_values = inputs['pixel_values'].numpy()
+        tf_tensor = tf.convert_to_tensor(pixel_values)
+        outputs = model(tf_tensor)
+        logits = outputs if isinstance(outputs, (np.ndarray, tf.Tensor)) else outputs[0]
         if isinstance(logits, tf.Tensor):
-            logits = torch.from_numpy(logits.numpy())
-        elif not isinstance(logits, torch.Tensor):
-            logits = torch.from_numpy(np.array(logits))
-
-        probs = torch.softmax(logits, dim=-1)
-
-    idx = probs.argmax().item()
-    emotion = emotions[idx]
-    confidence = probs[0, idx].item() * 100
+            logits = logits.numpy()
+        probs = tf.nn.softmax(logits).numpy()
+    else:
+        with torch.no_grad():
+            outputs = model(inputs['pixel_values'])
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=-1).numpy()
 
-    output_text = f"Detected Emotion: **{emotion}** (Confidence: {confidence:.2f}%)"
-    playlist_url = spotify_playlist_mapping.get(emotion, spotify_playlist_mapping["Neutral"])
-    spotify_link = f"**Listen on Spotify:** <a href='{playlist_url}' target='_blank'>🎧 {emotion} Vibes</a>"
+    predicted_class = int(np.argmax(probs))
+    confidence = float(np.max(probs)) * 100
+    emotion = emotions[predicted_class]
 
-    return output_text, spotify_link
+    spotify_link = spotify_playlist_mapping.get(emotion, spotify_playlist_mapping["Neutral"])
+    return (
+        f"Detected Emotion: **{emotion}** (Confidence: {confidence:.2f}%)",
+        f"**Listen on Spotify:** <a href='{spotify_link}' target='_blank'>🎧 {emotion} Vibes</a>"
+    )
 
-# --- Gradio UI ---
+# --- Gradio Interface ---
 iface = gr.Interface(
     fn=classify_expression_and_suggest_music,
-    inputs=gr.Image(type="numpy", source="webcam", streaming=True, label="Live Webcam"),
+    inputs=gr.Image(
+        type="numpy",
+        source="webcam",
+        streaming=True,
+        label="Webcam Input"
+    ),
     outputs=[
-        gr.Textbox(label="Emotion Detected"),
-        gr.Markdown(label="Suggested Music")
+        gr.Textbox(label="Detected Emotion"),
+        gr.Markdown(label="Suggested Spotify Playlist")
     ],
     live=True,
-    title="🎭 MoodTune: Your Emotional DJ 🎶",
-    description="Real-time facial expression detector that plays music to match your mood!",
+    title="🎭 MoodTune: Emotion-Based Music Recommender",
+    description="This app detects your mood from your face and plays music to match it! Allow webcam access to begin."
 )
 
-iface.launch()  # Automatically runs in Hugging Face Spaces
+if __name__ == "__main__":
+    iface.launch()