Update visualization.py

visualization.py

@@ -217,7 +217,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames):
+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
@@ -230,7 +230,7 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
 
-    fig, ax = plt.subplots(figsize=(10, 2))
+    fig, ax = plt.subplots(figsize=(video_width / 100, 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'])
@@ -272,8 +272,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     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)
-        combined_frame = np.vstack((video_frame, heatmap_frame))
+        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]))
+        combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame
 
     final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)