Update image_segmenter.py

image_segmenter.py

@@ -3,10 +3,12 @@ import numpy as np
 from ultralytics import YOLO
 import random
 import spaces
+import os
+import torch
 
 class ImageSegmenter:
     def __init__(self, model_type="yolov8s-seg") -> None:
-        # Don't initialize any CUDA/GPU stuff here
+        # Initialize parameters
         self.model_type = model_type
         self.is_show_bounding_boxes = True
         self.is_show_segmentation_boundary = False
@@ -17,6 +19,17 @@ class ImageSegmenter:
         self.bb_clr = (255, 0, 0)
         self.masks = {}
         self.model = None
+        
+        # Ensure model directory exists
+        os.makedirs('models', exist_ok=True)
+        
+        # Check if model file exists, if not download it
+        model_path = os.path.join('models', f'{model_type}.pt')
+        if not os.path.exists(model_path):
+            print(f"Downloading {model_type} model...")
+            self.model = YOLO(model_type)
+            self.model.export()
+            print("Model downloaded successfully")
 
     def get_cls_clr(self, cls_id):
         if cls_id in self.cls_clr:
@@ -29,59 +42,84 @@ class ImageSegmenter:
 
     @spaces.GPU
     def predict(self, image):            
-        # Initialize model if needed
-        if self.model is None:
-            print("Loading YOLO model...")
-            self.model = YOLO('models/' + self.model_type + '.pt')
-            print("Model loaded successfully")
+        try:
+            # Initialize model if needed
+            if self.model is None:
+                print("Loading YOLO model...")
+                model_path = os.path.join('models', f'{self.model_type}.pt')
+                # Force CPU mode for YOLO initialization
+                self.model = YOLO(model_path)
+                self.model.to('cpu')  # Explicitly move to CPU
+                print("Model loaded successfully")
 
-        # params
-        objects_data = [] 
-        image = image.copy()
-        predictions = self.model.predict(image)
+            # Ensure image is in correct format
+            if isinstance(image, np.ndarray):
+                image = image.copy()
+            else:
+                raise ValueError("Input image must be a numpy array")
 
-        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()
-        
-        if predictions[0].masks:
-            seg_mask_boundary = predictions[0].masks.xy
-            seg_mask = predictions[0].masks.data.cpu().numpy()  
-        else:
-            seg_mask_boundary, seg_mask = [], np.array([])    
-        
-        for id, cls in enumerate(cls_ids):
-            cls_clr = self.get_cls_clr(cls)
+            # Make prediction using CPU
+            predictions = self.model.predict(image, device='cpu')
 
-            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                
-                    colored_mask = np.expand_dims(seg_mask[id], 0).repeat(3, axis=0)
-                    colored_mask = np.moveaxis(colored_mask, 0, -1)
+            # Process results
+            objects_data = []
+            
+            if len(predictions) == 0 or not predictions[0].boxes:
+                return image, objects_data
+
+            cls_ids = predictions[0].boxes.cls.numpy()  # Changed from cpu().numpy()
+            bounding_boxes = predictions[0].boxes.xyxy.int().numpy()        
+            cls_conf = predictions[0].boxes.conf.numpy()
+            
+            if predictions[0].masks is not None:
+                seg_mask_boundary = predictions[0].masks.xy
+                seg_mask = predictions[0].masks.data.numpy()  # Changed from cpu().numpy()
+            else:
+                seg_mask_boundary, seg_mask = [], np.array([])
+
+            for id, cls in enumerate(cls_ids):
+                if cls_conf[id] <= self.confidence_threshold:
+                    continue
+                    
+                cls_clr = self.get_cls_clr(int(cls))
+
+                if seg_mask.size > 0:
+                    self.masks[id] = seg_mask[id]
+                    
+                    if self.is_show_segmentation:
+                        alpha = 0.8                
+                        colored_mask = np.expand_dims(seg_mask[id], 0).repeat(3, axis=0)
+                        colored_mask = np.moveaxis(colored_mask, 0, -1)
 
-                    if image.shape[:2] != seg_mask[id].shape[:2]:
-                        colored_mask = cv2.resize(colored_mask, (image.shape[1], image.shape[0]))
+                        if image.shape[:2] != seg_mask[id].shape[:2]:
+                            colored_mask = cv2.resize(colored_mask, (image.shape[1], image.shape[0]))
 
-                    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)
+                        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)
 
-                if self.is_show_bounding_boxes and cls_conf[id] > self.confidence_threshold:
+                if self.is_show_bounding_boxes:
                     (x1, y1, x2, y2) = bounding_boxes[id]
-                    cls_name = self.model.names[cls]
+                    cls_name = self.model.names[int(cls)]
                     cls_confidence = cls_conf[id]
-                    disp_str = cls_name +' '+ str(round(cls_confidence, 2))
+                    disp_str = f"{cls_name} {cls_confidence:.2f}"
                     cv2.rectangle(image, (x1, y1), (x2, y2), cls_clr, self.bb_thickness)
-                    cv2.rectangle(image, (x1, y1), (x1+(len(disp_str)*9), y1+15), cls_clr, -1)
+                    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)
                 
-                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)
+                if len(seg_mask_boundary) > 0 and self.is_show_segmentation_boundary:
+                    cv2.polylines(image, [np.array(seg_mask_boundary[id], dtype=np.int32)], 
+                                isClosed=True, color=cls_clr, thickness=2)
 
                 (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])
+                center = (x1+(x2-x1)//2, y1+(y2-y1)//2)
+                objects_data.append([int(cls), self.model.names[int(cls)], center, 
+                                  self.masks.get(id, None), cls_clr])
 
-        return image, objects_data
+            return image, objects_data
+            
+        except Exception as e:
+            print(f"Error in predict: {str(e)}")
+            import traceback
+            print(traceback.format_exc())
+            raise