Update visualization.py

visualization.py

@@ -216,56 +216,60 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.tight_layout()
     plt.close()
     return fig
-
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, analysis_fps, video_width):
-    video_frame_count = int(t * video_fps)
-    analysis_frame_count = int(t * analysis_fps)
-    
-    # Ensure we don't go out of bounds
-    analysis_frame_count = min(analysis_frame_count, len(mse_embeddings) - 1)
-    
-    # Normalize MSE values
-    def safe_normalize(arr):
-        min_val, max_val = np.min(arr), np.max(arr)
-        if min_val == max_val:
-            return np.zeros_like(arr)
-        return (arr - min_val) / (max_val - min_val)
-
-    mse_embeddings_norm = safe_normalize(mse_embeddings)
-    mse_posture_norm = safe_normalize(mse_posture)
-    mse_voice_norm = safe_normalize(mse_voice)
-    
-    # Create heatmap data
-    heatmap_width = int(video_width / 240 * 100)  # Adjust this multiplier as needed
-    combined_mse = np.zeros((3, heatmap_width))
-    
-    # Map analysis frames to heatmap width
-    for i in range(heatmap_width):
-        frame_index = int(i * len(mse_embeddings) / heatmap_width)
-        combined_mse[0, i] = mse_embeddings_norm[frame_index]
-        combined_mse[1, i] = mse_posture_norm[frame_index]
-        combined_mse[2, i] = mse_voice_norm[frame_index]
-
-    fig, ax = plt.subplots(figsize=(video_width / 240, 0.5))
-    ax.imshow(combined_mse, aspect='auto', cmap='Reds', vmin=0, vmax=1, extent=[0, heatmap_width, 0, 3])
-    ax.set_yticks([0.5, 1.5, 2.5])
-    ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7) 
-    ax.set_xticks([])
-    
-    # Calculate the position of the vertical line
-    line_pos = (video_frame_count / video_fps) / (len(mse_embeddings) / analysis_fps) * heatmap_width
-    ax.axvline(x=line_pos, 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, analysis_fps, largest_cluster):
+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):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
     os.makedirs(output_folder, exist_ok=True)
@@ -280,30 +284,19 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     # Get video properties
     width, height = video.w, video.h
-    video_duration = video.duration
-    video_fps = video.fps
-    
-    # Calculate the number of analysis frames
-    analysis_frames = int(video_duration * analysis_fps)
-    
-    # Ensure MSE arrays match the number of analysis frames
-    mse_embeddings = pad_or_trim(mse_embeddings, analysis_frames)
-    mse_posture = pad_or_trim(mse_posture, analysis_frames)
-    mse_voice = pad_or_trim(mse_voice, analysis_frames)
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        analysis_frame = int(t * analysis_fps)
-        heatmap_frame = create_heatmap(analysis_frame, mse_embeddings, mse_posture, mse_voice, analysis_fps, analysis_frames, width)
-        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
+        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)))
         combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame
 
-    final_clip = VideoClip(combine_video_and_heatmap, duration=video_duration)
+    final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
     final_clip = final_clip.set_audio(video.audio)
 
     # Write the final video
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video_fps)
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
     
     # Close the video clips
     video.close()
@@ -317,12 +310,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
         return None
 
-def pad_or_trim(mse_array, target_length):
-    if len(mse_array) < target_length:
-        return np.pad(mse_array, (0, target_length - len(mse_array)), 'constant', constant_values=0)
-    else:
-        return mse_array[:target_length]
-
 
 
 # Function to create the correlation heatmap