Update visualization.py

visualization.py

@@ -267,16 +267,14 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     ax.set_xticks([])
     plt.tight_layout()
 
-    def make_frame(t):
+    def make_frame(frame):
+        t = frame / video.fps if isinstance(frame, np.ndarray) else frame
         frame_count = int(t * video.fps)
         
-        # Get the frame from the original video
-        frame = video.get_frame(t)
-        
         # Update heatmap
         if hasattr(ax, 'line'):
-            ax.line.remove()
-        ax.line = ax.axvline(x=frame_count, color='r', linewidth=2)
+            ax.lines.pop(0)
+        ax.axvline(x=frame_count, color='r', linewidth=2)
         
         canvas = FigureCanvasAgg(fig)
         canvas.draw()
@@ -291,7 +289,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         seconds = t
         timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
         
-        pil_img = Image.fromarray(combined_frame)
+        pil_img = Image.fromarray(combined_frame.astype('uint8'))
         draw = ImageDraw.Draw(pil_img)
         font = ImageFont.load_default()
         draw.text((10, 30), f"Time: {timecode}", font=font, fill=(255, 255, 255))