Update visualization.py

visualization.py

@@ -9,6 +9,7 @@ import cv2
 from matplotlib.patches import Rectangle
 from utils import seconds_to_timecode
 from anomaly_detection import determine_anomalies
+import gradio as gr
 
 def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_threshold=4):
     plt.figure(figsize=(16, 8), dpi=300)
@@ -206,7 +207,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps, largest_cluster):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps, largest_cluster, progress=gr.Progress()):
     # Filter the DataFrame to only include frames from the largest cluster
     df_largest_cluster = df[df['Cluster'] == largest_cluster]
     mse_embeddings = mse_embeddings[df['Cluster'] == largest_cluster]
@@ -254,6 +255,10 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     plt.tight_layout()
     
     line = None
+    
+    # Add progress tracking
+    progress(0.9, desc="Generating video with heatmap")
+    
     for frame_count in range(total_frames):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
@@ -262,7 +267,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
         
         if line:
             line.remove()
-        line = ax.axvline(x=frame_count, color='r', linewidth=2)
+        line = ax.axvline(x=frame_count, color='blue', linewidth=3)
         
         canvas = FigureCanvasAgg(fig)
         canvas.draw()
@@ -280,6 +285,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
         cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
+        
+        # Update progress
+        progress(0.9 + (0.1 * (frame_count + 1) / total_frames), desc="Generating video with heatmap")
     
     cap.release()
     out.release()