Update visualization.py

visualization.py

@@ -6,6 +6,7 @@ import seaborn as sns
 import numpy as np
 import pandas as pd
 import cv2
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip
 from matplotlib.patches import Rectangle
 from utils import seconds_to_timecode
 from anomaly_detection import determine_anomalies
@@ -216,11 +217,10 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster, progress=None):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
-    # Create output folder if it doesn't exist
     os.makedirs(output_folder, exist_ok=True)
     
-    # Define output path
     output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
+    temp_video_path = os.path.join(output_folder, 'temp_' + output_filename)
     heatmap_video_path = os.path.join(output_folder, output_filename)
     
     print(f"Heatmap video will be saved at: {heatmap_video_path}")
@@ -231,8 +231,8 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(heatmap_video_path, fourcc, original_fps, (width, height + 200))
+    fourcc = cv2.VideoWriter_fourcc(*'avc1')
+    out = cv2.VideoWriter(temp_video_path, fourcc, original_fps, (width, height + 200))
     print(f"VideoWriter initialized. FPS: {original_fps}, Size: {(width, height + 200)}")
     
     # Ensure all MSE arrays have the same length as total_frames
@@ -252,7 +252,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
 
-    # Custom colormap definition
     cdict = {
         'red':   [(0.0,  0.0, 0.0),
                   (1.0,  1.0, 1.0)],
@@ -274,7 +273,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     line = None
     try:
         for frame_count in range(total_frames):
-            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
             ret, frame = cap.read()
             if not ret:
                 print(f"Failed to read frame {frame_count}")
@@ -290,13 +288,23 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
             heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
             heatmap_img = cv2.resize(heatmap_img, (width, 200))
             
-            combined_frame = np.vstack((frame, heatmap_img))
+            # Convert frame from BGR to RGB
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            
+            # Combine RGB frame with RGB heatmap
+            combined_frame = np.vstack((frame_rgb, heatmap_img))
             
             seconds = frame_count / original_fps
             timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
-            cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
             
-            out.write(combined_frame)
+            # Add timecode to the combined frame (which is now in RGB)
+            combined_frame = cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), 
+                                         cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+            
+            # Convert back to BGR for OpenCV VideoWriter
+            combined_frame_bgr = cv2.cvtColor(combined_frame, cv2.COLOR_RGB2BGR)
+            
+            out.write(combined_frame_bgr)
 
             if progress is not None:
                 try:
@@ -312,6 +320,33 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         cap.release()
         out.release()
         plt.close(fig)
+
+    # Add audio to the video
+    try:
+        original_video = VideoFileClip(video_path)
+        heatmap_video = VideoFileClip(temp_video_path)
+        
+        # Resize the heatmap video to match the original video's duration
+        heatmap_video = heatmap_video.set_duration(original_video.duration)
+        
+        # Add the audio from the original video
+        final_video = heatmap_video.set_audio(original_video.audio)
+        
+        # Write the final video
+        final_video.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac')
+        
+        # Close the video clips
+        original_video.close()
+        heatmap_video.close()
+        final_video.close()
+        
+        # Remove the temporary video file
+        os.remove(temp_video_path)
+        
+    except Exception as e:
+        print(f"Error in adding audio to video: {str(e)}")
+        import traceback
+        traceback.print_exc()
     
     if os.path.exists(heatmap_video_path):
         print(f"Heatmap video created at: {heatmap_video_path}")