Update visualization.py

visualization.py

@@ -199,3 +199,63 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.tight_layout()
     plt.close()
     return fig
+
+
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_path):
+    from matplotlib.backends.backend_agg import FigureCanvasAgg
+    # 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))
+    
+    # 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
+    
+    # Prepare the heatmap data
+    heatmap_data = np.vstack((mse_embeddings, mse_posture, mse_voice))
+    
+    # Create a figure for the heatmap
+    fig, ax = plt.subplots(figsize=(width/100, 2))
+    im = ax.imshow(heatmap_data, aspect='auto', cmap='YlOrRd')
+    ax.set_yticks([])
+    ax.set_xticks([])
+    plt.tight_layout()
+    
+    frame_count = 0
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        # Update the heatmap with the current frame position
+        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
+        
+        # Remove the vertical line for the next iteration
+        ax.lines.pop()
+    
+    cap.release()
+    out.release()
+    plt.close(fig)
+    
+    return output_path