Update app.py

app.py

@@ -19,11 +19,18 @@ img_seg = None
 depth_estimator = None
 
 def initialize_models():
+    """Loads models onto GPU if available, otherwise falls back to CPU."""
     global img_seg, depth_estimator
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     if img_seg is None:
-        img_seg = ImageSegmenter(model_type="yolov8s-seg")
+        print(f"🔹 Loading ImageSegmenter model on {device}...")
+        img_seg = ImageSegmenter(model_type="yolov8s-seg", device=device)
+
     if depth_estimator is None:
-        depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
+        print(f"🔹 Loading Depth Estimator model on {device}...")
+        depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256", device=device)
+
 
 def safe_gpu_decorator(func):
     """Custom decorator to handle GPU operations safely"""
@@ -42,22 +49,38 @@ def safe_gpu_decorator(func):
 @safe_gpu_decorator
 def process_image(image):
     try:
-        print("Starting image processing")
+        print("🚀 Starting image processing...")
         initialize_models()
-        
+
+        if torch.cuda.is_available():
+            print("✅ Using GPU for processing")
+            torch.set_default_tensor_type(torch.cuda.FloatTensor)
+        else:
+            print("⚠️ Using CPU for processing")
+
+        # Process image
         image = utils.resize(image)
         image_segmentation, objects_data = img_seg.predict(image)
         depthmap, depth_colormap = depth_estimator.make_prediction(image)
         dist_image = utils.draw_depth_info(image, depthmap, objects_data)
         objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
         plot_fig = display_pcd(objs_pcd)
+
         return image_segmentation, depth_colormap, dist_image, plot_fig
-    except Exception as e:
-        print(f"Error in process_image: {str(e)}")
+
+    except RuntimeError as e:
+        print(f"🚨 RuntimeError in process_image: {e}")
+
+        if "cuda" in str(e).lower():
+            print("⚠️ CUDA error detected. Switching to CPU mode.")
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
         import traceback
         print(traceback.format_exc())
         raise
 
+
+
 @safe_gpu_decorator
 def test_process_img(image):
     initialize_models()
@@ -102,41 +125,43 @@ def update_confidence_threshold(thres_val):
 @safe_gpu_decorator
 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)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model_dict = {
+        "Small - Better performance and less accuracy": ("midas_v21_small_256", "yolov8s-seg"),
+        "Medium - Balanced performance and accuracy": ("dpt_hybrid_384", "yolov8m-seg"),
+        "Large - Slow performance and high accuracy": ("dpt_large_384", "yolov8l-seg"),
+    }
+
+    midas_model, yolo_model = model_dict.get(model_type, ("midas_v21_small_256", "yolov8s-seg"))
+
+    print(f"🔹 Switching to models: YOLO={yolo_model}, MiDaS={midas_model} on {device}")
+
+    img_seg = ImageSegmenter(model_type=yolo_model, device=device)
+    depth_estimator = MonocularDepthEstimator(model_type=midas_model, device=device)
+
 
 def cancel():
     global CANCEL_PROCESSING
     CANCEL_PROCESSING = True
 
 if __name__ == "__main__":
-    # Try to initialize CUDA early to catch any issues
+        # Ensure CUDA is properly initialized
     try:
         if torch.cuda.is_available():
-            print("CUDA is available. Using GPU.")
-            # Test CUDA initialization
-            torch.cuda.init()
+            print(f"✅ CUDA is available: {torch.cuda.get_device_name(0)}")
             device = torch.device("cuda")
+            torch.cuda.empty_cache()  # Clear GPU cache
         else:
-            print("CUDA is not available. Using CPU.")
-            os.environ['CUDA_VISIBLE_DEVICES'] = ''
+            print("❌ No CUDA available. Falling back to CPU.")
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
             device = torch.device("cpu")
     except RuntimeError as e:
-        print(f"CUDA initialization failed: {e}")
-        print("Falling back to CPU mode")
-        os.environ['CUDA_VISIBLE_DEVICES'] = ''
+        print(f"🚨 CUDA initialization failed: {e}. Switching to CPU mode.")
+        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
         device = torch.device("cpu")
 
+
     with gr.Blocks() as my_app:
         # title
         gr.Markdown("<h1><center>Simultaneous Segmentation and Depth Estimation</center></h1>")