Update video_processing.py

video_processing.py

@@ -9,7 +9,7 @@ import math
 from face_analysis import get_face_embedding, cluster_faces, organize_faces_by_person, draw_facial_landmarks
 from pose_analysis import pose, calculate_posture_score, draw_pose_landmarks
 from anomaly_detection import anomaly_detection
-from visualization import plot_mse, plot_mse_histogram, plot_mse_heatmap, create_video_with_heatmap, fig_to_img
+from visualization import plot_mse, plot_mse_histogram, plot_mse_heatmap, create_video_with_heatmap
 from utils import frame_to_timecode
 import pandas as pd
 from facenet_pytorch import MTCNN
@@ -182,14 +182,9 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
 
             progress(0.9, "Generating video with heatmap")
 
-            # Convert heatmaps to images
-            mse_heatmap_embeddings_img = fig_to_img(mse_heatmap_embeddings)
-            mse_heatmap_posture_img = fig_to_img(mse_heatmap_posture)
-
             # Create video with heatmap
-            heatmap_video_path = os.path.join(output_folder, "video_with_heatmap.mp4")
-            create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatmap_posture_img, heatmap_video_path, desired_fps)
-
+            heatmap_video_path = os.path.join(output_folder, "heatmap_video.mp4")
+            heatmap_video_path = create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, heatmap_video_path, desired_fps, largest_cluster)
             
         except Exception as e:
             print(f"Error details: {str(e)}")