init models first

app.py

@@ -50,38 +50,34 @@ POSE_WEIGHTS = "models/ProbPose-s.pth"
 DEVICE = 'cuda:0'
 # DEVICE = 'cpu'
 
+# Global variables for models
+print("Initializing MMDetection detector...")
+det_model = init_detector(DET_CFG, DET_WEIGHTS, device=DEVICE)
+det_model.cfg = adapt_mmdet_pipeline(det_model.cfg)
+print("Detector initialized successfully!")
+
+print("Initializing MMPose estimator...")
+pose_model = init_pose_estimator(
+    POSE_CFG,
+    POSE_WEIGHTS,
+    device=DEVICE,
+    cfg_options=dict(model=dict(test_cfg=dict(output_heatmaps=True)))
+)
+print("Pose estimator initialized successfully!")
+
+pose_model.cfg.visualizer.radius = 4
+pose_model.cfg.visualizer.alpha = 0.8
+pose_model.cfg.visualizer.line_width = 2
+visualizer = VISUALIZERS.build(pose_model.cfg.visualizer)
+visualizer.set_dataset_meta(
+    pose_model.dataset_meta, skeleton_style='mmpose'
+)
+        
 @spaces.GPU
-def init_models():
-    """Initialize and return models for pose estimation"""
-    # Init detector
-    print("Initializing MMDetection detector...")
-    det_model = init_detector(DET_CFG, DET_WEIGHTS, device=DEVICE)
-    det_model.cfg = adapt_mmdet_pipeline(det_model.cfg)
-    print("Detector initialized successfully!")
-
-    # Init pose estimator
-    print("Initializing MMPose estimator...")
-    pose_model = init_pose_estimator(
-        POSE_CFG,
-        POSE_WEIGHTS,
-        device=DEVICE,
-        cfg_options=dict(model=dict(test_cfg=dict(output_heatmaps=True)))
-    )
-    
-    # Build visualizer
-    pose_model.cfg.visualizer.radius = 4
-    pose_model.cfg.visualizer.alpha = 0.8
-    pose_model.cfg.visualizer.line_width = 2
-    visualizer = VISUALIZERS.build(pose_model.cfg.visualizer)
-    visualizer.set_dataset_meta(
-        pose_model.dataset_meta, skeleton_style='mmpose'
-    )
-    print("Pose estimator initialized successfully!")
-    return det_model, pose_model, visualizer
-
-@spaces.GPU
-def process_frame(frame, det_model, pose_model, visualizer, bbox_thr=0.3, nms_thr=0.8, kpt_thr=0.3):
+def process_frame(frame, bbox_thr=0.3, nms_thr=0.8, kpt_thr=0.3):
     """Process a single frame with pose estimation"""
+    global det_model, pose_model, visualizer
+        
     processing_start = time.time()
 
     # Mirror the frame
@@ -173,14 +169,6 @@ class AsyncFrameProcessor:
         self.stop_event = threading.Event()
         self.new_frame_signal = threading.Event()
         
-        # Detector and pose estimator models
-        self.pose_model = None
-        self.det_model = None
-        self.visualizer = None
-        
-        # Initialize models using the standalone function
-        self.det_model, self.pose_model, self.visualizer = init_models()
-
         # Start background processing
         self._start_processing_thread()
     
@@ -213,20 +201,8 @@ class AsyncFrameProcessor:
 
                 print(f"Processing frame number: {frame_number}")
 
-                # Initialize models if not already done
-                if self.det_model is None or self.pose_model is None or self.visualizer is None:
-                    self.det_model, self.pose_model, self.visualizer = init_models()
-
-                print("Models initialized, starting frame processing...")
-
-                
-                # Process the frame using the standalone function
-                processed_frame = process_frame(
-                    frame_to_process,
-                    self.det_model,
-                    self.pose_model,
-                    self.visualizer
-                )
+                # Process the frame using the global models
+                processed_frame = process_frame(frame_to_process)
 
                 # Write frame number in the top left corner
                 processed_frame = cv2.putText(
@@ -341,7 +317,7 @@ with gr.Blocks(css=css) as demo:
 
 
 if __name__ == "__main__":
-
+    
     demo.launch(
         # server_name="0.0.0.0",
         # server_port=17860,