Update visualization.py

visualization.py

@@ -204,7 +204,12 @@ 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):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, 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]
+
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -214,24 +219,20 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
     
-    # Create custom colormap
     cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
         [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
     
-    # Ensure heatmap data covers all frames
     mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
                                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)
     
-    # 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))
     
-    # Combine MSEs
     combined_mse = np.zeros((2, total_frames, 3))
-    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, mse_embeddings_norm]).T  # RGB for facial
-    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T  # RGB for posture
+    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, mse_embeddings_norm]).T
+    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
     im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
@@ -272,3 +273,4 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     plt.close(fig)
     
     return output_path
+