Update visualization.py

visualization.py

@@ -217,8 +217,8 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice):
-    frame_count = int(t * video.fps)
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames):
+    frame_count = int(t * video_fps)
     
     # Normalize MSE values
     mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
@@ -245,7 +245,6 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice):
     plt.close(fig)
     return heatmap_img
 
-# Function to create video with heatmap
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
@@ -271,20 +270,11 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     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((3, total_frames))
-    combined_mse[0] = mse_embeddings_norm
-    combined_mse[1] = mse_posture_norm
-    combined_mse[2] = mse_voice_norm
-
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice)
+        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames)
         combined_frame = np.vstack((video_frame, heatmap_frame))
-        return add_timecode(combined_frame, t)
+        return combined_frame
 
     final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
     final_clip = final_clip.set_audio(video.audio)