Update image_segmenter.py

image_segmenter.py

@@ -2,124 +2,103 @@ import cv2
 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", device="cpu"):
-        self.device = device
-        self.model = YOLO(model_type).to(self.device)
+    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.is_show_bounding_boxes = True
         self.is_show_segmentation_boundary = False
         self.is_show_segmentation = False
         self.confidence_threshold = 0.5
         self.cls_clr = {}
+
+        # params
         self.bb_thickness = 2
         self.bb_clr = (255, 0, 0)
+
+        # variables
         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:
             return self.cls_clr[cls_id]
+        
+        # gen rand color
         r = random.randint(50, 200)
         g = random.randint(50, 200)
         b = random.randint(50, 200)
         self.cls_clr[cls_id] = (r, g, b)
         return (r, g, b)
 
-    @spaces.GPU
     def predict(self, image):            
-        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")
-
-            # Ensure image is in correct format
-            if isinstance(image, np.ndarray):
-                image = image.copy()
-            else:
-                raise ValueError("Input image must be a numpy array")
-
-            # Make prediction using CPU
-            predictions = self.model.predict(image, device='cpu')
-
-            # 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]))
-
-                        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:
+        # params
+        objects_data = [] 
+        image = image.copy()
+        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
+            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)
+
+            # 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[int(cls)]
+                    cls_name = self.model.names[cls]
                     cls_confidence = cls_conf[id]
-                    disp_str = f"{cls_name} {cls_confidence:.2f}"
+                    disp_str = cls_name +' '+ str(round(cls_confidence, 2))
                     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) > 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)
+                
+                # 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([int(cls), self.model.names[int(cls)], center, 
-                                  self.masks.get(id, None), cls_clr])
-
-            return image, objects_data
-            
-        except Exception as e:
-            print(f"Error in predict: {str(e)}")
-            import traceback
-            print(traceback.format_exc())
-            raise
+                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
+    
+
+
+