Update app.py

app.py

@@ -24,6 +24,14 @@ from point_cloud_generator import display_pcd
 
 device = torch.device("cpu")  # Start in CPU mode
 
+# Global instances (can be reinitialized dynamically)
+img_seg = ImageSegmenter(model_type="yolov8s-seg")
+depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
+
+
+
+
+
 def initialize_gpu():
     """Ensure ZeroGPU assigns a GPU before initializing CUDA"""
     global device
@@ -91,13 +99,34 @@ def update_segmentation_options(options):
     img_seg.is_show_segmentation = True if 'Show Segmentation Region' in options else False
     img_seg.is_show_segmentation_boundary = True if 'Show Segmentation Boundary' in options else False
 
-def update_confidence_threshold(thres_val):
-    img_seg.confidence_threshold = thres_val/100
+#def update_confidence_threshold(thres_val):
+#    img_seg.confidence_threshold = thres_val/100
 
-@spaces.GPU  # Ensures YOLO + MiDaS get GPU access
-def model_selector(model_type):
-    global img_seg, depth_estimator
+def update_confidence_threshold(thres_val, img_seg_instance):
+    """Update confidence threshold in ImageSegmenter"""
+    img_seg_instance.confidence_threshold = thres_val
+    print(f"Confidence threshold updated to: {thres_val}")
+
+#@spaces.GPU  # Ensures YOLO + MiDaS get GPU access
+#def model_selector(model_type):
+#    global img_seg, depth_estimator
+
+#    if "Small - Better performance and less accuracy" == model_type:
+#        midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"
+#    elif "Medium - Balanced performance and accuracy" == model_type:
+#        midas_model, yolo_model = "dpt_hybrid_384", "yolov8m-seg"
+#    elif "Large - Slow performance and high accuracy" == model_type:
+#        midas_model, yolo_model = "dpt_large_384", "yolov8l-seg"
+#    else:
+#        midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"
+#
+#    img_seg = ImageSegmenter(model_type=yolo_model)
+#    depth_estimator = MonocularDepthEstimator(model_type=midas_model)
 
+
+# Updated model_selector to accept img_seg and depth_estimator instances
+@spaces.GPU  # Ensures YOLO + MiDaS get GPU access
+def model_selector(model_type, img_seg_instance, depth_estimator_instance):
     if "Small - Better performance and less accuracy" == model_type:
         midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"
     elif "Medium - Balanced performance and accuracy" == model_type:
@@ -107,8 +136,10 @@ def model_selector(model_type):
     else:
         midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"
 
-    img_seg = ImageSegmenter(model_type=yolo_model)
-    depth_estimator = MonocularDepthEstimator(model_type=midas_model)
+    # Reinitialize the provided instances with the selected model types
+    img_seg_instance.__init__(model_type=yolo_model)
+    depth_estimator_instance.__init__(model_type=midas_model)
+    print(f"Model updated: YOLO={yolo_model}, MiDaS={midas_model}")
 
     # START 
     # added for lens studio 
@@ -228,6 +259,8 @@ def get_bbox_from_mask(mask):
 
 @spaces.GPU
 def get_detection_data(image_data):
+    global img_seg, depth_estimator  # Still reference global instances
+    
     try:
         if isinstance(image_data, dict):
             nested_dict = image_data.get("image", {}).get("image", {})
@@ -263,7 +296,7 @@ def get_detection_data(image_data):
         print(f"Debug - Original image shape: {image.shape}")
 
         # Dynamically update model size and confidence threshold
-        model_selector(model_size, img_seg, depth_estimator)
+        model_selector(model_size, img_seg, depth_estimator)  # Pass the global instances
         update_confidence_threshold(confidence_threshold, img_seg)
         
         image_segmentation, objects_data = img_seg.predict(resized_image)