Update visualization.py

visualization.py

@@ -203,7 +203,15 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
-def create_video_with_heatmap(video_path, mse_heatmap_embeddings, mse_heatmap_posture, output_path, desired_fps):
+def fig_to_img(fig):
+    canvas = FigureCanvasAgg(fig)
+    canvas.draw()
+    img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+    img = img.reshape(canvas.get_width_height()[::-1] + (3,))
+    plt.close(fig)
+    return img
+
+def create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatmap_posture_img, 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))
@@ -219,13 +227,7 @@ def create_video_with_heatmap(video_path, mse_heatmap_embeddings, mse_heatmap_po
         if not ret:
             break
         
-        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))
-        
-        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))
+        heatmap_combined_img = np.vstack((mse_heatmap_embeddings_img, mse_heatmap_posture_img))
         
         combined_frame = np.vstack((frame, heatmap_combined_img))