Update visualization.py

visualization.py

@@ -207,12 +207,13 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps, largest_cluster, progress=gr.Progress()):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_path, desired_fps, largest_cluster):
     # Filter the DataFrame to only include frames from the largest cluster
     df_largest_cluster = df[df['Cluster'] == largest_cluster]
     mse_embeddings = mse_embeddings[df['Cluster'] == largest_cluster]
     mse_posture = mse_posture[df['Cluster'] == largest_cluster]
-
+    mse_voice = mse_voice[df['Cluster'] == largest_cluster]
+    
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -226,14 +227,19 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
                                np.arange(len(mse_embeddings)), mse_embeddings)
     mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
                             np.arange(len(mse_posture)), mse_posture)
+    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
+                          np.arange(len(mse_voice)), mse_voice)
+
     
     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((2, total_frames))
     combined_mse[0] = mse_embeddings_norm  # Use normalized MSE values for facial
     combined_mse[1] = mse_posture_norm  # Use normalized MSE values for posture
-
+    combined_mse[2] = mse_voice_norm
+    
     # Custom colormap definition
     cdict = {
         'red':   [(0.0,  0.5, 0.5),  # Low MSE: 50% red (gray)
@@ -248,9 +254,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
     # Use the custom colormap in the heatmap generation
-    im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 2])
-    ax.set_yticks([0.5, 1.5])
-    ax.set_yticklabels(['Face', 'Posture'])
+    im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 3])
+    ax.set_yticks([0.5, 1.5, 2.5])
+    ax.set_yticklabels(['Face', 'Posture', 'Voice'])
     ax.set_xticks([])
     plt.tight_layout()