Update visualization.py

visualization.py

@@ -298,13 +298,24 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
-
-        if progress:
-            progress((frame_count + 1) / total_frames, desc="Generating video with heatmap")
+        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)}")
     
-    cap.release()
-    out.release()
-    plt.close(fig)
+    except Exception as e:
+        print(f"Error in frame processing: {str(e)}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        cap.release()
+        out.release()
+        plt.close(fig)
     
-    print(f"Heatmap video created at: {heatmap_video_path}")
-    return heatmap_video_path
+    if os.path.exists(heatmap_video_path):
+        print(f"Heatmap video created at: {heatmap_video_path}")
+        return heatmap_video_path
+    else:
+        print(f"Failed to create heatmap video at: {heatmap_video_path}")
+        return None