Refactor Streamlit app to support video uploads for emotion analysis and enhance audio processing capabilities

requirements.txt

@@ -1,12 +1,26 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 
-# Core app library
+# Pin the main UI components to recent, stable versions
 streamlit==1.35.0
 
-# AI libraries for this test
+# Library for video/audio file handling
+moviepy==1.0.3
+
+# Pin ML/AI libraries to modern, known-good versions
+transformers==4.40.1
+deepface==0.0.94
 openai-whisper==20231117
+
+# Pin frameworks to ensure CPU versions and prevent build timeouts
+tensorflow-cpu==2.16.1
+tf-keras==2.16.0
 torch==2.7.0
 torchaudio==2.7.0
 
-# Utility for handling audio files
-soundfile
+# Pin data/audio libraries for stability
+pandas==2.2.2
+numpy==1.26.4
+soundfile==0.12.1
+librosa==0.10.1
+scipy==1.13.0
+Pillow==10.3.0

src/streamlit_app.py

@@ -1,60 +1,174 @@
 import os
 import streamlit as st
+import numpy as np
+import torch
 import whisper
-import tempfile
+from transformers import pipeline, AutoModelForAudioClassification, AutoFeatureExtractor
+from deepface import DeepFace
 import logging
+import soundfile as sf
+import tempfile
+from PIL import Image
+import cv2
+from moviepy.editor import VideoFileClip
 
-# --- Set cache directory for models ---
-# Point the cache to a writable directory in the cloud environment
+# Set home directories for model caching to the writable /tmp folder
+os.environ['DEEPFACE_HOME'] = '/tmp/.deepface'
 os.environ['HF_HOME'] = '/tmp/huggingface'
 
 # --- Page Configuration ---
 st.set_page_config(
-    page_title="Audio Transcription Test",
-    page_icon="🎤",
-    layout="centered"
+    page_title="AffectLink Demo",
+    page_icon="😊",
+    layout="wide"
 )
 
-st.title("Step 2: Audio Transcription (Whisper) Test")
-st.write("Upload a short audio file (.wav, .mp3, .m4a) to test the Whisper model.")
+st.title("AffectLink: Post-Hoc Emotion Analysis")
+st.write("Upload a short video clip (under 30 seconds) to analyze facial expressions, speech-to-text, and the emotional tone of the audio.")
 
 # --- Logger Configuration ---
 logging.basicConfig(level=logging.INFO)
+logging.getLogger('deepface').setLevel(logging.ERROR)
 logging.getLogger('huggingface_hub').setLevel(logging.WARNING)
+logging.getLogger('moviepy').setLevel(logging.ERROR)
+
+# --- Emotion Mappings ---
+UNIFIED_EMOTIONS = ['neutral', 'happy', 'sad', 'angry']
+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 ---
 @st.cache_resource
-def load_whisper_model():
-    with st.spinner("Loading Whisper model... This may take a moment."):
-        # Load the base model and specify a writable cache directory
-        model = whisper.load_model("base", download_root="/tmp/whisper_cache")
-    return model
-
-whisper_model = load_whisper_model()
+def load_models():
+    with st.spinner("Loading AI models, this may take a moment..."):
+        whisper_model = whisper.load_model("base", download_root="/tmp/whisper_cache")
+        text_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None)
+        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)
+        # DeepFace loads its own models on first use, no need to preload here.
+        return whisper_model, text_classifier, ser_model, ser_feature_extractor
 
+whisper_model, text_classifier, ser_model, ser_feature_extractor = load_models()
 
 # --- UI and Processing Logic ---
-uploaded_file = st.file_uploader("Choose an audio file...", type=["wav", "mp3", "m4a"])
+uploaded_file = st.file_uploader("Choose a video file...", type=["mp4", "mov", "avi", "mkv"])
 
 if uploaded_file is not None:
     # Save the uploaded file to a temporary location
-    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tfile:
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
         tfile.write(uploaded_file.read())
-        temp_audio_path = tfile.name
+        temp_video_path = tfile.name
+
+    st.video(temp_video_path)
+
+    if st.button("Analyze Video"):
+        facial_analysis_results = []
+        audio_analysis_results = {}
+
+        # --- Video Processing for Facial Emotion ---
+        with st.spinner("Analyzing video for facial expressions... (1 frame per second)"):
+            try:
+                cap = cv2.VideoCapture(temp_video_path)
+                fps = cap.get(cv2.CAP_PROP_FPS)
+                if fps == 0:
+                    fps = 30 # Default to 30 fps if not available
+                frame_count = 0
+                while cap.isOpened():
+                    ret, frame = cap.read()
+                    if not ret:
+                        break
+                    
+                    if frame_count % int(fps) == 0:
+                        timestamp = frame_count / fps
+                        analysis = DeepFace.analyze(frame, actions=['emotion'], enforce_detection=False, silent=True)
+                        if isinstance(analysis, list) and len(analysis) > 0:
+                            dominant_emotion = analysis[0]['dominant_emotion']
+                            facial_analysis_results.append((timestamp, dominant_emotion.capitalize()))
+                    frame_count += 1
+                cap.release()
+            except Exception as e:
+                st.error(f"An error occurred during facial analysis: {e}")
+
+        # --- Audio Extraction and Processing ---
+        with st.spinner("Extracting and analyzing audio..."):
+            temp_audio_path = None
+            video_clip = None
+            try:
+                video_clip = VideoFileClip(temp_video_path)
+                if video_clip.audio is not None:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as taudio:
+                        video_clip.audio.write_audiofile(taudio.name, fps=AUDIO_SAMPLE_RATE, logger=None)
+                        temp_audio_path = taudio.name
+
+                    # 1. Speech-to-Text (Whisper)
+                    result = whisper_model.transcribe(temp_audio_path, fp16=False)
+                    transcribed_text = result['text'] if result['text'] else "No speech detected."
+                    audio_analysis_results['Transcription'] = transcribed_text
+
+                    # 2. Text-based Emotion
+                    if "No speech detected" not in transcribed_text:
+                        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)
+                        audio_analysis_results['Text Emotion'] = dominant_text_emotion.capitalize()
+                    
+                    # 3. Speech Emotion Recognition (SER)
+                    audio_array, _ = sf.read(temp_audio_path)
+                    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)
+                    audio_analysis_results['Speech Emotion'] = dominant_ser_emotion.capitalize()
+                else:
+                    audio_analysis_results['Transcription'] = "No audio track found in the video."
+
+            except Exception as e:
+                st.error(f"An error occurred during audio analysis: {e}")
+            finally:
+                if video_clip:
+                    video_clip.close()
+                if temp_audio_path and os.path.exists(temp_audio_path):
+                    os.unlink(temp_audio_path)
 
-    st.audio(temp_audio_path)
+        # --- Display Results ---
+        st.header("Analysis Results")
+        col1, col2 = st.columns(2)
 
-    with st.spinner("Transcribing audio..."):
-        try:
-            # Transcribe the audio file using Whisper
-            result = whisper_model.transcribe(temp_audio_path, fp16=False)
-            transcribed_text = result.get("text", "No text transcribed.")
+        with col1:
+            st.subheader("Audio Analysis")
+            if audio_analysis_results:
+                st.write(f"**Transcription:** \"{audio_analysis_results.get('Transcription', 'N/A')}\"")
+                st.metric("Emotion from Text", audio_analysis_results.get('Text Emotion', 'N/A'))
+                st.metric("Emotion from Speech", audio_analysis_results.get('Speech Emotion', 'N/A'))
+            else:
+                st.write("No audio results to display.")
 
-            st.success("Transcription Complete!")
-            st.markdown(f"> {transcribed_text}")
+        with col2:
+            st.subheader("Facial Expression Timeline")
+            if facial_analysis_results:
+                for timestamp, emotion in facial_analysis_results:
+                    st.write(f"**Time {int(timestamp // 60):02d}:{int(timestamp % 60):02d}:** {emotion}")
+            else:
+                st.write("No faces detected or video processing failed.")
 
-        except Exception as e:
-            st.error(f"An error occurred during transcription: {e}")
-        finally:
-            # Clean up the temporary file
-            os.unlink(temp_audio_path)
+    # Clean up temp video file after analysis is done
+    if os.path.exists(temp_video_path):
+        os.unlink(temp_video_path)