Update visualization.py

visualization.py

@@ -227,7 +227,7 @@ def fill_with_zeros(mse_array, total_frames):
 
 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(300, total_frames)
+    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)
 
@@ -237,17 +237,25 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total
         mse_voice[start_frame:end_frame]
     ])
 
-    fig, ax = plt.subplots(figsize=(video_width / 300, 0.4))
+    # 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=0, vmax=max(np.max(mse_embeddings), np.max(mse_posture), np.max(mse_voice)))
+                   vmin=vmin, vmax=vmax, interpolation='nearest')
+    
     ax.set_yticks([0.5, 1.5, 2.5])
-    ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7)
+    ax.set_yticklabels(['Face', 'Posture', 'Voice'], fontsize=7)
     
-    ax.axvline(x=t, color='black', linewidth=2)
+    # Add vertical line for current time
+    current_time = t
+    ax.axvline(x=current_time, color='black', linewidth=2)
 
-    ax.set_xticks([start_frame/desired_fps, t, end_frame/desired_fps])
-    ax.set_xticklabels([f'{start_frame/desired_fps:.2f}', f'{t:.2f}', f'{end_frame/desired_fps:.2f}'], fontsize=6)
+    # 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)
 
     plt.tight_layout(pad=0.5)
     
@@ -272,23 +280,22 @@ 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)
 
-    # Fill gaps with zeros
-    mse_embeddings = fill_with_zeros(mse_embeddings, total_frames)
-    mse_posture = fill_with_zeros(mse_posture, total_frames)
-    mse_voice = fill_with_zeros(mse_voice, total_frames)
+    # 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)
+
+    mse_embeddings = interpolate_mse(mse_embeddings)
+    mse_posture = interpolate_mse(mse_posture)
+    mse_voice = interpolate_mse(mse_voice)
 
     def combine_video_and_heatmap(t):
         original_frame = int(t * video.fps)
         video_frame = video.get_frame(original_frame / video.fps)
         
         heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, width)
-        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
-        
-        if video_frame.shape[2] != heatmap_frame_resized.shape[2]:
-            if video_frame.shape[2] == 3:
-                heatmap_frame_resized = heatmap_frame_resized[:, :, :3]
-            else:
-                video_frame = cv2.cvtColor(video_frame, cv2.COLOR_RGB2RGBA)
+        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, int(height * 0.2)))
         
         combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame