Update app.py

app.py

@@ -1,96 +1,81 @@
 import streamlit as st
-import requests
-import json
-import os
+from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Tokenizer
+import torch
+import numpy as np
+import io
 
-# Get the API token from environment variables or secrets
-API_URL = "https://api-inference.huggingface.co/models/ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition"
-headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}
+# Load the model and tokenizer
+model_name = "harshit345/xlsr-wav2vec-speech-emotion-recognition"
+model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
+tokenizer = Wav2Vec2Tokenizer.from_pretrained(model_name)
 
-# Function to query the Hugging Face API
-def query_api(audio_file):
-    with open(audio_file, "rb") as f:
-        response = requests.post(API_URL, headers=headers, data=f)
-    try:
-        response_json = response.json()
-    except json.JSONDecodeError:
-        st.error("Error decoding JSON from API response.")
-        response_json = {}
-    return response_json
-
-# Function to compute the neg_rate from API response
-def compute_neg_rate(emotion_data):
-    if isinstance(emotion_data, list):  # Ensure the data is a list
-        neg_emotions = ["anger", "disgust", "fear", "sadness"]  # Example negative emotions
-        neg_rate = sum([emotion['score'] for emotion in emotion_data if emotion['label'].lower() in neg_emotions])
-        return neg_rate * 100  # Scale to a percentage
-    else:
-        st.error("Unexpected API response format.")
-        return 0
-
-# Depression scale logic
-def analyze_depression(neg_rate):
-    responses = []
+def analyze_emotion(audio_bytes):
+    audio_input = tokenizer(audio_bytes, return_tensors="pt", padding="longest")
+    with torch.no_grad():
+        logits = model(**audio_input).logits
+    prob = torch.sigmoid(logits).numpy().flatten()
+    neg_rate = prob[0] * 100  # Assuming the first element is the negative emotion probability
+    
     if neg_rate <= 30.0:
         if neg_rate <= 5.0:
-            responses.append("just chill and relax")
-        else:
-            responses.append("***   MILD DEPRESSION  ***")
-        if neg_rate <= 10.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: Everything is a-okay! There's absolutely nothing wrong 😁😁\n\t\tYou're probably cuddling a fluffy kitten right now")
-        elif neg_rate > 10.0 and neg_rate <= 20.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: You are a bit frustrated and disappointed\n\t\tBut you're easily distracted and cheered with little effort 😁😁")
-        elif neg_rate > 20.0 and neg_rate <= 30.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: Things are bothering you but you're coping up\n\t\tYou might be overtired and hungry\n\t\tThe emotional equivalent of a headache 😁😁")
-    elif neg_rate > 30.0 and neg_rate <= 60.0:
-        responses.append("***   MODERATE DEPRESSION  ***")
-        if neg_rate > 30.0 and neg_rate <= 40.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: Today is slightly a bad day for you.\n\t\tYou still have the skills to get through it, but be gentle with yourself\n\t\tUse self-care strategies 😁😁")
-        elif neg_rate > 40.0 and neg_rate <= 50.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: Your mental health is starting to impact your everyday life.\n\t\tEasy things are becoming difficult")
-        elif neg_rate > 50.0 and neg_rate <= 60.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: You are not able to do things the way you usually do them due to your mental health.\n\t\tImpulsive and compulsive thoughts might be difficult to cope with")
-    elif neg_rate > 60.0 and neg_rate <= 100.0:
-        responses.append("***   SEVERE DEPRESSION  ***")
-        if neg_rate > 60.0 and neg_rate <= 70.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: You are losing interest in activities that used to be enjoyable.\n\t\tYou should definitely seek help\n\t\tThis is becoming serious ☹☹")
-        elif neg_rate > 70.0 and neg_rate <= 80.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: You can't ignore your struggles now, It's HIGH time!.\n\t\tYou may have issues sleeping, eating, having fun, socializing, and work/study\n\t\tYour mental health is affecting almost all parts of your life ☹☹")
-        elif neg_rate > 80.0 and neg_rate <= 90.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: You are at a critical point!!!.\n\t\tYou aren't functioning anymore and need urgent help.\n\t\tYou may be a risk to yourself or others if left untreated ☹☹")
-        elif neg_rate > 90.0 and neg_rate <= 100.0:
-            responses.append(f"\nDepression Scale Detected: {round(neg_rate) / 10}\n")
-            responses.append("My Response: The worst mental and emotional distress possible.\n\t\tYou can't imagine things getting any better now and you might think it's all over for you (SUICIDE).\n\t\tContact a crisis line or get started with treatment immediately. ☹☹")
-    return "\n".join(responses)
-
-# Streamlit app
-def main():
-    st.title("Emotion Detection and Depression Scale Analysis")
-
-    audio_file = st.file_uploader("Upload an audio file", type=["wav"])
-    if audio_file:
-        st.audio(audio_file, format="audio/wav")
+            return "just chill and relax"
+        elif neg_rate <= 10.0:
+            return (f"***   MILD DEPRESSION  ***\n"
+                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tEverything is a-okay! There's absolutely nothing wrong 😁😁\n"
+                    f"\t\tYou're probably cuddling a fluffy kitten right now")
+        elif neg_rate <= 20.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYou are a bit frustrated and disappointed\n"
+                    f"\t\tBut you're easily distracted and cheered with little effort 😁😁")
+        elif neg_rate <= 30.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tThings are bothering you but you're coping up\n"
+                    f"\t\tYou might be over tired and hungry\n"
+                    f"\t\tThe emotional equivalent of a headache 😁😁")
+    elif neg_rate <= 60.0:
+        if neg_rate <= 40.0:
+            return (f"***   MODERATE DEPRESSION  ***\n"
+                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tToday is slightly bad for you.\n"
+                    f"\t\tYou still have the skills to get through it, but be gentle with yourself\n"
+                    f"\t\tUse self-care strategies 😁😁")
+        elif neg_rate <= 50.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYour mental health is starting to impact your everyday life.\n"
+                    f"\t\tEasy things are becoming difficult")
+        elif neg_rate <= 60.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYou are not able to do things the way usually you do them due to your mental health.\n"
+                    f"\t\tImpulsive and compulsive thoughts might be difficult to cope with")
+    elif neg_rate <= 100.0:
+        if neg_rate <= 70.0:
+            return (f"***   SEVERE DEPRESSION  ***\n"
+                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYou are losing interest in the activities that used to be enjoyable.\n"
+                    f"\t\tYou should definitely seek help\n"
+                    f"\t\tThis is becoming serious ☹☹")
+        elif neg_rate <= 80.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYou can't ignore your struggles now, It's HIGH time!.\n"
+                    f"\t\tYou may have issues sleeping, eating, having fun, socializing, and work/study\n"
+                    f"\t\tYour mental health is affecting almost all parts of your life ☹☹")
+        elif neg_rate <= 90.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tYou are at a critical point !!!\n"
+                    f"\t\tYou aren't functioning anymore and need urgent help.\n"
+                    f"\t\tYou may be a risk to yourself or others if left untreated ☹☹")
+        elif neg_rate <= 100.0:
+            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
+                    f"My Response :\tThe worst mental and emotional distress possible.\n"
+                    f"\t\tYou can't imagine things getting any better now and you might think it's all over for you (SUICIDE).\n"
+                    f"\t\tContact crisis line or get started for the treatment immediately. ☹☹")
 
-        if st.button("Analyze"):
-            # Save uploaded file to a temporary file
-            with open("temp_audio.wav", "wb") as f:
-                f.write(audio_file.read())
+st.title("Emotion Detection App")
 
-            # Query API and compute neg_rate
-            emotion_data = query_api("temp_audio.wav")
-            st.write("API Response:", emotion_data)  # Debugging line to inspect response
-            neg_rate = compute_neg_rate(emotion_data)
-            result = analyze_depression(neg_rate)
-            st.text(result)
+uploaded_file = st.file_uploader("Choose an audio file", type=["wav"])
 
-if __name__ == "__main__":
-    main()
+if uploaded_file is not None:
+    audio_bytes = uploaded_file.read()
+    result = analyze_emotion(audio_bytes)
+    st.write(result)