Update visualization.py

visualization.py

@@ -271,39 +271,39 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     ax.set_xticks([])
     plt.tight_layout()
     
-    if progress:
-        progress(0, desc="Generating video with heatmap")
-
     line = None
-    for frame_count in range(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)
-        
-        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))
-        
-        combined_frame = np.vstack((frame, heatmap_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)
-        
-        out.write(combined_frame)
-        if progress is not None:
-            try:
-                progress((frame_count + 1) / total_frames, desc="Generating video with heatmap")
-            except Exception as e:
-                print(f"Error updating progress: {str(e)}")
-    
+    try:
+        for frame_count in range(total_frames):
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
+            ret, frame = cap.read()
+            if not ret:
+                print(f"Failed to read frame {frame_count}")
+                break
+            
+            if line:
+                line.remove()
+            line = ax.axvline(x=frame_count, color='r', linewidth=2)
+            
+            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))
+            
+            combined_frame = np.vstack((frame, heatmap_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)
+            
+            out.write(combined_frame)
+
+            if progress is not None:
+                try:
+                    progress((frame_count + 1) / total_frames, desc="Generating video with heatmap")
+                except Exception as e:
+                    print(f"Error updating progress: {str(e)}")
+
     except Exception as e:
         print(f"Error in frame processing: {str(e)}")
         import traceback
@@ -315,6 +315,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     if os.path.exists(heatmap_video_path):
         print(f"Heatmap video created at: {heatmap_video_path}")
+        print(f"Heatmap video size: {os.path.getsize(heatmap_video_path)} bytes")
         return heatmap_video_path
     else:
         print(f"Failed to create heatmap video at: {heatmap_video_path}")