Update visualization.py

visualization.py

@@ -203,7 +203,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps):
+def create_video_with_heatmap(video_path, mse_heatmap_embeddings, mse_heatmap_posture, output_path, desired_fps):
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -211,63 +211,32 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
+    out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 400))
     
-    # Create custom colormap
-    cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
-        [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
-    
-    # Ensure heatmap data covers all frames
-    mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
-                               np.arange(len(mse_embeddings)), mse_embeddings)
-    mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
-                            np.arange(len(mse_posture)), mse_posture)
-    
-    # Normalize MSE values
-    mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
-    mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
-    
-    # Combine MSEs
-    combined_mse = np.zeros((2, total_frames, 3))
-    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, mse_embeddings_norm]).T  # RGB for facial
-    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T  # RGB for posture
-    
-    fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
-    ax.set_yticks([0.5, 1.5])
-    ax.set_yticklabels(['Face', 'Posture'])
-    ax.set_xticks([])
-    plt.tight_layout()
-    
-    line = None
-    frame_interval = int(original_fps / desired_fps)
-    
-    for frame_count in range(0, total_frames, frame_interval):
+    for frame_count in range(0, total_frames):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
         if not ret:
             break
         
-        if line:
-            line.remove()
-        line = ax.axvline(x=frame_count, color='r', linewidth=2)
+        heatmap_embeddings_img = np.frombuffer(mse_heatmap_embeddings, dtype='uint8')
+        heatmap_embeddings_img = heatmap_embeddings_img.reshape((mse_heatmap_embeddings.get_height(), mse_heatmap_embeddings.get_width(), 3))
         
-        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))
+        heatmap_posture_img = np.frombuffer(mse_heatmap_posture, dtype='uint8')
+        heatmap_posture_img = heatmap_posture_img.reshape((mse_heatmap_posture.get_height(), mse_heatmap_posture.get_width(), 3))
+
+        heatmap_combined_img = np.vstack((heatmap_embeddings_img, heatmap_posture_img))
         
-        combined_frame = np.vstack((frame, heatmap_img))
+        combined_frame = np.vstack((frame, heatmap_combined_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)
+        cv2.putText(combined_frame, f"Time: {timecode}", (10, height + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
     
     cap.release()
     out.release()
-    plt.close(fig)
+    plt.close('all')
     
     return output_path