Update visualization.py

visualization.py

@@ -219,26 +219,35 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 
 
 
-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
+def fill_with_zeros(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
+    return result
+
+def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, video_width):
+    frame_count = int(t * desired_fps)
+    window_size = min(300, total_frames)
+    start_frame = max(0, frame_count - window_size // 2)
+    end_frame = min(total_frames, start_frame + window_size)
+
+    combined_mse = np.array([
+        mse_embeddings[start_frame:end_frame],
+        mse_posture[start_frame:end_frame],
+        mse_voice[start_frame:end_frame]
+    ])
 
     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])
+    im = ax.imshow(combined_mse, aspect='auto', cmap='Reds', 
+                   extent=[start_frame/desired_fps, end_frame/desired_fps, 0, 3],
+                   vmin=0, vmax=max(np.max(mse_embeddings), np.max(mse_posture), np.max(mse_voice)))
     ax.set_yticks([0.5, 1.5, 2.5])
-    ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7) 
-    ax.set_xticks([])
+    ax.set_yticklabels(['Voice', 'Posture', 'Face'], fontsize=7)
+    
+    ax.axvline(x=t, color='black', linewidth=2)
 
-    ax.axvline(x=frame_count, color='black', linewidth=2)
+    ax.set_xticks([start_frame/desired_fps, t, end_frame/desired_fps])
+    ax.set_xticklabels([f'{start_frame/desired_fps:.2f}', f'{t:.2f}', f'{end_frame/desired_fps:.2f}'], fontsize=6)
 
     plt.tight_layout(pad=0.5)
     
@@ -261,35 +270,23 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
 
     # 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)
+    total_frames = int(video.duration * desired_fps)
+
+    # Fill gaps with zeros
+    mse_embeddings = fill_with_zeros(mse_embeddings, total_frames)
+    mse_posture = fill_with_zeros(mse_posture, total_frames)
+    mse_voice = fill_with_zeros(mse_voice, total_frames)
 
     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]))
+        original_frame = int(t * video.fps)
+        video_frame = video.get_frame(original_frame / video.fps)
         
-        # Convert heatmap frame to RGB if it's RGBA
-        if heatmap_frame_resized.shape[2] == 4:
-            heatmap_frame_resized = cv2.cvtColor(heatmap_frame_resized, cv2.COLOR_RGBA2RGB)
+        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_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]:
             if video_frame.shape[2] == 3:
-                heatmap_frame_resized = heatmap_frame_resized[:, :, :3]  # Use only RGB channels
+                heatmap_frame_resized = heatmap_frame_resized[:, :, :3]
             else:
                 video_frame = cv2.cvtColor(video_frame, cv2.COLOR_RGB2RGBA)
         
@@ -297,10 +294,12 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         return combined_frame
 
     final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
-    final_clip = final_clip.set_audio(video.audio)
+    final_clip = final_clip.set_fps(desired_fps)
+
+    if video.audio is not None:
+        final_clip = final_clip.set_audio(video.audio.set_fps(desired_fps))
 
-    # Write the final video
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=desired_fps)
 
     # Close the video clips
     video.close()