Update visualization.py

visualization.py

@@ -1,4 +1,5 @@
 import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 import matplotlib.colors as mcolors
 from matplotlib.colors import LinearSegmentedColormap
@@ -315,4 +316,35 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         return heatmap_video_path
     else:
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
-        return None
+        return None
+
+
+# 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()
+
+    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()
+    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.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