Update visualization.py

visualization.py

@@ -7,7 +7,7 @@ import seaborn as sns
 import numpy as np
 import pandas as pd
 import cv2
-from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip, VideoClip
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip, VideoClip, concatenate_videoclips
 from moviepy.video.fx.all import resize
 from PIL import Image, ImageDraw, ImageFont
 from matplotlib.patches import Rectangle
@@ -216,107 +216,45 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-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)
-    
-    # Normalize the 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.full((3, total_frames), np.nan)
-    combined_mse[0] = mse_embeddings_norm
-    combined_mse[1] = mse_posture_norm
-    combined_mse[2] = mse_voice_norm
-
-    # Create custom colormap
-    cdict = {
-        'red':   [(0.0,  0.5, 0.5), (1.0,  1.0, 1.0)],
-        'green': [(0.0,  0.5, 0.5), (1.0,  0.0, 0.0)],
-        'blue':  [(0.0,  0.5, 0.5), (1.0,  0.0, 0.0)],
-    }
-    custom_cmap = LinearSegmentedColormap('custom_cmap', segmentdata=cdict, N=256)
-    
-    fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 3], vmin=0, vmax=1)
-    ax.set_yticks([0.5, 1.5, 2.5])
-    ax.set_yticklabels(['Face', 'Posture', 'Voice'])
-    ax.set_xticks([])
-    plt.tight_layout()
 
-    def create_heatmap(t):
-        frame_count = int(t * video.fps)
-        
-        # Clear previous lines
-        if hasattr(ax, 'lines') and len(ax.lines) > 0:
-            ax.lines.pop(0)
-        ax.axvline(x=frame_count, color='blue', linewidth=2)
-        
-        canvas = FigureCanvasAgg(fig)
-        canvas.draw()
-        heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
-        heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
-        return heatmap_img
-
-    def add_timecode(frame, t):
-        seconds = t
-        timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
-        
-        pil_img = Image.fromarray(frame.astype('uint8'))
-        draw = ImageDraw.Draw(pil_img)
-        font = ImageFont.load_default()
-        draw.text((10, 30), f"Time: {timecode}", font=font, fill=(255, 255, 255))
-        
-        return np.array(pil_img)
-
-    heatmap_clip = VideoClip(create_heatmap, duration=video.duration)
-    heatmap_clip = heatmap_clip.resize(height=200)
-
-    def combine_video_and_heatmap(t):
-        video_frame = video.get_frame(t)
-        heatmap_frame = heatmap_clip.get_frame(t)
-        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()
+def create_heatmap(frame_time, mse_embeddings, mse_posture, mse_voice):
+    fig = Figure(figsize=(10, 1))
+    canvas = FigureCanvas(fig)
+    ax = fig.add_subplot(111)
+    time_index = int(frame_time)
     
-    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
+    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:
-        print(f"Failed to create heatmap video at: {heatmap_video_path}")
-        return None
+        mse_values = [0, 0, 0]  # Default values if the index is out of bounds
+
+    ax.barh(['Face', 'Posture', 'Voice'], mse_values, color=['navy', 'purple', 'green'])
+    ax.set_xlim(0, 1)  # Normalize the MSE values
+
+    canvas.draw()
+    img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+    img = img.reshape(fig.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
 
 
 # Function to create the correlation heatmap