Update visualization.py

visualization.py

@@ -231,26 +231,24 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     plt.tight_layout()
     return plt.gcf()
 
-def plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df, title="Stacked MSE Heatmaps"):
-    plt.figure(figsize=(20, 9), dpi=300)
-    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(20, 9), sharex=True)
+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})
 
     # Face heatmap
-    sns.heatmap(mse_face.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax1)
-    ax1.set_yticks([0.5])
-    ax1.set_yticklabels(['Face'], rotation=0, va='center')
-    ax1.set_xticks([])
+    sns.heatmap(mse_face.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax1, xticklabels=False, yticklabels=False)
+    ax1.set_ylabel('Face', rotation=0, ha='right', va='center')
+    ax1.yaxis.set_label_coords(-0.01, 0.5)
 
     # Posture heatmap
-    sns.heatmap(mse_posture.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax2)
-    ax2.set_yticks([0.5])
-    ax2.set_yticklabels(['Posture'], rotation=0, va='center')
-    ax2.set_xticks([])
+    sns.heatmap(mse_posture.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax2, xticklabels=False, yticklabels=False)
+    ax2.set_ylabel('Posture', rotation=0, ha='right', va='center')
+    ax2.yaxis.set_label_coords(-0.01, 0.5)
 
     # Voice heatmap
-    sns.heatmap(mse_voice.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax3)
-    ax3.set_yticks([0.5])
-    ax3.set_yticklabels(['Voice'], rotation=0, va='center')
+    sns.heatmap(mse_voice.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax3, yticklabels=False)
+    ax3.set_ylabel('Voice', rotation=0, ha='right', va='center')
+    ax3.yaxis.set_label_coords(-0.01, 0.5)
 
     # Set x-axis ticks to timecodes for the bottom subplot
     num_ticks = min(60, len(mse_voice))
@@ -259,6 +257,13 @@ def plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df, title="Stack
     ax3.set_xticks(tick_locations)
     ax3.set_xticklabels(tick_labels, rotation=90, ha='center', va='top')
 
+    # Remove spines
+    for ax in [ax1, ax2, ax3]:
+        ax.spines['top'].set_visible(False)
+        ax.spines['right'].set_visible(False)
+        ax.spines['bottom'].set_visible(False)
+        ax.spines['left'].set_visible(False)
+
     plt.suptitle(title)
     plt.tight_layout()
     plt.close()