Update visualization.py

visualization.py

@@ -216,60 +216,38 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.tight_layout()
     plt.close()
     return fig
-def create_combined_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, analysis_fps, video_width):
-    def plot_single_mse_heatmap(mse_values, height=1):
-        plt.figure(figsize=(video_width / 100, height), dpi=100)
-        fig, ax = plt.subplots(figsize=(video_width / 100, height))
-        
-        # Reshape MSE values to 2D array for heatmap
-        mse_2d = mse_values.reshape(1, -1)
-        
-        # Create heatmap
-        sns.heatmap(mse_2d, cmap='YlOrRd', cbar=False, ax=ax)
-        
-        # Remove all axes
-        ax.set_xticks([])
-        ax.set_yticks([])
-        ax.axis('off')
-        
-        plt.tight_layout(pad=0)
-        
-        # Convert plot to image
-        canvas = FigureCanvas(fig)
-        canvas.draw()
-        image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
-        image = image.reshape(canvas.get_width_height()[::-1] + (3,))
-        plt.close(fig)
-        return image
-
-    # Create individual heatmaps
-    face_heatmap = plot_single_mse_heatmap(mse_embeddings)
-    posture_heatmap = plot_single_mse_heatmap(mse_posture)
-    voice_heatmap = plot_single_mse_heatmap(mse_voice)
-
-    # Combine heatmaps vertically
-    combined_heatmap = np.vstack((face_heatmap, posture_heatmap, voice_heatmap))
-
-    # Add labels
-    label_height = 20
-    label_image = np.ones((label_height, combined_heatmap.shape[1], 3), dtype=np.uint8) * 255
-    cv2.putText(label_image, 'Face', (5, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    cv2.putText(label_image, 'Posture', (5, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    cv2.putText(label_image, 'Voice', (5, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-
-    combined_heatmap = np.vstack((label_image, combined_heatmap))
-
-    # Calculate position of vertical line
-    video_frame = int(t * video_fps)
-    total_analysis_frames = len(mse_embeddings)
-    line_pos = int((video_frame / (video_fps / analysis_fps)) * combined_heatmap.shape[1] / total_analysis_frames)
-
-    # Draw vertical line
-    cv2.line(combined_heatmap, (line_pos, 0), (line_pos, combined_heatmap.shape[0]), (0, 0, 0), 2)
-
-    return combined_heatmap
-
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, analysis_fps):
+
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
+    frame_count = int(t * video_fps)
+    
+    # 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
+    combined_mse[1] = mse_posture_norm
+    combined_mse[2] = mse_voice_norm
+
+    fig, ax = plt.subplots(figsize=(video_width / 240, 0.6))
+    ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, total_frames, 0, 3])
+    ax.set_yticks([0.5, 1.5, 2.5])
+    ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7) 
+    ax.set_xticks([])
+
+    ax.axvline(x=frame_count, color='black', linewidth=3)
+
+    plt.tight_layout(pad=0.5)
+    
+    canvas = FigureCanvas(fig)
+    canvas.draw()
+    heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+    heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
+    plt.close(fig)
+    return heatmap_img
+
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
     os.makedirs(output_folder, exist_ok=True)
@@ -284,11 +262,20 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     # Get video properties
     width, height = video.w, video.h
+    total_frames = int(video.duration * video.fps)
+    
+    # Ensure all MSE arrays have the same length as total_frames
+    mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
+                               np.arange(len(mse_embeddings)), mse_embeddings)
+    mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
+                            np.arange(len(mse_posture)), mse_posture)
+    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
+                          np.arange(len(mse_voice)), mse_voice)
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        heatmap_frame = create_combined_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, analysis_fps, width)
-        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, int(height * 0.2)))
+        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
 
@@ -311,7 +298,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         return None
 
 
-
 # Function to create the correlation heatmap
 def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     data = np.vstack((mse_embeddings, mse_posture, mse_voice)).T