remaking src folder

app.py → src/app.py

@@ -1,213 +1,213 @@
-import streamlit as st
-from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
-import av
-import numpy as np
-import torch
-import whisper
-from transformers import pipeline, AutoModelForAudioClassification, AutoFeatureExtractor
-from deepface import DeepFace
-import logging
-import queue
-import soundfile as sf
-from scipy.io.wavfile import write as write_wav
-import os
-import tempfile
-
-# --- Page Configuration ---
-st.set_page_config(
-    page_title="AffectLink Online Demo",
-    page_icon="😊",
-    layout="wide"
-)
-
-st.title("AffectLink: Real-time Emotion Recognition")
-st.write("This demo analyzes your facial expressions in real-time and processes short audio clips for speech and text-based emotion.")
-
-# --- Logger Configuration ---
-# Suppress noisy logs from libraries
-logging.basicConfig(level=logging.INFO)
-logging.getLogger('deepface').setLevel(logging.ERROR)
-logging.getLogger('huggingface_hub').setLevel(logging.WARNING)
-
-# --- Emotion Mappings ---
-UNIFIED_EMOTIONS = ['neutral', 'happy', 'sad', 'angry']
-FACIAL_TO_UNIFIED = {
-    'neutral': 'neutral', 'happy': 'happy', 'sad': 'sad', 'angry': 'angry',
-    'fear': None, 'surprise': None, 'disgust': None
-}
-TEXT_TO_UNIFIED = {
-    'neutral': 'neutral', 'joy': 'happy', 'sadness': 'sad', 'anger': 'angry',
-    'fear': None, 'surprise': None, 'disgust': None
-}
-SER_TO_UNIFIED = {
-    'neu': 'neutral', 'hap': 'happy', 'sad': 'sad', 'ang': 'angry'
-}
-AUDIO_SAMPLE_RATE = 16000
-
-# --- Model Loading ---
-# Use Streamlit's caching to load models only once.
-@st.cache_resource
-def load_models():
-    with st.spinner("Loading AI models... This may take a moment on first run."):
-        # Whisper for Speech-to-Text
-        whisper_model = whisper.load_model("base")
-
-        # Text Emotion Classifier
-        text_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None)
-
-        # Speech Emotion Recognition (SER)
-        ser_model_name = "superb/hubert-large-superb-er"
-        ser_feature_extractor = AutoFeatureExtractor.from_pretrained(ser_model_name)
-        ser_model = AutoModelForAudioClassification.from_pretrained(ser_model_name)
-
-        # NEW: Pre-load the DeepFace model to prevent lag on first use
-        DeepFace.build_model("Emotion")
-
-        return whisper_model, text_classifier, ser_model, ser_feature_extractor
-
-whisper_model, text_classifier, ser_model, ser_feature_extractor = load_models()
-
-# --- WebRTC and Video Processing ---
-webrtc_ctx = webrtc_streamer(
-    key="affectlink-video",
-    mode=WebRtcMode.SENDRECV,
-    rtc_configuration=RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}),
-    media_stream_constraints={"video": True, "audio": False},
-    async_processing=True,
-)
-
-if 'facial_emotion' not in st.session_state:
-    st.session_state.facial_emotion = "Neutral"
-if 'last_emotion_time' not in st.session_state:
-    st.session_state.last_emotion_time = 0
-
-def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
-    img = frame.to_ndarray(format="bgr24")
-    
-    # Process a frame every 5 seconds to conserve CPU
-    current_time = st.session_state.get('last_emotion_time', 0)
-    if torch.cuda.is_available() or (hasattr(st.session_state, 'last_emotion_time') and (torch.tensor(current_time).item() + 5 < torch.tensor(frame.time).item())):
-        try:
-            analysis = DeepFace.analyze(
-                img,
-                actions=['emotion'],
-                enforce_detection=False,
-                silent=True
-            )
-            if isinstance(analysis, list) and len(analysis) > 0:
-                dominant_emotion = analysis[0]['dominant_emotion']
-                st.session_state.facial_emotion = dominant_emotion.capitalize()
-            else:
-                st.session_state.facial_emotion = "Unknown"
-        except Exception as e:
-            logging.error(f"DeepFace analysis failed: {e}")
-            st.session_state.facial_emotion = "Error"
-        
-        st.session_state.last_emotion_time = frame.time
-
-    return av.VideoFrame.from_ndarray(img, format="bgr24")
-
-if webrtc_ctx.video_processor:
-    webrtc_ctx.video_processor.video_frame_callback = video_frame_callback
-
-# --- Audio Processing ---
-if "audio_buffer" not in st.session_state:
-    st.session_state.audio_buffer = []
-
-def audio_frame_callback(frame: av.AudioFrame):
-    sound = np.frombuffer(frame.to_ndarray(), dtype=np.int16)
-    st.session_state.audio_buffer.append(sound)
-
-webrtc_streamer(
-    key="affectlink-audio",
-    mode=WebRtcMode.RECVONLY,
-    media_stream_constraints={"video": False, "audio": True},
-    audio_frame_callback=audio_frame_callback,
-    async_processing=True,
-)
-
-# --- UI Layout ---
-st.sidebar.header("Facial Emotion")
-st.sidebar.metric("Current Expression", st.session_state.get('facial_emotion', 'N/A'))
-st.sidebar.info("Facial emotion is updated every 5 seconds to optimize performance.")
-st.sidebar.divider()
-
-st.sidebar.header("Audio Analysis")
-is_recording = st.sidebar.checkbox("Start Recording Audio")
-
-# NEW: Set up placeholders for audio results
-st.sidebar.subheader("Transcription:")
-transcription_placeholder = st.sidebar.empty()
-transcription_placeholder.write("_Waiting for audio..._")
-
-st.sidebar.subheader("Text Emotion:")
-text_emotion_placeholder = st.sidebar.empty()
-text_emotion_placeholder.write("_Waiting for audio..._")
-
-st.sidebar.subheader("Speech Emotion:")
-ser_placeholder = st.sidebar.empty()
-ser_placeholder.write("_Waiting for audio..._")
-
-
-if not is_recording and st.session_state.audio_buffer:
-    # Combine audio chunks
-    audio_data = np.concatenate(st.session_state.audio_buffer)
-    st.session_state.audio_buffer = []  # Clear buffer
-
-    # Save to a temporary file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio_file:
-        write_wav(tmp_audio_file.name, AUDIO_SAMPLE_RATE, audio_data)
-        
-        # 1. Speech-to-Text (Whisper)
-        with st.spinner("Transcribing audio..."):
-            try:
-                result = whisper_model.transcribe(tmp_audio_file.name, fp16=False)
-                transcribed_text = result['text']
-            except Exception as e:
-                transcribed_text = f"Transcription failed: {e}"
-            transcription_placeholder.write(f'"{transcribed_text}"')
-
-        # 2. Text-based Emotion
-        with st.spinner("Analyzing text emotion..."):
-            if transcribed_text:
-                try:
-                    text_emotions = text_classifier(transcribed_text)[0]
-                    unified_text_scores = {e: 0.0 for e in UNIFIED_EMOTIONS}
-                    for EMO in text_emotions:
-                        unified_emo = TEXT_TO_UNIFIED.get(EMO['label'])
-                        if unified_emo:
-                            unified_text_scores[unified_emo] += EMO['score']
-                    
-                    dominant_text_emotion = max(unified_text_scores, key=unified_text_scores.get)
-                except Exception as e:
-                    dominant_text_emotion = f"Text analysis failed: {e}"
-            else:
-                dominant_text_emotion = "No text to analyze."
-            text_emotion_placeholder.write(dominant_text_emotion.capitalize())
-
-        # 3. Speech Emotion Recognition (SER)
-        with st.spinner("Analyzing speech emotion..."):
-            try:
-                audio_array, _ = sf.read(tmp_audio_file.name)
-                inputs = ser_feature_extractor(audio_array, sampling_rate=AUDIO_SAMPLE_RATE, return_tensors="pt", padding=True)
-                with torch.no_grad():
-                    logits = ser_model(**inputs).logits
-                scores = torch.nn.functional.softmax(logits, dim=1).squeeze()
-                
-                unified_ser_scores = {e: 0.0 for e in UNIFIED_EMOTIONS}
-                for i, score in enumerate(scores):
-                    raw_emo = ser_model.config.id2label[i]
-                    unified_emo = SER_TO_UNIFIED.get(raw_emo)
-                    if unified_emo:
-                        unified_ser_scores[unified_emo] += score.item()
-                
-                dominant_ser_emotion = max(unified_ser_scores, key=unified_ser_scores.get)
-            except Exception as e:
-                dominant_ser_emotion = f"Speech analysis failed: {e}"
-            ser_placeholder.write(dominant_ser_emotion.capitalize())
-
-        # Clean up the temporary file
-        os.unlink(tmp_audio_file.name)
-
-elif is_recording:
+import streamlit as st
+from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
+import av
+import numpy as np
+import torch
+import whisper
+from transformers import pipeline, AutoModelForAudioClassification, AutoFeatureExtractor
+from deepface import DeepFace
+import logging
+import queue
+import soundfile as sf
+from scipy.io.wavfile import write as write_wav
+import os
+import tempfile
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="AffectLink Online Demo",
+    page_icon="😊",
+    layout="wide"
+)
+
+st.title("AffectLink: Real-time Emotion Recognition")
+st.write("This demo analyzes your facial expressions in real-time and processes short audio clips for speech and text-based emotion.")
+
+# --- Logger Configuration ---
+# Suppress noisy logs from libraries
+logging.basicConfig(level=logging.INFO)
+logging.getLogger('deepface').setLevel(logging.ERROR)
+logging.getLogger('huggingface_hub').setLevel(logging.WARNING)
+
+# --- Emotion Mappings ---
+UNIFIED_EMOTIONS = ['neutral', 'happy', 'sad', 'angry']
+FACIAL_TO_UNIFIED = {
+    'neutral': 'neutral', 'happy': 'happy', 'sad': 'sad', 'angry': 'angry',
+    'fear': None, 'surprise': None, 'disgust': None
+}
+TEXT_TO_UNIFIED = {
+    'neutral': 'neutral', 'joy': 'happy', 'sadness': 'sad', 'anger': 'angry',
+    'fear': None, 'surprise': None, 'disgust': None
+}
+SER_TO_UNIFIED = {
+    'neu': 'neutral', 'hap': 'happy', 'sad': 'sad', 'ang': 'angry'
+}
+AUDIO_SAMPLE_RATE = 16000
+
+# --- Model Loading ---
+# Use Streamlit's caching to load models only once.
+@st.cache_resource
+def load_models():
+    with st.spinner("Loading AI models... This may take a moment on first run."):
+        # Whisper for Speech-to-Text
+        whisper_model = whisper.load_model("base")
+
+        # Text Emotion Classifier
+        text_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None)
+
+        # Speech Emotion Recognition (SER)
+        ser_model_name = "superb/hubert-large-superb-er"
+        ser_feature_extractor = AutoFeatureExtractor.from_pretrained(ser_model_name)
+        ser_model = AutoModelForAudioClassification.from_pretrained(ser_model_name)
+
+        # NEW: Pre-load the DeepFace model to prevent lag on first use
+        DeepFace.build_model("Emotion")
+
+        return whisper_model, text_classifier, ser_model, ser_feature_extractor
+
+whisper_model, text_classifier, ser_model, ser_feature_extractor = load_models()
+
+# --- WebRTC and Video Processing ---
+webrtc_ctx = webrtc_streamer(
+    key="affectlink-video",
+    mode=WebRtcMode.SENDRECV,
+    rtc_configuration=RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}),
+    media_stream_constraints={"video": True, "audio": False},
+    async_processing=True,
+)
+
+if 'facial_emotion' not in st.session_state:
+    st.session_state.facial_emotion = "Neutral"
+if 'last_emotion_time' not in st.session_state:
+    st.session_state.last_emotion_time = 0
+
+def video_frame_callback(frame: av.VideoFrame) -> av.VideoFrame:
+    img = frame.to_ndarray(format="bgr24")
+    
+    # Process a frame every 5 seconds to conserve CPU
+    current_time = st.session_state.get('last_emotion_time', 0)
+    if torch.cuda.is_available() or (hasattr(st.session_state, 'last_emotion_time') and (torch.tensor(current_time).item() + 5 < torch.tensor(frame.time).item())):
+        try:
+            analysis = DeepFace.analyze(
+                img,
+                actions=['emotion'],
+                enforce_detection=False,
+                silent=True
+            )
+            if isinstance(analysis, list) and len(analysis) > 0:
+                dominant_emotion = analysis[0]['dominant_emotion']
+                st.session_state.facial_emotion = dominant_emotion.capitalize()
+            else:
+                st.session_state.facial_emotion = "Unknown"
+        except Exception as e:
+            logging.error(f"DeepFace analysis failed: {e}")
+            st.session_state.facial_emotion = "Error"
+        
+        st.session_state.last_emotion_time = frame.time
+
+    return av.VideoFrame.from_ndarray(img, format="bgr24")
+
+if webrtc_ctx.video_processor:
+    webrtc_ctx.video_processor.video_frame_callback = video_frame_callback
+
+# --- Audio Processing ---
+if "audio_buffer" not in st.session_state:
+    st.session_state.audio_buffer = []
+
+def audio_frame_callback(frame: av.AudioFrame):
+    sound = np.frombuffer(frame.to_ndarray(), dtype=np.int16)
+    st.session_state.audio_buffer.append(sound)
+
+webrtc_streamer(
+    key="affectlink-audio",
+    mode=WebRtcMode.RECVONLY,
+    media_stream_constraints={"video": False, "audio": True},
+    audio_frame_callback=audio_frame_callback,
+    async_processing=True,
+)
+
+# --- UI Layout ---
+st.sidebar.header("Facial Emotion")
+st.sidebar.metric("Current Expression", st.session_state.get('facial_emotion', 'N/A'))
+st.sidebar.info("Facial emotion is updated every 5 seconds to optimize performance.")
+st.sidebar.divider()
+
+st.sidebar.header("Audio Analysis")
+is_recording = st.sidebar.checkbox("Start Recording Audio")
+
+# NEW: Set up placeholders for audio results
+st.sidebar.subheader("Transcription:")
+transcription_placeholder = st.sidebar.empty()
+transcription_placeholder.write("_Waiting for audio..._")
+
+st.sidebar.subheader("Text Emotion:")
+text_emotion_placeholder = st.sidebar.empty()
+text_emotion_placeholder.write("_Waiting for audio..._")
+
+st.sidebar.subheader("Speech Emotion:")
+ser_placeholder = st.sidebar.empty()
+ser_placeholder.write("_Waiting for audio..._")
+
+
+if not is_recording and st.session_state.audio_buffer:
+    # Combine audio chunks
+    audio_data = np.concatenate(st.session_state.audio_buffer)
+    st.session_state.audio_buffer = []  # Clear buffer
+
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio_file:
+        write_wav(tmp_audio_file.name, AUDIO_SAMPLE_RATE, audio_data)
+        
+        # 1. Speech-to-Text (Whisper)
+        with st.spinner("Transcribing audio..."):
+            try:
+                result = whisper_model.transcribe(tmp_audio_file.name, fp16=False)
+                transcribed_text = result['text']
+            except Exception as e:
+                transcribed_text = f"Transcription failed: {e}"
+            transcription_placeholder.write(f'"{transcribed_text}"')
+
+        # 2. Text-based Emotion
+        with st.spinner("Analyzing text emotion..."):
+            if transcribed_text:
+                try:
+                    text_emotions = text_classifier(transcribed_text)[0]
+                    unified_text_scores = {e: 0.0 for e in UNIFIED_EMOTIONS}
+                    for EMO in text_emotions:
+                        unified_emo = TEXT_TO_UNIFIED.get(EMO['label'])
+                        if unified_emo:
+                            unified_text_scores[unified_emo] += EMO['score']
+                    
+                    dominant_text_emotion = max(unified_text_scores, key=unified_text_scores.get)
+                except Exception as e:
+                    dominant_text_emotion = f"Text analysis failed: {e}"
+            else:
+                dominant_text_emotion = "No text to analyze."
+            text_emotion_placeholder.write(dominant_text_emotion.capitalize())
+
+        # 3. Speech Emotion Recognition (SER)
+        with st.spinner("Analyzing speech emotion..."):
+            try:
+                audio_array, _ = sf.read(tmp_audio_file.name)
+                inputs = ser_feature_extractor(audio_array, sampling_rate=AUDIO_SAMPLE_RATE, return_tensors="pt", padding=True)
+                with torch.no_grad():
+                    logits = ser_model(**inputs).logits
+                scores = torch.nn.functional.softmax(logits, dim=1).squeeze()
+                
+                unified_ser_scores = {e: 0.0 for e in UNIFIED_EMOTIONS}
+                for i, score in enumerate(scores):
+                    raw_emo = ser_model.config.id2label[i]
+                    unified_emo = SER_TO_UNIFIED.get(raw_emo)
+                    if unified_emo:
+                        unified_ser_scores[unified_emo] += score.item()
+                
+                dominant_ser_emotion = max(unified_ser_scores, key=unified_ser_scores.get)
+            except Exception as e:
+                dominant_ser_emotion = f"Speech analysis failed: {e}"
+            ser_placeholder.write(dominant_ser_emotion.capitalize())
+
+        # Clean up the temporary file
+        os.unlink(tmp_audio_file.name)
+
+elif is_recording:
     st.sidebar.warning("Recording audio... Uncheck to stop and process.")