Update visualization.py

visualization.py

@@ -217,27 +217,45 @@ 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):
-    frame_count = int(t * video_fps)
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, analysis_fps, video_width):
+    video_frame_count = int(t * video_fps)
+    analysis_frame_count = int(t * analysis_fps)
+    
+    # Ensure we don't go out of bounds
+    analysis_frame_count = min(analysis_frame_count, len(mse_embeddings) - 1)
     
     # 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))
-    mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
+    def safe_normalize(arr):
+        min_val, max_val = np.min(arr), np.max(arr)
+        if min_val == max_val:
+            return np.zeros_like(arr)
+        return (arr - min_val) / (max_val - min_val)
+
+    mse_embeddings_norm = safe_normalize(mse_embeddings)
+    mse_posture_norm = safe_normalize(mse_posture)
+    mse_voice_norm = safe_normalize(mse_voice)
+    
+    # Create heatmap data
+    heatmap_width = int(video_width / 240 * 100)  # Adjust this multiplier as needed
+    combined_mse = np.zeros((3, heatmap_width))
     
-    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
+    # Map analysis frames to heatmap width
+    for i in range(heatmap_width):
+        frame_index = int(i * len(mse_embeddings) / heatmap_width)
+        combined_mse[0, i] = mse_embeddings_norm[frame_index]
+        combined_mse[1, i] = mse_posture_norm[frame_index]
+        combined_mse[2, i] = mse_voice_norm[frame_index]
 
     fig, ax = plt.subplots(figsize=(video_width / 240, 0.5))
-    ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, total_frames, 0, 3])
+    ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, heatmap_width, 0, 3])
     ax.set_yticks([0.5, 1.5, 2.5])
     ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7) 
     ax.set_xticks([])
-
-    ax.axvline(x=frame_count, color='black', linewidth=3)
-
+    
+    # Calculate the position of the vertical line
+    line_pos = (video_frame_count / video_fps) / (len(mse_embeddings) / analysis_fps) * heatmap_width
+    ax.axvline(x=line_pos, color='black', linewidth=3)
+    
     plt.tight_layout(pad=0.5)
     
     canvas = FigureCanvas(fig)
@@ -247,7 +265,7 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
     plt.close(fig)
     return heatmap_img
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, analysis_fps, largest_cluster):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
     os.makedirs(output_folder, exist_ok=True)
@@ -262,31 +280,30 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     # Get video properties
     width, height = video.w, video.h
-    total_frames = int(video.duration * video.fps)
+    video_duration = video.duration
+    video_fps = video.fps
+    
+    # Calculate the number of analysis frames
+    analysis_frames = int(video_duration * analysis_fps)
+    
+    # Ensure MSE arrays match the number of analysis frames
+    mse_embeddings = pad_or_trim(mse_embeddings, analysis_frames)
+    mse_posture = pad_or_trim(mse_posture, analysis_frames)
+    mse_voice = pad_or_trim(mse_voice, analysis_frames)
     
-    # Fill missing MSE values with 0.001
-    def pad_with_zeros(mse_array, total_frames):
-        if len(mse_array) < total_frames:
-            return np.pad(mse_array, (0, total_frames - len(mse_array)), 'constant', constant_values=0.001)
-        else:
-            return mse_array[:total_frames]
-
-    mse_embeddings = pad_with_zeros(mse_embeddings, total_frames)
-    mse_posture = pad_with_zeros(mse_posture, total_frames)
-    mse_voice = pad_with_zeros(mse_voice, total_frames)
-
     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)
+        analysis_frame = int(t * analysis_fps)
+        heatmap_frame = create_heatmap(analysis_frame, mse_embeddings, mse_posture, mse_voice, analysis_fps, analysis_frames, width)
         heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
         combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame
 
-    final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
+    final_clip = VideoClip(combine_video_and_heatmap, duration=video_duration)
     final_clip = final_clip.set_audio(video.audio)
 
     # Write the final video
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video_fps)
     
     # Close the video clips
     video.close()
@@ -300,6 +317,13 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
         return None
 
+def pad_or_trim(mse_array, target_length):
+    if len(mse_array) < target_length:
+        return np.pad(mse_array, (0, target_length - len(mse_array)), 'constant', constant_values=0)
+    else:
+        return mse_array[:target_length]
+
+
 
 # Function to create the correlation heatmap
 def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):