Update visualization.py

visualization.py

@@ -218,104 +218,6 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-
-def fill_with_zeros(mse_array, total_frames):
-    result = np.zeros(total_frames)
-    indices = np.linspace(0, total_frames - 1, len(mse_array)).astype(int)
-    result[indices] = mse_array
-    return result
-
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, video_width):
-    fig, ax = plt.subplots(figsize=(video_width / 250, 0.4))
-
-    # Create the full heatmap for the entire video duration
-    combined_mse = np.array([mse_embeddings, mse_posture, mse_voice])
-    
-    # Use pcolormesh for better performance with large datasets
-    im = ax.pcolormesh(np.arange(total_frames) / desired_fps, [0, 1, 2], combined_mse, 
-                       cmap='Reds', vmin=0, vmax=np.max(combined_mse))
-
-    ax.set_ylim(0, 3)
-    ax.set_yticks([0.5, 1.5, 2.5])
-    ax.set_yticklabels(['Face', 'Posture', 'Voice'], fontsize=7)
-    
-    # Set x-axis to show full video duration
-    ax.set_xlim(0, total_frames / desired_fps)
-    
-    # Add vertical line for current time
-    current_time = t
-    ax.axvline(x=current_time, color='black', linewidth=2)
-
-    # Set x-axis ticks and labels
-    ax.set_xticks([0, current_time, total_frames / desired_fps])
-    ax.set_xticklabels(['0:00', f'{current_time:.2f}', f'{total_frames / desired_fps:.2f}'], fontsize=6)
-
-    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)
-    output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
-    heatmap_video_path = os.path.join(output_folder, output_filename)
-    print(f"Heatmap video will be saved at: {heatmap_video_path}")
-
-    # Load the original video
-    video = VideoFileClip(video_path)
-
-    # Get video properties
-    width, height = video.w, video.h
-    total_frames = int(video.duration * video.fps)
-
-    # Ensure MSE arrays align with original video frames
-    def align_mse_array(mse_array, original_fps, desired_fps, total_frames):
-        original_times = np.arange(len(mse_array)) / original_fps
-        desired_times = np.arange(total_frames) / desired_fps
-        interpolated_mse = np.interp(desired_times, original_times, mse_array)
-        return interpolated_mse
-
-    original_fps = len(mse_embeddings) / video.duration
-    mse_embeddings = align_mse_array(mse_embeddings, original_fps, desired_fps, total_frames)
-    mse_posture = align_mse_array(mse_posture, original_fps, desired_fps, total_frames)
-    mse_voice = align_mse_array(mse_voice, original_fps, desired_fps, total_frames)
-
-    def combine_video_and_heatmap(t):
-        frame_index = int(t * desired_fps)
-        video_frame = video.get_frame(t)
-        
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, 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 = final_clip.set_fps(desired_fps)
-
-    if video.audio is not None:
-        final_clip = final_clip.set_audio(video.audio.set_fps(desired_fps))
-
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=desired_fps)
-
-    # Close the video clips
-    video.close()
-    final_clip.close()
-
-    if os.path.exists(heatmap_video_path):
-        print(f"Heatmap video created at: {heatmap_video_path}")
-        print(f"Heatmap video size: {os.path.getsize(heatmap_video_path)} bytes")
-        return heatmap_video_path
-    else:
-        print(f"Failed to create heatmap video at: {heatmap_video_path}")
-        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
@@ -328,3 +230,36 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     plt.title('Correlation Heatmap of MSEs')
     plt.tight_layout()
     return plt.gcf()
+
+def plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df, title="Stacked MSE Heatmaps"):
+    plt.figure(figsize=(20, 9), dpi=300)
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(20, 9), sharex=True)
+
+    # Face heatmap
+    sns.heatmap(mse_face.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax1)
+    ax1.set_yticks([0.5])
+    ax1.set_yticklabels(['Face'], rotation=0, va='center')
+    ax1.set_xticks([])
+
+    # Posture heatmap
+    sns.heatmap(mse_posture.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax2)
+    ax2.set_yticks([0.5])
+    ax2.set_yticklabels(['Posture'], rotation=0, va='center')
+    ax2.set_xticks([])
+
+    # Voice heatmap
+    sns.heatmap(mse_voice.reshape(1, -1), cmap='YlOrRd', cbar=False, ax=ax3)
+    ax3.set_yticks([0.5])
+    ax3.set_yticklabels(['Voice'], rotation=0, va='center')
+
+    # Set x-axis ticks to timecodes for the bottom subplot
+    num_ticks = min(60, len(mse_voice))
+    tick_locations = np.linspace(0, len(mse_voice) - 1, num_ticks).astype(int)
+    tick_labels = [df['Timecode'].iloc[i] if i < len(df) else '' for i in tick_locations]
+    ax3.set_xticks(tick_locations)
+    ax3.set_xticklabels(tick_labels, rotation=90, ha='center', va='top')
+
+    plt.suptitle(title)
+    plt.tight_layout()
+    plt.close()
+    return fig