Update visualization.py

visualization.py

@@ -226,36 +226,29 @@ def fill_with_zeros(mse_array, total_frames):
     return result
 
 def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, video_width):
-    frame_count = int(t * desired_fps)
-    window_size = min(600, total_frames)  # Increased window size for better context
-    start_frame = max(0, frame_count - window_size // 2)
-    end_frame = min(total_frames, start_frame + window_size)
-
-    combined_mse = np.array([
-        mse_embeddings[start_frame:end_frame],
-        mse_posture[start_frame:end_frame],
-        mse_voice[start_frame:end_frame]
-    ])
-
-    # Calculate global min and max for consistent scaling
-    vmin = 0
-    vmax = max(np.max(mse_embeddings), np.max(mse_posture), np.max(mse_voice))
-
-    fig, ax = plt.subplots(figsize=(video_width / 100, 0.4))  # Adjusted figure size
-    im = ax.imshow(combined_mse, aspect='auto', cmap='Reds', 
-                   extent=[start_frame/desired_fps, end_frame/desired_fps, 0, 3],
-                   vmin=vmin, vmax=vmax, interpolation='nearest')
+    fig, ax = plt.subplots(figsize=(video_width / 100, 0.4))
+
+    # Create the full heatmap for the entire video duration
+    combined_mse = np.array([mse_embeddings, mse_posture, mse_voice])
     
+    # Use pcolormesh for better performance with large datasets
+    im = ax.pcolormesh(np.arange(total_frames) / desired_fps, [0, 1, 2], combined_mse, 
+                       cmap='Reds', vmin=0, vmax=np.max(combined_mse))
+
+    ax.set_ylim(0, 3)
     ax.set_yticks([0.5, 1.5, 2.5])
     ax.set_yticklabels(['Face', 'Posture', 'Voice'], fontsize=7)
     
+    # Set x-axis to show full video duration
+    ax.set_xlim(0, total_frames / desired_fps)
+    
     # Add vertical line for current time
     current_time = t
     ax.axvline(x=current_time, color='black', linewidth=2)
 
     # Set x-axis ticks and labels
-    ax.set_xticks([start_frame/desired_fps, current_time, end_frame/desired_fps])
-    ax.set_xticklabels([f'{start_frame/desired_fps:.2f}', f'{current_time:.2f}', f'{end_frame/desired_fps:.2f}'], fontsize=6)
+    ax.set_xticks([0, current_time, total_frames / desired_fps])
+    ax.set_xticklabels(['0:00', f'{current_time:.2f}', f'{total_frames / desired_fps:.2f}'], fontsize=6)
 
     plt.tight_layout(pad=0.5)
     
@@ -280,15 +273,15 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     width, height = video.w, video.h
     total_frames = int(video.duration * desired_fps)
 
-    # Interpolate MSE values to match the desired fps
-    def interpolate_mse(mse_array):
-        original_indices = np.linspace(0, total_frames - 1, len(mse_array))
-        new_indices = np.arange(total_frames)
-        return np.interp(new_indices, original_indices, mse_array)
+    # Ensure MSE arrays have the same length as total_frames
+    def pad_mse_array(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)
+        return mse_array[:total_frames]
 
-    mse_embeddings = interpolate_mse(mse_embeddings)
-    mse_posture = interpolate_mse(mse_posture)
-    mse_voice = interpolate_mse(mse_voice)
+    mse_embeddings = pad_mse_array(mse_embeddings, total_frames)
+    mse_posture = pad_mse_array(mse_posture, total_frames)
+    mse_voice = pad_mse_array(mse_voice, total_frames)
 
     def combine_video_and_heatmap(t):
         original_frame = int(t * video.fps)