Update visualization.py

visualization.py

@@ -217,9 +217,15 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width, largest_cluster):
     frame_count = int(t * video_fps)
     
+    # Replace MSE values outside of the largest cluster with zeros
+    mask = (largest_cluster == 1)
+    mse_embeddings[~mask] = 0
+    mse_posture[~mask] = 0
+    mse_voice[~mask] = 0
+
     # 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))
@@ -274,7 +280,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames, width)
+        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames, width, largest_cluster)
         heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
         combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame