Update video_processing.py

video_processing.py

@@ -15,8 +15,6 @@ import pandas as pd
 from facenet_pytorch import MTCNN
 import torch
 import mediapipe as mp
-from voice_analysis import process_audio
-from pydub import AudioSegment
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
@@ -24,6 +22,7 @@ mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_fac
 mp_face_mesh = mp.solutions.face_mesh
 face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.8)
 
+
 def extract_frames(video_path, output_folder, desired_fps, progress_callback=None):
     os.makedirs(output_folder, exist_ok=True)
     clip = VideoFileClip(video_path)
@@ -46,6 +45,7 @@ def extract_frames(video_path, output_folder, desired_fps, progress_callback=Non
     clip.close()
     return frame_count, original_fps
 
+
 def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
     embeddings_by_frame = {}
     posture_scores_by_frame = {}
@@ -88,8 +88,8 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
         progress((i + 1) / len(frame_files), f"Processing frame {i + 1} of {len(frame_files)}")
 
     return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths, facial_landmarks_by_frame
-    
-    
+
+
 def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
     start_time = time.time()
     output_folder = "output"
@@ -124,7 +124,6 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
             frames_folder, aligned_faces_folder,
             frame_count,
             progress)
-    
 
         if not aligned_face_paths:
             raise ValueError("No faces were extracted from the video.")
@@ -155,45 +154,40 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
         try:
             X_posture = np.array([posture_scores_by_frame.get(frame, None) for frame in df['Frame']])
             X_posture = X_posture[X_posture != None].reshape(-1, 1)
-    
+
             if len(X_posture) == 0:
                 raise ValueError("No valid posture data found")
-    
+
             mse_embeddings, mse_posture = anomaly_detection(X_embeddings, X_posture)
-    
-            # Extract audio from video
-            video = AudioSegment.from_file(video_path, "mp4")
-            audio_path = os.path.join(temp_dir, "audio.wav")
-            video.export(audio_path, format="wav")
-    
-            # Process audio
-            most_frequent_voice, voice_features, voice_clusters = process_audio(audio_path, desired_fps, video_duration)
-    
-            if len(voice_features) < 2:
-                print("Not enough voice segments for processing. Skipping voice analysis.")
-                mse_voice = None
-                mse_plot_voice = None
-                mse_histogram_voice = None
-                mse_heatmap_voice = None
-                anomaly_segments_voice = None
-            else:
-                # Perform anomaly detection on voice
-                X_voice = np.array(most_frequent_voice)
-                mse_voice = anomaly_detection(X_voice, X_voice)
-    
-                # Generate plots for voice
-                mse_plot_voice, anomaly_segments_voice = plot_mse(df, mse_voice, "Voice",
-                                                                  color='green',
-                                                                  anomaly_threshold=anomaly_threshold)
-                mse_histogram_voice = plot_mse_histogram(mse_voice, "MSE Distribution: Voice",
-                                                         anomaly_threshold, color='green')
-                mse_heatmap_voice = plot_mse_heatmap(mse_voice, "Voice MSE Heatmap", df)
-    
+
+            progress(0.95, "Generating plots")
+            mse_plot_embeddings, anomaly_frames_embeddings = plot_mse(df, mse_embeddings, "Facial Features",
+                                                                      color=GRAPH_COLORS['facial_embeddings'],
+                                                                      anomaly_threshold=anomaly_threshold)
+
+            mse_histogram_embeddings = plot_mse_histogram(mse_embeddings, "MSE Distribution: Facial Features",
+                                                          anomaly_threshold, color=GRAPH_COLORS['facial_embeddings'])
+
+            mse_plot_posture, anomaly_frames_posture = plot_mse(df, mse_posture, "Body Posture",
+                                                                color=GRAPH_COLORS['body_posture'],
+                                                                anomaly_threshold=anomaly_threshold)
+
+            mse_histogram_posture = plot_mse_histogram(mse_posture, "MSE Distribution: Body Posture",
+                                                       anomaly_threshold, color=GRAPH_COLORS['body_posture'])
+
+            mse_heatmap_posture = plot_mse_heatmap(mse_posture, "Body Posture MSE Heatmap", df)
+
+            mse_heatmap_embeddings = plot_mse_heatmap(mse_embeddings, "Facial Features MSE Heatmap", df)
+
+            # Create video with heatmap
+            heatmap_video_path = os.path.join(output_folder, "video_with_heatmap.mp4")
+            create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, heatmap_video_path)
+
         except Exception as e:
             print(f"Error details: {str(e)}")
             import traceback
             traceback.print_exc()
-            return (f"Error in video processing: {str(e)}",) + (None,) * 20
+            return (f"Error in video processing: {str(e)}",) + (None,) * 15
 
         progress(1.0, "Preparing results")
         results = f"Number of persons detected: {num_clusters}\n\n"
@@ -251,18 +245,15 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
             mse_histogram_posture,
             mse_heatmap_embeddings,
             mse_heatmap_posture,
-            mse_voice,
-            mse_plot_voice,
-            mse_histogram_voice,
-            mse_heatmap_voice,
-            anomaly_segments_voice,
             face_samples["most_frequent"],
             anomaly_faces_embeddings,
             anomaly_frames_posture_images,
             aligned_faces_folder,
-            frames_folder
+            frames_folder,
+            heatmap_video_path
         )
 
+
 def is_frontal_face(landmarks, threshold=60):
     nose_tip = landmarks[4]
     left_chin = landmarks[234]
@@ -277,6 +268,7 @@ def is_frontal_face(landmarks, threshold=60):
     angle_degrees = math.degrees(angle)
     return abs(180 - angle_degrees) < threshold
 
+
 def save_person_data_to_csv(embeddings_by_frame, clusters, desired_fps, original_fps, output_folder, video_duration):
     person_data = {}
 
@@ -310,6 +302,7 @@ def save_person_data_to_csv(embeddings_by_frame, clusters, desired_fps, original
 
     return df, largest_cluster
 
+
 def get_all_face_samples(organized_faces_folder, output_folder, largest_cluster, max_samples=100):
     face_samples = {"most_frequent": [], "others": []}
     for cluster_folder in sorted(os.listdir(organized_faces_folder)):