Update visualization.py

visualization.py

@@ -218,9 +218,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 
 
 def create_heatmap(frame_time, mse_embeddings, mse_posture, mse_voice):
-    fig = Figure(figsize=(10, 1))
-    canvas = FigureCanvas(fig)
-    ax = fig.add_subplot(111)
+    fig, ax = plt.subplots(figsize=(10, 1))
     time_index = int(frame_time)
     
     if time_index < len(mse_embeddings) and time_index < len(mse_posture) and time_index < len(mse_voice):
@@ -231,6 +229,7 @@ def create_heatmap(frame_time, mse_embeddings, mse_posture, mse_voice):
     ax.barh(['Face', 'Posture', 'Voice'], mse_values, color=['navy', 'purple', 'green'])
     ax.set_xlim(0, 1)  # Normalize the MSE values
 
+    canvas = FigureCanvas(fig)
     canvas.draw()
     img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
     img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
@@ -263,21 +262,27 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     df = pd.DataFrame(data, columns=["Facial Features", "Body Posture", "Voice"])
     corr = df.corr()
 
-    plt.figure(figsize=(10, 8))
+    plt.figure(figsize=(10, 8), dpi=300)
+
     heatmap = sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
     plt.title('Correlation Heatmap of MSEs')
     plt.tight_layout()
     return plt.gcf()
 
 def plot_3d_scatter(mse_embeddings, mse_posture, mse_voice):
-    fig = plt.figure(figsize=(10, 8))
+    fig = plt.figure()
+    plt.figure(figsize=(16, 8), dpi=300)
     ax = fig.add_subplot(111, projection='3d')
 
-    ax.scatter(mse_posture, mse_embeddings, mse_voice, c='r', marker='o')
+    # Scatter plot
+    ax.scatter(mse_posture, mse_embeddings, mse_voice, c=['purple']*len(mse_posture), label='Body Posture', alpha=0.6)
+    ax.scatter(mse_posture, mse_embeddings, mse_voice, c=['navy']*len(mse_embeddings), label='Facial Features', alpha=0.6)
+    ax.scatter(mse_posture, mse_embeddings, mse_voice, c=['green']*len(mse_voice), label='Voice', alpha=0.6)
+
     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.tight_layout()
-    return plt.gcf()
+    ax.legend()
+    plt.close(fig)
+    return fig