Update app.py

app.py

@@ -1,130 +1,107 @@
-# app.py
 import streamlit as st
 import cv2
-import numpy as np
-import moviepy.editor as mp
-from transformers import ViTImageProcessor, ViTForImageClassification
 import torch
-from speechbrain.pretrained import EncoderClassifier
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from facenet_pytorch import MTCNN
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from PIL import Image
+from collections import Counter
 import tempfile
 import os
 
-# 1. Load Models with caching
-@st.cache_resource
-def load_models():
-    # Load ViT model for facial emotion detection
-    vit_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
-    vit_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
-    
-    # Load SpeechBrain model for audio emotion recognition
-    audio_classifier = EncoderClassifier.from_hparams(
-        source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
-        savedir="pretrained_models/emotion-audio"
-    )
-    
-    return vit_processor, vit_model, audio_classifier
-
-# 2. Video Processing Functions
-def analyze_frame(frame, processor, model):
-    """Analyze single frame using ViT model"""
-    inputs = processor(images=frame, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return model.config.id2label[outputs.logits.argmax(-1).item()]
-
-def process_video(video_path, processor, model, audio_classifier):
-    """Process video and return combined results"""
-    # Extract audio from video
-    video = mp.VideoFileClip(video_path)
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_audio:
-        audio_path = tmp_audio.name
-        video.audio.write_audiofile(audio_path)
-    
-    # Analyze audio
-    audio_signal = audio_classifier.load_audio(audio_path)
-    audio_prediction = audio_classifier.classify_batch(audio_signal)
-    audio_emotion = audio_prediction[3][0]
-    
-    # Analyze video frames
-    cap = cv2.VideoCapture(video_path)
-    emotions = []
+# Load models
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+mtcnn = MTCNN(device=device)
+model = AutoModelForImageClassification.from_pretrained("trpakov/vit-face-expression").to(device)
+extractor = AutoFeatureExtractor.from_pretrained("trpakov/vit-face-expression")
+
+# Emotion labels
+affectnet_labels = {
+    0: "neutral", 1: "happy", 2: "sad", 3: "surprise", 4: "fear",
+    5: "disgust", 6: "anger", 7: "contempt"
+}
+
+def detect_emotions(frame):
+    """Detects facial emotions in a given frame."""
+    img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    faces, _ = mtcnn.detect(img)
+    if faces is None or len(faces) == 0:
+        return "No Face Detected"
+    
+    face = img.crop(faces[0])
+    inputs = extractor(images=face, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    
+    return model.config.id2label[torch.argmax(probs).item()]
+
+def process_video(input_path):
+    """Processes video, overlays emotions, and creates a summary chart."""
+    cap = cv2.VideoCapture(input_path)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    frame_width, frame_height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    # Create a temporary output video file
+    output_path = "output_video.mp4"
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+
+    emotion_counts = []
     
-    # Process every 5th frame to reduce computation
-    frame_count = 0
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-            
-        if frame_count % 5 == 0:  # Sample every 5th frame
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            emotions.append(analyze_frame(frame_rgb, processor, model))
-            
-        frame_count += 1
-    
+
+        emotion = detect_emotions(frame)
+        emotion_counts.append(emotion)
+
+        # Overlay emotion
+        overlay = frame.copy()
+        cv2.rectangle(overlay, (10, 10), (350, 80), (255, 255, 255), -1)
+        cv2.putText(overlay, f'Emotion: {emotion}', (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
+        cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)
+
+        out.write(frame)
+
     cap.release()
-    os.unlink(audio_path)  # Clean up temporary audio file
-    
-    # Get most common visual emotion
-    visual_emotion = max(set(emotions), key=emotions.count)
-    
-    return {
-        'audio_emotion': audio_emotion,
-        'visual_emotion': visual_emotion,
-        'frame_emotions': emotions
-    }
-
-# 3. Streamlit UI
-st.set_page_config(page_title="Video Sentiment Analyzer", layout="wide")
-
-st.title("🎥 Video Sentiment Analysis")
-st.markdown("""
-Analyze emotions from:
-- **Facial Expressions** using ViT (Vision Transformer)
-- **Speech Tone** using wav2vec2
-""")
-
-uploaded_file = st.file_uploader("Upload a video file (max 30 seconds)", type=["mp4", "mov", "avi"])
-
-if uploaded_file:
-    # Display video preview
-    st.video(uploaded_file)
-    
-    # Save to temporary file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
-        tmp_video.write(uploaded_file.getbuffer())
-        video_path = tmp_video.name
-    
-    # Load models
-    vit_processor, vit_model, audio_classifier = load_models()
-    
+    out.release()
+    cv2.destroyAllWindows()
+
+    # Find major emotion
+    emotion_counter = Counter(emotion_counts)
+    major_emotion = emotion_counter.most_common(1)[0][0] if emotion_counter else "No Face Detected"
+
+    # Generate emotion distribution pie chart
+    plt.figure(figsize=(5, 5))
+    labels, sizes = zip(*emotion_counter.items())
+    plt.pie(sizes, labels=labels, autopct='%1.1f%%', colors=sns.color_palette('pastel'))
+    plt.title("Emotion Distribution")
+    plt.savefig("emotion_distribution.jpg")
+    
+    return output_path, plt, major_emotion
+
+# Streamlit Web Interface
+st.set_page_config(page_title="Emotion Analysis from Video", layout="wide")
+
+st.title("🎭 Emotion Analysis from Video 🎥")
+st.markdown("Upload a video, and the AI will detect emotions in each frame, providing a processed video, an emotion distribution chart, and the major detected emotion.")
+
+# File uploader
+video_input = st.file_uploader("📤 Upload Video (MP4, MOV, AVI)", type=["mp4", "mov", "avi"])
+
+if video_input is not None:
+    # Save uploaded video to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file:
+        tmp_file.write(video_input.read())
+        video_path = tmp_file.name
+
     # Process video
-    with st.spinner("Analyzing video content..."):
-        try:
-            results = process_video(video_path, vit_processor, vit_model, audio_classifier)
-        finally:
-            os.unlink(video_path)  # Clean up temporary video file
-    
-    # Display results
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        st.subheader("🎧 Audio Analysis")
-        st.metric("Dominant Emotion", results['audio_emotion'])
-    
-    with col2:
-        st.subheader("👁️ Visual Analysis")
-        st.metric("Dominant Emotion", results['visual_emotion'])
-    
-    # Show emotion timeline
-    st.subheader("📈 Emotion Timeline")
-    st.line_chart(
-        data={ "Frame Emotions": results['frame_emotions'] },
-        use_container_width=True
-    )
-    
-    st.success("Analysis complete!")
+    if st.button("🚀 Analyze"):
+        with st.spinner("Processing video..."):
+            output_video, emotion_chart, major_emotion = process_video(video_path)
 
-# Footer
-st.markdown("---")
-st.markdown("Built with [Hugging Face](https://huggingface.co/) 🤗 & [Streamlit](https://streamlit.io/) 🎈")
+        # Display results
+        st.subheader("📥 Processed Video")
+