Update visualization.py

visualization.py

@@ -1,6 +1,6 @@
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
-from matplotlib.backends.backend_agg import FigureCanvasAgg
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 import matplotlib.colors as mcolors
 from matplotlib.colors import LinearSegmentedColormap
 import seaborn as sns
@@ -217,43 +217,92 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_heatmap(frame_time, mse_embeddings, mse_posture, mse_voice):
-    fig, ax = plt.subplots(figsize=(10, 1))
-    time_index = int(frame_time)
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice):
+    frame_count = int(t * video.fps)
     
-    if time_index < len(mse_embeddings) and time_index < len(mse_posture) and time_index < len(mse_voice):
-        mse_values = [mse_embeddings[time_index], mse_posture[time_index], mse_voice[time_index]]
-    else:
-        mse_values = [0, 0, 0]  # Default values if the index is out of bounds
+    # 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
 
-    ax.barh(['Face', 'Posture', 'Voice'], mse_values, color=['navy', 'purple', 'green'])
-    ax.set_xlim(0, 1)  # Normalize the MSE values
+    fig, ax = plt.subplots(figsize=(10, 2))
+    ax.imshow(combined_mse, aspect='auto', cmap='coolwarm', vmin=0, vmax=1, extent=[0, total_frames, 0, 3])
+    ax.set_yticks([0.5, 1.5, 2.5])
+    ax.set_yticklabels(['Face', 'Posture', 'Voice'])
+    ax.set_xticks([])
 
-    canvas = FigureCanvasAgg(fig)
+    ax.axvline(x=frame_count, color='blue', linewidth=2)
+
+    canvas = FigureCanvas(fig)
     canvas.draw()
-    img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
-    img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+    heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
     plt.close(fig)
-    return img
-
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, fps, largest_cluster):
-    original_clip = VideoFileClip(video_path)
-    duration = original_clip.duration
-    heatmap_clips = []
-
-    for t in np.arange(0, duration, 1.0 / fps):
-        heatmap_img = create_heatmap(t, mse_embeddings, mse_posture, mse_voice)
-        heatmap_img_bgr = cv2.cvtColor(heatmap_img, cv2.COLOR_RGB2BGR)
-        heatmap_filename = os.path.join(output_folder, f"heatmap_{int(t * fps)}.png")
-        cv2.imwrite(heatmap_filename, heatmap_img_bgr)
-        heatmap_clips.append(ImageClip(heatmap_filename).set_duration(1.0 / fps).set_start(t).resize(height=100))
-
-    heatmap_clip = concatenate_videoclips(heatmap_clips, method="compose")
-    final_clip = CompositeVideoClip([original_clip, heatmap_clip.set_position(('center', 'bottom'))])
-    heatmap_video_path = os.path.join(output_folder, "heatmap_video.mp4")
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', fps=fps, audio_codec='aac')
-
-    return heatmap_video_path
+    return heatmap_img
+
+# Function to create video with heatmap
+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 all MSE arrays have the same length as total_frames
+    mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
+                               np.arange(len(mse_embeddings)), mse_embeddings)
+    mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
+                            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)
+    
+    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 combine_video_and_heatmap(t):
+        video_frame = video.get_frame(t)
+        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice)
+        combined_frame = np.vstack((video_frame, heatmap_frame))
+        return add_timecode(combined_frame, t)
+
+    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