Update visualization.py

visualization.py

@@ -216,3 +216,113 @@ 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, 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 / 300, 0.4))
+    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=2)
+
+    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)
+
+    def fill_with_previous_values(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
+        for i in range(1, total_frames):
+            if result[i] == 0:
+                result[i] = result[i-1]
+        return result
+
+    # Fill gaps with previous values
+    mse_embeddings = fill_with_previous_values(mse_embeddings, total_frames)
+    mse_posture = fill_with_previous_values(mse_posture, total_frames)
+    mse_voice = fill_with_previous_values(mse_voice, total_frames)
+
+    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)
+        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
+        
+        # Convert heatmap frame to RGB if it's RGBA
+        if heatmap_frame_resized.shape[2] == 4:
+            heatmap_frame_resized = cv2.cvtColor(heatmap_frame_resized, cv2.COLOR_RGBA2RGB)
+        
+        # Ensure both frames have the same number of channels
+        if video_frame.shape[2] != heatmap_frame_resized.shape[2]:
+            if video_frame.shape[2] == 3:
+                heatmap_frame_resized = heatmap_frame_resized[:, :, :3]  # Use only RGB channels
+            else:
+                video_frame = cv2.cvtColor(video_frame, cv2.COLOR_RGB2RGBA)
+        
+        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_audio(video.audio)
+
+    # Write the final video
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.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
+    df = pd.DataFrame(data, columns=["Facial Features", "Body Posture", "Voice"])
+    corr = df.corr()
+
+    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()