Update image_segmenter.py

image_segmenter.py

@@ -3,14 +3,12 @@ import numpy as np
 from ultralytics import YOLO
 import random
 import torch
+import spaces
 
 class ImageSegmenter:
     def __init__(self, model_type="yolov8s-seg") -> None:
-        
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.model = YOLO('models/'+ model_type +'.pt')
-        self.model.to(self.device)
-
+        self.model_type = model_type
+        self.device = 'cuda'  # ZeroGPU will always use CUDA
         self.is_show_bounding_boxes = True
         self.is_show_segmentation_boundary = False
         self.is_show_segmentation = False
@@ -23,7 +21,9 @@ class ImageSegmenter:
 
         # variables
         self.masks = {}
-
+        
+        # Model will be loaded in predict to work with ZeroGPU
+        self.model = None
 
     def get_cls_clr(self, cls_id):
         if cls_id in self.cls_clr:
@@ -36,15 +36,24 @@ class ImageSegmenter:
         self.cls_clr[cls_id] = (r, g, b)
         return (r, g, b)
 
+    @spaces.GPU(duration=30)  # Adjust duration based on your needs
     def predict(self, image):            
+        # Load model if not loaded (will happen on first prediction)
+        if self.model is None:
+            self.model = YOLO('models/' + self.model_type + '.pt')
+            self.model.to(self.device)
+
         # params
         objects_data = [] 
         image = image.copy()
+        
+        # Run prediction
         predictions = self.model.predict(image)
 
         cls_ids = predictions[0].boxes.cls.cpu().numpy()
         bounding_boxes = predictions[0].boxes.xyxy.int().cpu().numpy()        
         cls_conf = predictions[0].boxes.conf.cpu().numpy()
+        
         # segmentation
         if predictions[0].masks:
             seg_mask_boundary = predictions[0].masks.xy
@@ -56,8 +65,7 @@ class ImageSegmenter:
             cls_clr = self.get_cls_clr(cls)
 
             # draw filled segmentation region
-            if seg_mask.any() and  cls_conf[id] > self.confidence_threshold:
-
+            if seg_mask.any() and cls_conf[id] > self.confidence_threshold:
                 self.masks[id] = seg_mask[id]
                 
                 if self.is_show_segmentation:
@@ -71,7 +79,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 mased area with class color
+                    # filling the masked 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)
@@ -86,19 +94,13 @@ class ImageSegmenter:
                     cv2.rectangle(image, (x1, y1), (x1+(len(disp_str)*9), y1+15), cls_clr, -1)
                     cv2.putText(image, disp_str, (x1+5, y1+10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                 
-                
                 # draw segmentation boundary
                 if len(seg_mask_boundary) and self.is_show_segmentation_boundary and cls_conf[id] > self.confidence_threshold:            
                     cv2.polylines(image, [np.array(seg_mask_boundary[id], dtype=np.int32)], isClosed=True, color=cls_clr, thickness=2)
 
-
                 # object variables
                 (x1, y1, x2, y2) = bounding_boxes[id]
                 center = x1+(x2-x1)//2, y1+(y2-y1)//2
                 objects_data.append([cls, self.model.names[cls], center, self.masks[id], cls_clr])
 
-        return image, objects_data
-    
-
-
-        
+        return image, objects_data