Update image_segmenter.py

image_segmenter.py

@@ -7,17 +7,17 @@ import spaces
 
 class ImageSegmenter:
     def __init__(self, model_type="yolov8s-seg") -> None:
-    # Don't initialize CUDA here
-    self.model_type = model_type
-    self.is_show_bounding_boxes = True
-    self.is_show_segmentation_boundary = False
-    self.is_show_segmentation = False
-    self.confidence_threshold = 0.5
-    self.cls_clr = {}
-    self.bb_thickness = 2
-    self.bb_clr = (255, 0, 0)
-    self.masks = {}
-    self.model = None  # Model will be loaded in predict
+        # Store parameters but don't initialize CUDA
+        self.model_type = model_type
+        self.is_show_bounding_boxes = True
+        self.is_show_segmentation_boundary = False
+        self.is_show_segmentation = False
+        self.confidence_threshold = 0.5
+        self.cls_clr = {}
+        self.bb_thickness = 2
+        self.bb_clr = (255, 0, 0)
+        self.masks = {}
+        self.model = None  # Model will be loaded in predict
 
     def get_cls_clr(self, cls_id):
         if cls_id in self.cls_clr:
@@ -30,18 +30,18 @@ class ImageSegmenter:
         self.cls_clr[cls_id] = (r, g, b)
         return (r, g, b)
 
-    @spaces.GPU(duration=30)  # Adjust duration based on your needs
+    @spaces.GPU
     def predict(self, image):            
-        # Load model if not loaded (will happen on first prediction)
+        # Load model if not loaded
         if self.model is None:
+            print("Loading YOLO model...")
             self.model = YOLO('models/' + self.model_type + '.pt')
-            self.model.to(self.device)
+            self.model.to('cuda')
+            print("Model loaded successfully")
 
         # params
         objects_data = [] 
         image = image.copy()
-        
-        # Run prediction
         predictions = self.model.predict(image)
 
         cls_ids = predictions[0].boxes.cls.cpu().numpy()
@@ -73,7 +73,7 @@ class ImageSegmenter:
                     if image.shape[:2] != seg_mask[id].shape[:2]:
                         colored_mask = cv2.resize(colored_mask, (image.shape[1], image.shape[0]))
 
-                    # filling the masked area with class color
+                    # filling the mased area with class color
                     masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=cls_clr)
                     image_overlay = masked.filled()                
                     image = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)