Update visualization.py

visualization.py

@@ -202,59 +202,49 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 
 
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path):
-    from matplotlib.backends.backend_agg import FigureCanvasAgg
-    import cv2
-    # Open the video
     cap = cv2.VideoCapture(video_path)
     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))
     
-    # Create the output video writer
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height + 200))  # Additional 200 pixels for heatmap
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height + 200))
     
-    # Prepare the heatmap data
-    heatmap_data = np.vstack((mse_embeddings, mse_posture))
+    # Ensure heatmap data covers all frames
+    heatmap_data = np.zeros((2, total_frames))
+    heatmap_data[0, :len(mse_embeddings)] = mse_embeddings
+    heatmap_data[1, :len(mse_posture)] = mse_posture
     
-    # Create a figure for the heatmap
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(heatmap_data, aspect='auto', cmap='YlOrRd')
+    im = ax.imshow(heatmap_data, aspect='auto', cmap='YlOrRd', extent=[0, total_frames, 0, 2])
     ax.set_yticks([])
     ax.set_xticks([])
     plt.tight_layout()
     
-    frame_count = 0
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    
-    line = None  # Initialize line object
+    line = None
     
-    while True:
+    for frame_count in range(total_frames):
         ret, frame = cap.read()
         if not ret:
             break
         
-        # Update the heatmap with the current frame position
         if line:
-            line.remove()  # Remove the previous line
+            line.remove()
         line = ax.axvline(x=frame_count, color='r', linewidth=2)
         
-        # Convert the matplotlib figure to an image
         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))
         
-        # Combine the video frame and the heatmap
         combined_frame = np.vstack((frame, heatmap_img))
         
-        # Add timecode to the frame
         timecode = df['Timecode'][frame_count] if frame_count < len(df) else "End"
         cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
-        frame_count += 1
     
     cap.release()
     out.release()