Update visualization.py

visualization.py

@@ -154,6 +154,41 @@ def plot_combined_mse(df, mse_embeddings, mse_posture, mse_voice, title):
     plt.close()
     return fig
 
+def plot_combined_heatmap(mse_embeddings, mse_posture, mse_voice, df):
+    # Normalize MSE values
+    mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
+    mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
+    mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
+
+    combined_mse = np.zeros((3, len(mse_embeddings)))
+    combined_mse[0] = mse_embeddings_norm
+    combined_mse[1] = mse_posture_norm
+    combined_mse[2] = mse_voice_norm
+
+    plt.figure(figsize=(20, 3), dpi=300)
+    fig, ax = plt.subplots(figsize=(20, 3))
+
+    # Create heatmap
+    sns.heatmap(combined_mse, cmap='Reds', cbar=False, ax=ax)
+
+    # Set y-axis labels
+    ax.set_yticks([0.5, 1.5, 2.5])
+    ax.set_yticklabels(['Facial Features', 'Body Posture', 'Voice'], fontsize=8) 
+
+    # Set x-axis ticks to timecodes
+    num_ticks = min(60, len(mse_embeddings))
+    tick_locations = np.linspace(0, len(mse_embeddings) - 1, num_ticks).astype(int)
+    tick_labels = [df['Timecode'].iloc[i] if i < len(df) else '' for i in tick_locations]
+
+    ax.set_xticks(tick_locations)
+    ax.set_xticklabels(tick_labels, rotation=90, ha='center', va='top')
+
+    ax.set_title('Combined MSE Heatmap')
+
+    plt.tight_layout()
+    plt.close(fig)
+    return fig
+
 def plot_mse_histogram(mse_values, title, anomaly_threshold, color='blue'):
     plt.figure(figsize=(16, 3), dpi=300)
     fig, ax = plt.subplots(figsize=(16, 3))