Update visualization.py

visualization.py

@@ -212,20 +212,17 @@ 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, mse_voice, output_path, desired_fps, largest_cluster, progress=None):
-    print(f"Creating heatmap video. Output path: {output_path}")
-    # Filter the DataFrame to only include frames from the largest cluster
-    df_largest_cluster = df[df['Cluster'] == largest_cluster]
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster, progress=None):
+    print(f"Creating heatmap video. Output folder: {output_folder}")
     
-    # Interpolate mse_voice to match the length of df
-    x_voice = np.linspace(0, len(mse_voice)-1, len(mse_voice))
-    x_new = np.linspace(0, len(mse_voice)-1, len(df))
-    f = interpolate.interp1d(x_voice, mse_voice)
-    mse_voice_interpolated = f(x_new)
+    # Create output folder if it doesn't exist
+    os.makedirs(output_folder, exist_ok=True)
     
-    mse_embeddings = mse_embeddings[df['Cluster'] == largest_cluster]
-    mse_posture = mse_posture[df['Cluster'] == largest_cluster]
-    mse_voice_interpolated = mse_voice_interpolated[df['Cluster'] == largest_cluster]
+    # Define output path
+    output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
+    heatmap_video_path = os.path.join(output_folder, output_filename)
+    
+    print(f"Heatmap video will be saved at: {heatmap_video_path}")
 
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
@@ -234,7 +231,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, original_fps, (width, height + 200))
+    out = cv2.VideoWriter(heatmap_video_path, fourcc, original_fps, (width, height + 200))
     print(f"VideoWriter initialized. FPS: {original_fps}, Size: {(width, height + 200)}")
     
     # Ensure all MSE arrays have the same length as total_frames
@@ -242,12 +239,12 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
                                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_interpolated = np.interp(np.linspace(0, len(mse_voice_interpolated) - 1, total_frames), 
-                                       np.arange(len(mse_voice_interpolated)), mse_voice_interpolated)
+    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_interpolated - np.min(mse_voice_interpolated)) / (np.max(mse_voice_interpolated) - np.min(mse_voice_interpolated))
+    mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
     
     combined_mse = np.zeros((3, total_frames))
     combined_mse[0] = mse_embeddings_norm
@@ -307,5 +304,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     cap.release()
     out.release()
     plt.close(fig)
-    print(f"Heatmap video created at: {output_path}")
-    return output_path
+    
+    print(f"Heatmap video created at: {heatmap_video_path}")
+    return heatmap_video_path