Update visualization.py

visualization.py

@@ -217,30 +217,13 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width, largest_cluster):
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
     frame_count = int(t * video_fps)
     
-    # Replace MSE values outside of the largest cluster with zeros
-    mask = (largest_cluster == 1)
-    mse_embeddings[~mask] = 0
-    mse_posture[~mask] = 0
-    mse_voice[~mask] = 0
-
-    # Check if all values are zero
-    if np.all(mse_embeddings == 0):
-        mse_embeddings_norm = mse_embeddings
-    else:
-        mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
-    
-    if np.all(mse_posture == 0):
-        mse_posture_norm = mse_posture
-    else:
-        mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
-    
-    if np.all(mse_voice == 0):
-        mse_voice_norm = mse_voice
-    else:
-        mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
+    # Normalize MSE values
+    mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
+    mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
+    mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
     
     combined_mse = np.zeros((3, total_frames))
     combined_mse[0] = mse_embeddings_norm
@@ -253,7 +236,7 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
     ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7) 
     ax.set_xticks([])
 
-    ax.axvline(x=frame_count, color='black', linewidth=2)
+    ax.axvline(x=frame_count, color='black', linewidth=3)
 
     plt.tight_layout(pad=0.5)
     
@@ -291,7 +274,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames, width, largest_cluster)
+        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames, width)
         heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
         combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame
@@ -326,4 +309,4 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     heatmap = sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
     plt.title('Correlation Heatmap of MSEs')
     plt.tight_layout()
-    return plt.gcf()
+    return plt.gcf()