Update app.py

app.py

@@ -9,7 +9,7 @@ from facenet_pytorch import MTCNN
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from PIL import Image
 import os
-import time
+from collections import Counter
 
 # Load models
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -18,13 +18,13 @@ model = AutoModelForImageClassification.from_pretrained("trpakov/vit-face-expres
 extractor = AutoFeatureExtractor.from_pretrained("trpakov/vit-face-expression")
 
 # Emotion labels
-emotion_labels = {
+affectnet_labels = {
     0: "neutral", 1: "happy", 2: "sad", 3: "surprise", 4: "fear",
     5: "disgust", 6: "anger", 7: "contempt"
 }
 
-# Emotion detection function
-def detect_emotion(frame):
+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:
@@ -34,29 +34,27 @@ def detect_emotion(frame):
     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()]
 
-# Process Video
-def process_video(video_path, progress=gr.Progress()):
-    cap = cv2.VideoCapture(video_path)
+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))
-    out_path = "output_video.mp4"
-    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
-
-    emotion_counts = {}
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    processed_frames = 0
+    out = cv2.VideoWriter("output_video.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
 
+    emotion_counts = []
+    
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
 
-        emotion = detect_emotion(frame)
-        emotion_counts[emotion] = emotion_counts.get(emotion, 0) + 1
+        emotion = detect_emotions(frame)
+        emotion_counts.append(emotion)
 
-        # Overlay emotion on frame
+        # 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)
@@ -64,78 +62,41 @@ def process_video(video_path, progress=gr.Progress()):
 
         out.write(frame)
 
-        processed_frames += 1
-        progress((processed_frames / total_frames) * 100)
-
     cap.release()
     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 Pie Chart
-    plt.figure(figsize=(6, 6))
-    labels, sizes = zip(*emotion_counts.items())
+    # 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")
     plt.close()
 
-    return emotion, "emotion_distribution.jpg", out_path
-
-# Custom CSS for styling
-css = """
-h1 {
-    text-align: center;
-    color: #ffffff;
-    font-size: 32px;
-    font-weight: bold;
-}
-
-.gradio-container {
-    background-color: #1E1E1E;
-    color: #ffffff;
-    padding: 20px;
-    font-family: 'Arial', sans-serif;
-}
-
-button {
-    font-size: 18px !important;
-    padding: 10px 15px !important;
-    background-color: #00BFFF !important;
-    color: white !important;
-    border-radius: 10px !important;
-}
-
-.gr-text-input {
-    background-color: #2E2E2E;
-    color: white;
-    border: 1px solid #00BFFF;
-}
-"""
-
-# Gradio Interface with Enhanced UI
-with gr.Blocks(css=css) as demo:
-    with gr.Row():
-        gr.Markdown("<h1>🎭 Emotion Analysis from Video 🎥</h1>")
-    
-    with gr.Row():
-        video_input = gr.File(label="📤 Upload your video", type="filepath")
-    
-    with gr.Row():
-        analyze_button = gr.Button("🚀 Analyze Video")
-
-    with gr.Row():
-        result_text = gr.Textbox(label="Detected Emotion", interactive=False)
-    
-    with gr.Row():
-        emotion_chart = gr.Image(label="📊 Emotion Distribution", interactive=False)
-    
-    with gr.Row():
-        processed_video = gr.Video(label="🎞 Processed Video with Emotion Detection")
-
-    analyze_button.click(
-        process_video,
-        inputs=[video_input],
-        outputs=[result_text, emotion_chart, processed_video]
-    )
-
-# Launch Gradio app
-demo.launch()
+    return "output_video.mp4", "emotion_distribution.jpg", major_emotion
+
+# Gradio Web Interface
+gr.Interface(
+    fn=process_video,
+    inputs=gr.File(type="filepath"),
+    outputs=[
+        gr.File(label="Processed Video"),
+        gr.File(label="Emotion Distribution Chart"),
+        gr.Textbox(label="Major Emotion Detected")
+    ],
+    title="Emotion Detection from Video",
+    description="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.",
+    css="""
+    .gradio-container { max-width: 800px !important; margin: auto; }
+    .gradio-container h1 { font-size: 22px; }
+    @media screen and (max-width: 768px) {
+        .gradio-container { width: 100%; padding: 10px; }
+        .gradio-container h1 { font-size: 18px; }
+    }
+    """
+).launch()