Update visualization.py

visualization.py

@@ -296,35 +296,45 @@ 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 create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
-    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))
-    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))
+def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
+    # Normalize the MSE values
+    mse_embeddings_norm = normalize_mse(mse_embeddings_filtered)
+    mse_posture_norm = normalize_mse(mse_posture_filtered)
+    mse_voice_norm = normalize_mse(mse_voice_filtered)
+
+    # Debug prints
+    print(f"mse_embeddings_norm shape: {mse_embeddings_norm.shape}")
+    print(f"mse_posture_norm shape: {mse_posture_norm.shape}")
+    print(f"mse_voice_norm shape: {mse_voice_norm.shape}")
+
+    # Ensure combined_mse has the correct shape
+    combined_mse = np.zeros((total_frames, width))
+
+    # Adjust shapes and pad with zeros if necessary
+    mse_embeddings_norm = pad_or_trim_array(mse_embeddings_norm, width)
+    mse_posture_norm = pad_or_trim_array(mse_posture_norm, width)
+    mse_voice_norm = pad_or_trim_array(mse_voice_norm, width)
+
     combined_mse[0] = mse_embeddings_norm
+    # Assuming you combine posture and voice MSEs similarly
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
 
-    fig, ax = plt.subplots(figsize=(video_width / 250, 0.6))
-    ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, total_frames, 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)
-
-    plt.tight_layout(pad=0.5)
-    
-    canvas = FigureCanvas(fig)
-    canvas.draw()
-    heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
-    heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
-    plt.close(fig)
-    return heatmap_img
+    # Return or use combined_mse as needed
+    return combined_mse
+
+def normalize_mse(mse):
+    # Your normalization logic here
+    return mse / np.max(mse)
+
+def pad_or_trim_array(arr, target_length):
+    if len(arr) > target_length:
+        # Trim the array
+        return arr[:target_length]
+    elif len(arr) < target_length:
+        # Pad the array with zeros
+        return np.pad(arr, (0, target_length - len(arr)), 'constant')
+    return arr
 
 def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     data = np.vstack((mse_embeddings, mse_posture, mse_voice)).T