feat: include per-mask confidence score in API GeoJSON output

docs/api.md

@@ -286,6 +286,17 @@ curl -X POST http://localhost:8000/segment/text \
 }
 ```
 
+## GeoJSON output
+
+With `output_format=geojson`, every segmentation endpoint returns one polygon
+feature per mask. Each feature's `properties` carry the mask's raster `value`
+and, when the model reports one, its confidence `score`, so a single request
+yields both geometry and confidence:
+
+```json
+{"type": "Feature", "geometry": {"type": "Polygon", "coordinates": [...]}, "properties": {"value": 1, "score": 0.887}}
+```
+
 ## Caching
 
 The API automatically caches models and image encodings for better performance:

samgeo/api.py

@@ -285,13 +285,66 @@ def _validate_output_format(output_format: str) -> None:
         )
 
 
-def _format_response(raster_path: str, output_format: str, tmpdir: str):
+def _attach_mask_scores(data: dict, mask_scores: Optional[dict]) -> dict:
+    """Add a ``score`` property to each GeoJSON feature from its mask value.
+
+    ``raster_to_geojson`` writes one feature per mask with the raster value in
+    the ``value`` property. The model records the confidence score of every
+    mask under that same value in ``mask_scores``, so the two are joined here
+    without a second inference run.
+
+    Args:
+        data: A GeoJSON FeatureCollection dict produced from the mask raster.
+        mask_scores: Mapping of raster value to confidence score, or None.
+
+    Returns:
+        The same FeatureCollection with ``score`` set where a score is known.
+    """
+    if not mask_scores or not isinstance(data, dict):
+        return data
+    for feature in data.get("features", []):
+        props = feature.get("properties") or {}
+        feature["properties"] = props
+        try:
+            key = int(props.get("value"))
+        except (TypeError, ValueError):
+            continue
+        if key in mask_scores:
+            props["score"] = mask_scores[key]
+    return data
+
+
+def _snapshot_mask_scores(model) -> Optional[dict]:
+    """Copy the model's per-mask scores while its lock is still held.
+
+    Cached models are shared between requests, so the mapping is copied before
+    the lock is released and another request can overwrite it.
+
+    Args:
+        model: The model instance that just saved its masks.
+
+    Returns:
+        A copy of ``model.mask_scores``, or None when the model has none.
+    """
+    scores = getattr(model, "mask_scores", None)
+    return dict(scores) if scores else None
+
+
+def _format_response(
+    raster_path: str,
+    output_format: str,
+    tmpdir: str,
+    mask_scores: Optional[dict] = None,
+):
     """Convert a raster mask to the requested output format.
 
     Args:
         raster_path: Path to the raster mask file.
         output_format: One of "geojson", "geotiff", "png".
         tmpdir: Temporary directory for intermediate files.
+        mask_scores: Mapping of raster value to confidence score recorded by
+            the model when it saved the masks. When given, each GeoJSON
+            feature gains a ``score`` property.
 
     Returns:
         FastAPI response object.
@@ -311,6 +364,7 @@ def _format_response(raster_path: str, output_format: str, tmpdir: str):
         raster_to_geojson(raster_path, geojson_path)
         with open(geojson_path) as f:
             data = json.load(f)
+        data = _attach_mask_scores(data, mask_scores)
         _cleanup_tmpdir(tmpdir)
         return JSONResponse(content=data)
 
@@ -617,6 +671,7 @@ async def segment_automatic(
         input_path, image_hash = await _save_upload(file, tmpdir)
         output_path = os.path.join(tmpdir, "mask.tif")
 
+        mask_scores = None
         t_start = time.time()
         if model_version == "sam3":
             model, lock, model_key = get_model(model_version, model_id)
@@ -634,6 +689,7 @@ async def segment_automatic(
                         detail="No objects found for automatic segmentation.",
                     )
                 model.save_masks(output=output_path, unique=unique)
+                mask_scores = _snapshot_mask_scores(model)
         else:
             sam_kwargs = {
                 "points_per_side": points_per_side,
@@ -656,14 +712,20 @@ async def segment_automatic(
                     min_size=min_size,
                     max_size=max_size,
                 )
+                mask_scores = _snapshot_mask_scores(model)
 
         t_inference = time.time() - t_start
         logger.info(
             "Automatic segmentation completed in %.2fs (model: %s)",
             t_inference,
             model_version,
         )
-        return _format_response(output_path, output_format, tmpdir)
+        return _format_response(
+            output_path,
+            output_format,
+            tmpdir,
+            mask_scores=mask_scores,
+        )
     except HTTPException:
         _cleanup_tmpdir(tmpdir)
         raise
@@ -745,6 +807,7 @@ async def segment_predict(
         if boxes is not None:
             parsed_boxes = np.array(json.loads(boxes))
 
+        mask_scores = None
         t_start = time.time()
 
         if model_version == "sam3":
@@ -781,6 +844,7 @@ async def segment_predict(
                     min_size=min_size,
                     max_size=max_size,
                 )
+                mask_scores = _snapshot_mask_scores(model)
         else:
             model, lock, model_key = get_model(
                 model_version, model_id, automatic=False
@@ -795,14 +859,20 @@ async def segment_predict(
                     multimask_output=multimask_output,
                     output=output_path,
                 )
+                mask_scores = _snapshot_mask_scores(model)
 
         t_inference = time.time() - t_start
         logger.info(
             "Prompt segmentation completed in %.2fs (model: %s)",
             t_inference,
             model_version,
         )
-        return _format_response(output_path, output_format, tmpdir)
+        return _format_response(
+            output_path,
+            output_format,
+            tmpdir,
+            mask_scores=mask_scores,
+        )
     except HTTPException:
         _cleanup_tmpdir(tmpdir)
         raise
@@ -859,6 +929,7 @@ async def segment_text(
             backend=backend,
             confidence_threshold=confidence_threshold,
         )
+        mask_scores = None
         t_start = time.time()
         with lock:
             _set_image_cached(model, model_key, input_path, image_hash)
@@ -883,6 +954,7 @@ async def segment_text(
                     det_result = _build_detections_json(model)
             else:
                 model.save_masks(output=output_path)
+                mask_scores = _snapshot_mask_scores(model)
 
         t_inference = time.time() - t_start
         logger.info(
@@ -893,7 +965,12 @@ async def segment_text(
         if output_format in ("detections", "json"):
             _cleanup_tmpdir(tmpdir)
             return JSONResponse(content=det_result)
-        return _format_response(output_path, output_format, tmpdir)
+        return _format_response(
+            output_path,
+            output_format,
+            tmpdir,
+            mask_scores=mask_scores,
+        )
     except HTTPException:
         _cleanup_tmpdir(tmpdir)
         raise

samgeo/samgeo.py

@@ -119,6 +119,7 @@ def __init__(
         self.prediction = None
         self.scores = None
         self.logits = None
+        self.mask_scores = None
 
         # Build the SAM model
         self.sam = sam_model_registry[self.model_type](checkpoint=self.checkpoint)
@@ -305,6 +306,7 @@ def save_masks(
             **kwargs: Other arguments for array_to_image().
 
         """
+        self.mask_scores = None
 
         if self.masks is None:
             raise ValueError("No masks found. Please run generate() first.")
@@ -334,13 +336,16 @@ def save_masks(
             )
             # Assign a unique value to each object
             count = len(sorted_masks)
+            self.mask_scores = {}
             for index, ann in enumerate(sorted_masks):
                 m = ann["segmentation"]
                 if min_size > 0 and ann["area"] < min_size:
                     continue
                 if max_size is not None and ann["area"] > max_size:
                     continue
                 objects[m] = count - index
+                if "predicted_iou" in ann:
+                    self.mask_scores[count - index] = float(ann["predicted_iou"])
 
         # Generate a binary mask
         else:
@@ -591,6 +596,7 @@ def predict(
             return_results (bool, optional): Whether to return the predicted masks, scores, and logits. Defaults to False.
 
         """
+        self.mask_scores = None
         out_of_bounds = []
 
         if isinstance(boxes, str):

samgeo/samgeo2.py

@@ -143,6 +143,7 @@ def __init__(
         self.model_id = model_id
         self.model_version = "sam2"
         self.device = device
+        self.mask_scores: Optional[Dict[int, float]] = None
 
         if video:
             automatic = False
@@ -307,6 +308,7 @@ def save_masks(
             max_size (int, optional): The maximum size of the object. Defaults to None.
             **kwargs: Additional keyword arguments for common.array_to_image().
         """
+        self.mask_scores = None
 
         if self.masks is None:
             raise ValueError("No masks found. Please run generate() first.")
@@ -336,13 +338,16 @@ def save_masks(
             )
             # Assign a unique value to each object
             count = len(sorted_masks)
+            self.mask_scores = {}
             for index, ann in enumerate(sorted_masks):
                 m = ann["segmentation"]
                 if min_size > 0 and ann["area"] < min_size:
                     continue
                 if max_size is not None and ann["area"] > max_size:
                     continue
                 objects[m] = count - index
+                if "predicted_iou" in ann:
+                    self.mask_scores[count - index] = float(ann["predicted_iou"])
 
         # Generate a binary mask
         else:
@@ -615,6 +620,7 @@ def predict(
             Tuple[np.ndarray, np.ndarray, np.ndarray]: The mask, the multimask,
                 and the logits.
         """
+        self.mask_scores = None
         import geopandas as gpd
 
         out_of_bounds = []
@@ -1157,6 +1163,10 @@ def save_prediction(
 
         array = self.masks[index] * mask_multiplier
         self.prediction = array
+        try:
+            self.mask_scores = {int(mask_multiplier): float(self.scores[index])}
+        except (TypeError, ValueError, IndexError):
+            self.mask_scores = None
         common.array_to_image(array, output, self.source, dtype=dtype, **kwargs)
 
         if vector is not None: