Update chart

app.py

@@ -11,6 +11,7 @@ import torch.nn.functional as F
 from torchvision import transforms
 from facenet_pytorch import MTCNN
 import gradio as gr
+import seaborn as sns
 
 class EmotionModel(torch.nn.Module):
     def __init__(self):
@@ -43,9 +44,6 @@ class EmotionDetector:
         ])
         self.softmax = torch.nn.Softmax(dim=1)
 
-        # Load pre-trained weights here if available
-        # self.model.load_state_dict(torch.load("emotion_model.pt", map_location=self.device))
-
     def detect_emotions_video(self, video_path, sample_rate=30, max_size_mb=50):
         try:
             if video_path is None:
@@ -57,7 +55,6 @@ class EmotionDetector:
             cap = cv2.VideoCapture(video_path)
             fps = int(cap.get(cv2.CAP_PROP_FPS))
             frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
             if frame_count == 0:
                 return None, "Invalid video file"
 
@@ -79,7 +76,7 @@ class EmotionDetector:
 
                 face_tensor = self.transform(face_tensor)  # Resize
                 face_tensor = face_tensor.mean(dim=0, keepdim=True)  # grayscale
-                face_tensor = face_tensor.unsqueeze(0).to(self.device)  # batch + channel
+                face_tensor = face_tensor.unsqueeze(0).to(self.device)
 
                 with torch.no_grad():
                     output = self.model(face_tensor)
@@ -95,25 +92,39 @@ class EmotionDetector:
                 return None, "No emotions detected."
 
             df = pd.DataFrame(emotions_over_time)
-
-            plt.figure(figsize=(12, 8))
-            for emotion in self.emotions:
-                if emotion in df.columns:
-                    plt.plot(df['timestamp'], df[emotion], label=emotion.title(), linewidth=2)
-
-            plt.xlabel('Time (seconds)')
-            plt.ylabel('Confidence (%)')
-            plt.title('Emotions Over Time')
-            plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
-            plt.grid(True)
-            plt.tight_layout()
-
-            img_buf = io.BytesIO()
-            plt.savefig(img_buf, format='png', dpi=150, bbox_inches='tight')
-            img_buf.seek(0)
+            df['dominant_emotion'] = df[self.emotions].idxmax(axis=1)
+
+            # --- Chart Plotting ---
+            fig, axs = plt.subplots(2, 1, figsize=(12, 10), constrained_layout=True)
+
+            # 1. Stacked Area Chart
+            df_sorted = df.sort_values("timestamp")
+            axs[0].stackplot(df_sorted["timestamp"], [df_sorted[e] for e in self.emotions], labels=[e.title() for e in self.emotions])
+            axs[0].set_title("Emotions Over Time")
+            axs[0].set_xlabel("Time (seconds)")
+            axs[0].set_ylabel("Confidence (%)")
+            axs[0].legend(loc="upper right")
+            axs[0].grid(True)
+
+            # 2. Dominant Emotion Timeline (Bar Chart)
+            color_palette = sns.color_palette("husl", len(self.emotions))
+            emotion_color_map = {e: color_palette[i] for i, e in enumerate(self.emotions)}
+
+            colors = df['dominant_emotion'].map(emotion_color_map)
+            axs[1].bar(df['timestamp'], 1, color=colors, width=sample_rate / fps)
+            axs[1].set_title("Dominant Emotion Timeline")
+            axs[1].set_xlabel("Time (seconds)")
+            axs[1].set_yticks([])
+            axs[1].legend(handles=[plt.Rectangle((0, 0), 1, 1, color=emotion_color_map[e]) for e in self.emotions],
+                          labels=[e.title() for e in self.emotions], loc="upper right", title="Emotion")
+
+            # Save chart
+            buf = io.BytesIO()
+            plt.savefig(buf, format="png")
+            buf.seek(0)
             plt.close()
 
-            chart_image = Image.open(img_buf)
+            chart_image = Image.open(buf)
             avg_emotions = df[self.emotions].mean().sort_values(ascending=False)
 
             result_text = f"**Video Analysis Complete**\n"
@@ -128,6 +139,7 @@ class EmotionDetector:
         except Exception as e:
             return None, f"Error: {str(e)}"
 
+# Gradio interface
 def create_interface():
     detector = EmotionDetector()
 
@@ -140,7 +152,7 @@ def create_interface():
     return gr.Interface(
         fn=process,
         inputs=[
-            gr.Video(label="Upload Video"),  # Removed type="file"
+            gr.Video(label="Upload Video"),
             gr.Slider(minimum=1, maximum=60, step=1, value=30, label="Sample Rate (Frames)")
         ],
         outputs=[
@@ -151,6 +163,5 @@ def create_interface():
         description="Upload a video to analyze emotions over time."
     )
 
-
 if __name__ == "__main__":
     create_interface().launch()