Update visualization.py

visualization.py

@@ -15,6 +15,7 @@ from anomaly_detection import determine_anomalies
 from scipy import interpolate
 import gradio as gr
 import os
+import logging
 
 def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_threshold=4):
     plt.figure(figsize=(16, 8), dpi=300)
@@ -302,8 +303,20 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
     final_clip = final_clip.set_audio(video.audio)
 
-    # Write the final video with progress bar
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps, progress_bar=progress)
+    # Define a logger to track progress
+    class ProgressBarLogger(logging.Logger):
+        def __init__(self, name, progress_callback):
+            super().__init__(name)
+            self.progress_callback = progress_callback
+
+        def debug(self, msg, *args, **kwargs):
+            if self.progress_callback:
+                self.progress_callback(float(msg.split('%')[0]) / 100, msg)
+
+    progress_logger = ProgressBarLogger('progress_logger', progress)
+
+    # Write the final video
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps, logger=progress_logger)
     
     # Close the video clips
     video.close()