Update visualization.py

visualization.py

@@ -321,30 +321,25 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
 
 # Function to create the correlation heatmap
 def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
-    mse_data = {
-        'Facial Features MSE': mse_embeddings,
-        'Body Posture MSE': mse_posture,
-        'Voice MSE': mse_voice
-    }
-    mse_df = pd.DataFrame(mse_data)
-    correlation_matrix = mse_df.corr()
+    data = np.vstack((mse_embeddings, mse_posture, mse_voice)).T
+    df = pd.DataFrame(data, columns=["Facial Features", "Body Posture", "Voice"])
+    corr = df.corr()
 
-    plt.figure(figsize=(8, 6))
-    sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
-    plt.title("Correlation Heatmap of MSEs")
-    plt.close()
+    plt.figure(figsize=(10, 8))
+    heatmap = sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
+    plt.title('Correlation Heatmap of MSEs')
+    plt.tight_layout()
     return plt.gcf()
 
-# Function to create the 3D scatter plot
 def plot_3d_scatter(mse_embeddings, mse_posture, mse_voice):
     fig = plt.figure(figsize=(10, 8))
     ax = fig.add_subplot(111, projection='3d')
-    ax.scatter(mse_posture, mse_embeddings, mse_voice, c='b', marker='o')
 
+    ax.scatter(mse_posture, mse_embeddings, mse_voice, c='r', marker='o')
     ax.set_xlabel('Body Posture MSE')
     ax.set_ylabel('Facial Features MSE')
     ax.set_zlabel('Voice MSE')
     ax.set_title('3D Scatter Plot of MSEs')
 
-    plt.close()
-    return fig
+    plt.tight_layout()
+    return plt.gcf()