Update visualization.py

visualization.py

@@ -11,12 +11,22 @@ 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))
@@ -67,7 +77,6 @@ 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')
 
@@ -127,7 +136,6 @@ 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))
@@ -147,7 +155,6 @@ 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))
@@ -192,7 +199,6 @@ 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)
@@ -243,6 +249,40 @@ def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voi
 
     return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered
 
+def normalize_mse(mse):
+    return mse / np.max(mse) if np.max(mse) > 0 else mse
+
+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
+
+def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
+    frame_index = int(t * fps)
+    
+    # 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):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
@@ -275,6 +315,11 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         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_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
 
@@ -296,46 +341,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 create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
-    # 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)
-
-    # Debug prints
-    print(f"mse_embeddings_norm shape: {mse_embeddings_norm.shape}")
-    print(f"mse_posture_norm shape: {mse_posture_norm.shape}")
-    print(f"mse_voice_norm shape: {mse_voice_norm.shape}")
-
-    # Ensure combined_mse has the correct shape
-    combined_mse = np.zeros((total_frames, width))
-
-    # Adjust shapes and pad with zeros if necessary
-    mse_embeddings_norm = pad_or_trim_array(mse_embeddings_norm, width)
-    mse_posture_norm = pad_or_trim_array(mse_posture_norm, width)
-    mse_voice_norm = pad_or_trim_array(mse_voice_norm, width)
-
-    combined_mse[0] = mse_embeddings_norm
-    # Assuming you combine posture and voice MSEs similarly
-    combined_mse[1] = mse_posture_norm
-    combined_mse[2] = mse_voice_norm
-
-    # Return or use combined_mse as needed
-    return combined_mse
-
-def normalize_mse(mse):
-    # Your normalization logic here
-    return mse / np.max(mse)
-
-def pad_or_trim_array(arr, target_length):
-    if len(arr) > target_length:
-        # Trim the array
-        return arr[:target_length]
-    elif len(arr) < target_length:
-        # Pad the array with zeros
-        return np.pad(arr, (0, target_length - len(arr)), 'constant')
-    return arr
-
 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"])