Update visualization.py

visualization.py

@@ -234,7 +234,7 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
 
 def plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df, title="Combined MSE Heatmaps"):
     plt.figure(figsize=(20, 6), dpi=300)
-    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(20, 6), sharex=True, gridspec_kw={'height_ratios': [1, 1, 1.2], 'hspace': 0})
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(20, 8), sharex=True, gridspec_kw={'height_ratios': [1, 1, 1.2], 'hspace': 0})
 
     # Face heatmap
     sns.heatmap(mse_face.reshape(1, -1), cmap='Reds', cbar=False, ax=ax1, xticklabels=False, yticklabels=False)
@@ -290,7 +290,7 @@ def create_video_with_heatmap(video_path, mse_face, mse_posture, mse_voice, df,
         # Calculate the position of the vertical line
         line_pos = int(t / video.duration * video.w)
         
-        # Add the vertical line to the heatmap
+        # Add the vertical line to the heatmap only
         heatmap_with_line = heatmap_resized.copy()
         cv2.line(heatmap_with_line, (line_pos, 0), (line_pos, heatmap_height), (0, 0, 0), 2)