Making it installabel package and switching to segmentation mode

src/model/yolo_wrapper.py

@@ -16,7 +16,7 @@ logger = get_logger(__name__)
 class YOLOWrapper:
     """Wrapper for YOLOv8 model operations."""
 
-    def __init__(self, model_path: str = "yolov8s.pt"):
+    def __init__(self, model_path: str = "yolov8s-seg.pt"):
         """
         Initialize YOLO model.
 
@@ -282,6 +282,10 @@ class YOLOWrapper:
                 boxes = result.boxes
                 image_path = str(result.path)
                 orig_shape = result.orig_shape  # (height, width)
+                height, width = orig_shape
+
+                # Check if this is a segmentation model with masks
+                has_masks = hasattr(result, "masks") and result.masks is not None
 
                 for i in range(len(boxes)):
                     # Get normalized coordinates
@@ -299,6 +303,33 @@ class YOLOWrapper:
                             float(v) for v in boxes.xyxy[i].cpu().numpy()
                         ],  # Absolute pixels
                     }
+
+                    # Extract segmentation mask if available
+                    if has_masks:
+                        try:
+                            # Get the mask for this detection
+                            mask_data = result.masks.xy[
+                                i
+                            ]  # Polygon coordinates in absolute pixels
+
+                            # Convert to normalized coordinates
+                            if len(mask_data) > 0:
+                                mask_normalized = []
+                                for point in mask_data:
+                                    x_norm = float(point[0]) / width
+                                    y_norm = float(point[1]) / height
+                                    mask_normalized.append([x_norm, y_norm])
+                                detection["segmentation_mask"] = mask_normalized
+                            else:
+                                detection["segmentation_mask"] = None
+                        except Exception as mask_error:
+                            logger.warning(
+                                f"Error extracting mask for detection {i}: {mask_error}"
+                            )
+                            detection["segmentation_mask"] = None
+                    else:
+                        detection["segmentation_mask"] = None
+
                     detections.append(detection)
 
             return detections