Update visualization.py

visualization.py

@@ -11,22 +11,12 @@ from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, Ima
 from moviepy.video.fx.all import resize
 from PIL import Image, ImageDraw, ImageFont
 from matplotlib.patches import Rectangle
+from utils import seconds_to_timecode
+from anomaly_detection import determine_anomalies
 from scipy import interpolate
+import gradio as gr
 import os
 
-# Utility functions
-def seconds_to_timecode(seconds):
-    hours = seconds // 3600
-    minutes = (seconds % 3600) // 60
-    seconds = seconds % 60
-    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}"
-
-def determine_anomalies(values, threshold):
-    mean = np.mean(values)
-    std = np.std(values)
-    anomalies = np.where(values > mean + threshold * std)[0]
-    return anomalies
-
 def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_threshold=4):
     plt.figure(figsize=(16, 8), dpi=300)
     fig, ax = plt.subplots(figsize=(16, 8))
@@ -77,6 +67,7 @@ def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_thre
             ax.plot(segment_df['Seconds'], mean, color=color, linewidth=0.5)
             ax.fill_between(segment_df['Seconds'], mean - std, mean + std, color=color, alpha=0.1)
 
+    # Rest of the function remains the same
     median = np.median(mse_values)
     ax.axhline(y=median, color='black', linestyle='--', label='Median Baseline')
 
@@ -136,6 +127,7 @@ def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_thre
     plt.close()
     return fig, anomaly_frames
 
+
 def plot_mse_histogram(mse_values, title, anomaly_threshold, color='blue'):
     plt.figure(figsize=(16, 3), dpi=300)
     fig, ax = plt.subplots(figsize=(16, 3))
@@ -155,6 +147,7 @@ def plot_mse_histogram(mse_values, title, anomaly_threshold, color='blue'):
     plt.close()
     return fig
 
+
 def plot_mse_heatmap(mse_values, title, df):
     plt.figure(figsize=(20, 3), dpi=300)
     fig, ax = plt.subplots(figsize=(20, 3))
@@ -199,6 +192,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     # Create a new dataframe for posture data
     posture_df = pd.DataFrame({'Frame': posture_frames, 'Score': posture_scores})
 
+
     posture_df = posture_df.merge(df[['Frame', 'Seconds']], on='Frame', how='inner')
 
     ax.scatter(posture_df['Seconds'], posture_df['Score'], color=color, alpha=0.3, s=5)
@@ -223,68 +217,37 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames):
-    # Ensure most_frequent_person_frames is a list
-    if not isinstance(most_frequent_person_frames, (list, np.ndarray)):
-        most_frequent_person_frames = [most_frequent_person_frames]
-
-    # Ensure df and mse arrays have the same length
-    min_length = min(len(df), len(mse_embeddings), len(mse_posture), len(mse_voice))
-    df = df.iloc[:min_length].copy()
-    mse_embeddings = mse_embeddings[:min_length]
-    mse_posture = mse_posture[:min_length]
-    mse_voice = mse_voice[:min_length]
-
-    # Create a mask for the most frequent person frames
-    mask = df['Frame'].isin(most_frequent_person_frames)
-
-    # Pad mask to match the length of the video frames
-    padded_mask = np.zeros(len(mse_embeddings), dtype=bool)
-    padded_mask[:len(mask)] = mask
-
-    # Apply the mask to filter the MSE arrays
-    mse_embeddings_filtered = np.where(padded_mask, mse_embeddings, 0)
-    mse_posture_filtered = np.where(padded_mask, mse_posture, 0)
-    mse_voice_filtered = np.where(padded_mask, mse_voice, 0)
-
-    return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered
+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
 
-def normalize_mse(mse):
-    return mse / np.max(mse) if np.max(mse) > 0 else mse
+    fig, ax = plt.subplots(figsize=(video_width / 250, 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([])
 
-def pad_or_trim_array(arr, target_length):
-    if len(arr) > target_length:
-        return arr[:target_length]
-    elif len(arr) < target_length:
-        return np.pad(arr, (0, target_length - len(arr)), 'constant')
-    return arr
+    ax.axvline(x=frame_count, color='black', linewidth=3)
 
-def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
-    frame_index = min(int(t * fps), len(mse_embeddings_filtered) - 1)
+    plt.tight_layout(pad=0.5)
     
-    # Normalize the MSE values
-    mse_embeddings_norm = normalize_mse(mse_embeddings_filtered)
-    mse_posture_norm = normalize_mse(mse_posture_filtered)
-    mse_voice_norm = normalize_mse(mse_voice_filtered)
-
-    # Ensure all arrays have the correct length
-    mse_embeddings_norm = pad_or_trim_array(mse_embeddings_norm, total_frames)
-    mse_posture_norm = pad_or_trim_array(mse_posture_norm, total_frames)
-    mse_voice_norm = pad_or_trim_array(mse_voice_norm, total_frames)
-
-    # Create a 3D array for the heatmap (height, width, channels)
-    heatmap_height = 3  # Assuming you want 3 rows in your heatmap
-    heatmap_frame = np.zeros((heatmap_height, width, 3), dtype=np.uint8)
-
-    # Fill the heatmap frame with color based on MSE values
-    heatmap_frame[0, :, 0] = (mse_embeddings_norm[frame_index] * 255).astype(np.uint8)  # Red channel for facial features
-    heatmap_frame[1, :, 1] = (mse_posture_norm[frame_index] * 255).astype(np.uint8)    # Green channel for body posture
-    heatmap_frame[2, :, 2] = (mse_voice_norm[frame_index] * 255).astype(np.uint8)      # Blue channel for voice
-
-    return heatmap_frame
-    
-
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, most_frequent_person_frames):
+    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)
@@ -308,32 +271,19 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
                             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)
-
-    print(f"Total frames: {total_frames}")
-    print(f"mse_embeddings length: {len(mse_embeddings)}")
-    print(f"mse_posture length: {len(mse_posture)}")
-    print(f"mse_voice length: {len(mse_voice)}")
-    
-    # Filter MSE arrays for the most frequent person frames
-    mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered = filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames)
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, video.fps, total_frames, width)
+        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]))
-        
-        # Ensure both frames have the same number of channels
-        if video_frame.shape[2] != heatmap_frame_resized.shape[2]:
-            heatmap_frame_resized = cv2.cvtColor(heatmap_frame_resized, cv2.COLOR_RGB2BGR)
-        
         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 using x264 codec
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps, preset='medium', ffmpeg_params=['-crf', '23'])
+    # 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()
@@ -347,6 +297,8 @@ 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
 
+
+# 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"])
@@ -357,4 +309,4 @@ def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     heatmap = sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
     plt.title('Correlation Heatmap of MSEs')
     plt.tight_layout()
-    return plt.gcf()
+    return plt.gcf()