Update visualization.py

visualization.py

@@ -7,6 +7,7 @@ import numpy as np
 import pandas as pd
 import cv2
 from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip
+from moviepy.video.fx.all import resize
 from matplotlib.patches import Rectangle
 from utils import seconds_to_timecode
 from anomaly_detection import determine_anomalies
@@ -220,20 +221,16 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     os.makedirs(output_folder, exist_ok=True)
     
     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}")
 
-    cap = cv2.VideoCapture(video_path)
-    original_fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    # Load the original video
+    video = VideoFileClip(video_path)
     
-    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)}")
+    # Get video properties
+    width, height = video.w, video.h
+    total_frames = int(video.duration * video.fps)
     
     # Ensure all MSE arrays have the same length as total_frames
     mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
@@ -269,84 +266,50 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     ax.set_yticklabels(['Face', 'Posture', 'Voice'])
     ax.set_xticks([])
     plt.tight_layout()
-    
-    line = None
-    try:
-        for frame_count in range(total_frames):
-            ret, frame = cap.read()
-            if not ret:
-                print(f"Failed to read frame {frame_count}")
-                break
-            
-            if line:
-                line.remove()
-            line = ax.axvline(x=frame_count, color='r', linewidth=2)
-            
-            canvas = FigureCanvasAgg(fig)
-            canvas.draw()
-            heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
-            heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
-            heatmap_img = cv2.resize(heatmap_img, (width, 200))
-            
-            # 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}"
-            
-            # 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:
-                    progress((frame_count + 1) / total_frames, desc="Generating video with heatmap")
-                except Exception as e:
-                    print(f"Error updating progress: {str(e)}")
-
-    except Exception as e:
-        print(f"Error in frame processing: {str(e)}")
-        import traceback
-        traceback.print_exc()
-    finally:
-        cap.release()
-        out.release()
-        plt.close(fig)
-
-    # Add audio to the video
-    try:
-        original_video = VideoFileClip(video_path)
-        heatmap_video = VideoFileClip(temp_video_path)
+
+    def make_frame(t):
+        frame_count = int(t * video.fps)
+        
+        # Get the frame from the original video
+        frame = video.get_frame(t)
         
-        # Resize the heatmap video to match the original video's duration
-        heatmap_video = heatmap_video.set_duration(original_video.duration)
+        # Update heatmap
+        if hasattr(ax, 'line'):
+            ax.line.remove()
+        ax.line = ax.axvline(x=frame_count, color='r', linewidth=2)
         
-        # Add the audio from the original video
-        final_video = heatmap_video.set_audio(original_video.audio)
+        canvas = FigureCanvasAgg(fig)
+        canvas.draw()
+        heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+        heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
+        heatmap_img = resize(ImageClip(heatmap_img), width=width).get_frame(0)
         
-        # Write the final video
-        final_video.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac')
+        # Combine frame with heatmap
+        combined_frame = np.vstack((frame, heatmap_img))
         
-        # Close the video clips
-        original_video.close()
-        heatmap_video.close()
-        final_video.close()
+        # Add timecode
+        seconds = t
+        timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
         
-        # Remove the temporary video file
-        os.remove(temp_video_path)
+        # We'll use PIL for text as it's easier to work with in this context
+        from PIL import Image, ImageDraw, ImageFont
+        pil_img = Image.fromarray(combined_frame)
+        draw = ImageDraw.Draw(pil_img)
+        font = ImageFont.load_default()
+        draw.text((10, 30), f"Time: {timecode}", font=font, fill=(255, 255, 255))
         
-    except Exception as e:
-        print(f"Error in adding audio to video: {str(e)}")
-        import traceback
-        traceback.print_exc()
+        return np.array(pil_img)
+
+    # Create the final video
+    final_clip = VideoFileClip(video_path).fl(make_frame)
+    final_clip = final_clip.set_duration(video.duration)
+
+    # Write the final video
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac')
+    
+    # Close the video clips
+    video.close()
+    final_clip.close()
     
     if os.path.exists(heatmap_video_path):
         print(f"Heatmap video created at: {heatmap_video_path}")
@@ -354,4 +317,22 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         return heatmap_video_path
     else:
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
-        return None
+        return None
+
+This version of the function:
+
+    Uses MoviePy for all video operations, avoiding OpenCV codec issues.
+    Creates the heatmap for each frame on-the-fly.
+    Combines the original video frame with the heatmap.
+    Adds the timecode to each frame.
+    Preserves the original audio.
+
+Make sure you have the necessary dependencies installed:
+
+pip install moviepy Pillow matplotlib numpy
+
+This approach should resolve the codec issues and create a video with the heatmap, original audio, and correct playback speed. If you still encounter any issues, please provide the new error message, and I'll be happy to help further.
+Claude does not have the ability to run the code it generates yet.
+Claude can make mistakes. Please double-check responses.
+
+