Update visualization.py

visualization.py

@@ -216,134 +216,3 @@ 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
-    combined_mse[1] = mse_posture
-    combined_mse[2] = mse_voice
-
-    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)
-
-    # Find max values for each variable
-    max_embeddings = np.max(mse_embeddings)
-    max_posture = np.max(mse_posture)
-    max_voice = np.max(mse_voice)
-
-    def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, fps, total_frames, width):
-        frame_index = int(t * fps)
-        
-        # Create separate color maps for each variable
-        cmap_embeddings = plt.get_cmap('Reds')
-        cmap_posture = plt.get_cmap('Reds')
-        cmap_voice = plt.get_cmap('Reds')
-
-        # Normalize values within their own scales
-        norm_embeddings = mse_embeddings[frame_index] / max_embeddings
-        norm_posture = mse_posture[frame_index] / max_posture
-        norm_voice = mse_voice[frame_index] / max_voice
-
-        # Create color arrays for each variable
-        color_embeddings = cmap_embeddings(norm_embeddings)
-        color_posture = cmap_posture(norm_posture)
-        color_voice = cmap_voice(norm_voice)
-
-        # Create the heatmap frame
-        heatmap_height = 100  # Adjust as needed
-        heatmap_frame = np.zeros((heatmap_height, width, 4))
-
-        # Fill the heatmap sections
-        heatmap_frame[:heatmap_height//3, :] = color_embeddings
-        heatmap_frame[heatmap_height//3:2*heatmap_height//3, :] = color_posture
-        heatmap_frame[2*heatmap_height//3:, :] = color_voice
-
-        return (heatmap_frame * 255).astype(np.uint8)
-
-    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]))
-        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()