add vectorize

app.py

@@ -124,16 +124,58 @@ def api_predict(image, text_prompt, confidence_threshold):
     except Exception as e:
         return {"error": str(e)}
 
+def _mask_to_polygons_original_size(binary_mask, epsilon=1.0):
+    """
+    Convert binary mask to vector polygons (mask is already at original image size)
+
+    Args:
+        binary_mask: Binary mask array (0 or 1) at original image size
+        epsilon: Polygon simplification epsilon
+
+    Returns:
+        List of polygons, where each polygon is a list of [x, y] points in pixel coordinates
+    """
+    import cv2
+    import numpy as np
+
+    try:
+        # Find contours using OpenCV
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        polygons = []
+
+        for contour in contours:
+            if len(contour) < 3:  # Skip small contours
+                continue
+
+            # Simplify polygon using Douglas-Peucker algorithm
+            simplified = cv2.approxPolyDP(contour, epsilon, True)
+
+            # Convert to list of [x, y] points
+            polygon_points = [[float(point[0][0]), float(point[0][1])] for point in simplified]
+
+            # Only add polygons with at least 3 points
+            if len(polygon_points) >= 3:
+                polygons.append(polygon_points)
+
+        return polygons
+
+    except Exception as e:
+        # Return empty list on error, but don't fail the entire request
+        print(f"Warning: Polygon extraction failed: {e}")
+        return []
+
 @spaces.GPU
 def sam2_compatible_api(data):
     """
     SAM2-compatible API endpoint with SAM3 extensions
     Supports text prompts (SAM3), points, and boxes (SAM2 compatible)
+    Includes vectorize option for polygon extraction
     """
     import numpy as np
     from PIL import Image
     import base64
     import io
+    import cv2
 
     try:
         inputs_data = data.get("inputs", {})
@@ -145,6 +187,8 @@ def sam2_compatible_api(data):
         input_labels = inputs_data.get("labels", [])
         input_boxes = inputs_data.get("boxes", [])
         confidence_threshold = inputs_data.get("confidence_threshold", 0.5)
+        vectorize = inputs_data.get("vectorize", False)
+        simplify_epsilon = inputs_data.get("simplify_epsilon", 1.0)
 
         # Validate inputs
         if not image_b64:
@@ -165,9 +209,11 @@ def sam2_compatible_api(data):
             image_b64 = image_b64.split(',')[1]
         image_bytes = base64.b64decode(image_b64)
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        original_image_size = image.size  # Store for response metadata
 
         all_masks = []
         all_scores = []
+        all_polygons = []
 
         # Process text prompts (SAM3 feature)
         if has_text:
@@ -188,6 +234,12 @@ def sam2_compatible_api(data):
                             all_masks.append(mask_b64)
                             all_scores.append(score)
 
+                            # Extract polygons if vectorize is enabled
+                            if vectorize:
+                                binary_mask = (mask_np > 0).astype(np.uint8)
+                                polygons = _mask_to_polygons_original_size(binary_mask, simplify_epsilon)
+                                all_polygons.append(polygons)
+
         # Process visual prompts (SAM2 compatibility) - Basic implementation
         if has_boxes or has_points:
             # For visual prompts, use a generic prompt to get masks
@@ -211,8 +263,14 @@ def sam2_compatible_api(data):
                         all_masks.append(mask_b64)
                         all_scores.append(score)
 
+                        # Extract polygons if vectorize is enabled
+                        if vectorize:
+                            binary_mask = (mask_np > 0).astype(np.uint8)
+                            polygons = _mask_to_polygons_original_size(binary_mask, simplify_epsilon)
+                            all_polygons.append(polygons)
+
         # Build SAM2-compatible response
-        return {
+        response = {
             "masks": all_masks,
             "scores": all_scores,
             "num_objects": len(all_masks),
@@ -220,6 +278,16 @@ def sam2_compatible_api(data):
             "success": True
         }
 
+        # Add polygon data if vectorize is enabled
+        if vectorize:
+            response.update({
+                "polygons": all_polygons,
+                "polygon_format": "pixel_coordinates",
+                "original_image_size": original_image_size
+            })
+
+        return response
+
     except Exception as e:
         return {"error": str(e), "success": False, "sam_version": "3.0"}
 
@@ -340,6 +408,22 @@ response = requests.post(
     }
 )
 
+# SAM3 with Vectorize (returns both masks and polygons)
+response = requests.post(
+    "https://your-username-sam3-api.hf.space/api/sam2_compatible",
+    json={
+        "data": {
+            "inputs": {
+                "image": image_b64,
+                "text_prompts": ["cat"],
+                "confidence_threshold": 0.5,
+                "vectorize": true,
+                "simplify_epsilon": 1.0
+            }
+        }
+    }
+)
+
 result = response.json()
         </pre>
 
@@ -361,7 +445,9 @@ result = response.json()
     "boxes": [[x1, y1, x2, y2], [x1, y1, x2, y2]],  // Bounding boxes
 
     "multimask_output": false,  // Optional, defaults to False
-    "confidence_threshold": 0.5  // Optional, minimum confidence for returned masks
+    "confidence_threshold": 0.5,  // Optional, minimum confidence for returned masks
+    "vectorize": false,  // Optional, return vector polygons instead of/in addition to bitmaps
+    "simplify_epsilon": 1.0  // Optional, polygon simplification factor
   }
 }
         </pre>
@@ -373,7 +459,12 @@ result = response.json()
   "scores": [0.95, 0.87],
   "num_objects": 2,
   "sam_version": "3.0",
-  "success": true
+  "success": true,
+
+  // If vectorize=true, additional fields:
+  "polygons": [[[x1,y1],[x2,y2],...], [[x1,y1],...]],  // Array of polygon arrays for each object
+  "polygon_format": "pixel_coordinates",
+  "original_image_size": [width, height]
 }
         </pre>
         """)