Update image_segmenter.py

image_segmenter.py

@@ -2,12 +2,11 @@ import cv2
 import numpy as np
 from ultralytics import YOLO
 import random
-import torch
 import spaces
 
 class ImageSegmenter:
     def __init__(self, model_type="yolov8s-seg") -> None:
-        # Store parameters but don't initialize CUDA
+        # Don't initialize any CUDA/GPU stuff here
         self.model_type = model_type
         self.is_show_bounding_boxes = True
         self.is_show_segmentation_boundary = False
@@ -17,13 +16,11 @@ class ImageSegmenter:
         self.bb_thickness = 2
         self.bb_clr = (255, 0, 0)
         self.masks = {}
-        self.model = None  # Model will be loaded in predict
+        self.model = None
 
     def get_cls_clr(self, cls_id):
         if cls_id in self.cls_clr:
             return self.cls_clr[cls_id]
-        
-        # gen rand color
         r = random.randint(50, 200)
         g = random.randint(50, 200)
         b = random.randint(50, 200)
@@ -32,11 +29,10 @@ class ImageSegmenter:
 
     @spaces.GPU
     def predict(self, image):            
-        # Load model if not loaded
+        # Initialize model if needed
         if self.model is None:
             print("Loading YOLO model...")
             self.model = YOLO('models/' + self.model_type + '.pt')
-            self.model.to('cuda')
             print("Model loaded successfully")
 
         # params
@@ -48,7 +44,6 @@ class ImageSegmenter:
         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
             seg_mask = predictions[0].masks.data.cpu().numpy()  
@@ -58,27 +53,21 @@ class ImageSegmenter:
         for id, cls in enumerate(cls_ids):
             cls_clr = self.get_cls_clr(cls)
 
-            # draw filled segmentation region
             if seg_mask.any() and cls_conf[id] > self.confidence_threshold:
                 self.masks[id] = seg_mask[id]
                 
                 if self.is_show_segmentation:
                     alpha = 0.8                
-
-                    # converting the mask from 1 channel to 3 channels
                     colored_mask = np.expand_dims(seg_mask[id], 0).repeat(3, axis=0)
                     colored_mask = np.moveaxis(colored_mask, 0, -1)
 
-                    # Resize the mask to match the image size, if necessary
                     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
                     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)
 
-                # draw bounding box with class name and score
                 if self.is_show_bounding_boxes and cls_conf[id] > self.confidence_threshold:
                     (x1, y1, x2, y2) = bounding_boxes[id]
                     cls_name = self.model.names[cls]
@@ -88,11 +77,9 @@ 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])