Update visualization.py

visualization.py

@@ -127,7 +127,7 @@ def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_thre
     plt.close()
     return fig, anomaly_frames
 
-def plot_combined_mse(df, mse_embeddings, mse_posture, mse_voice, title):
+def plot_combined_mse(df, mse_embeddings, mse_posture, mse_voice, title, anomaly_threshold=4, time_threshold=3):
     plt.figure(figsize=(16, 8), dpi=300)
     fig, ax = plt.subplots(figsize=(16, 8))
 
@@ -135,17 +135,57 @@ def plot_combined_mse(df, mse_embeddings, mse_posture, mse_voice, title):
         df['Seconds'] = df['Timecode'].apply(
             lambda x: sum(float(t) * 60 ** i for i, t in enumerate(reversed(x.split(':')))))
 
-    # Normalize MSE values to the same scale
-    mse_embeddings = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
-    mse_posture = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
-    mse_voice = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
+    # Function to identify continuous segments
+    def get_continuous_segments(seconds, values, max_gap=1):
+        segments = []
+        current_segment = []
+        for i, (sec, val) in enumerate(zip(seconds, values)):
+            if not current_segment or (sec - current_segment[-1][0] <= max_gap):
+                current_segment.append((sec, val))
+            else:
+                segments.append(current_segment)
+                current_segment = [(sec, val)]
+        if current_segment:
+            segments.append(current_segment)
+        return segments
+
+    # Plot each data series
+    for mse_values, color, label in zip([mse_embeddings, mse_posture, mse_voice], 
+                                        ['navy', 'purple', 'green'], 
+                                        ['Facial Features', 'Body Posture', 'Voice']):
+        segments = get_continuous_segments(df['Seconds'], mse_values)
+        
+        for segment in segments:
+            segment_seconds, segment_mse = zip(*segment)
+            ax.scatter(segment_seconds, segment_mse, color=color, alpha=0.3, s=5, label=label if segment == segments[0] else "")
+            
+            if len(segment) > 1:
+                segment_df = pd.DataFrame({'Seconds': segment_seconds, 'MSE': segment_mse})
+                segment_df = segment_df.sort_values('Seconds')
+                mean = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).mean()
+                std = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).std()
+                
+                ax.plot(segment_df['Seconds'], mean, color=color, linewidth=0.5)
+                ax.fill_between(segment_df['Seconds'], mean - std, mean + std, color=color, alpha=0.1)
+
+        # Plot median baseline for each series
+        median = np.median(mse_values)
+        ax.axhline(y=median, color=color, linestyle='--', alpha=0.5, label=f'{label} Median')
+
+        # Plot threshold for each series
+        threshold = np.mean(mse_values) + anomaly_threshold * np.std(mse_values)
+        ax.axhline(y=threshold, color=color, linestyle=':', alpha=0.5, label=f'{label} Threshold')
+
+    max_seconds = df['Seconds'].max()
+    num_ticks = 100
+    tick_locations = np.linspace(0, max_seconds, num_ticks)
+    tick_labels = [seconds_to_timecode(int(s)) for s in tick_locations]
 
-    ax.plot(df['Seconds'], mse_embeddings, color='navy', label='Facial Features')
-    ax.plot(df['Seconds'], mse_posture, color='purple', label='Body Posture')
-    ax.plot(df['Seconds'], mse_voice, color='green', label='Voice')
+    ax.set_xticks(tick_locations)
+    ax.set_xticklabels(tick_labels, rotation=90, ha='center', fontsize=6)
 
     ax.set_xlabel('Timecode')
-    ax.set_ylabel('Normalized Mean Squared Error')
+    ax.set_ylabel('Mean Squared Error')
     ax.set_title(title)
 
     ax.grid(True, linestyle='--', alpha=0.7)
@@ -257,7 +297,7 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
 
-    fig, ax = plt.subplots(figsize=(video_width / 200, 0.5))
+    fig, ax = plt.subplots(figsize=(video_width / 200, 0.4))
     ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, total_frames, 0, 3])
     ax.set_yticks([0.5, 1.5, 2.5])
     ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7)