diff --git a/docs/api.md b/docs/api.md index 89a45d5d..56be7307 100644 --- a/docs/api.md +++ b/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: diff --git a/samgeo/api.py b/samgeo/api.py index f4fa26a2..c7490126 100644 --- a/samgeo/api.py +++ b/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,6 +712,7 @@ 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( @@ -663,7 +720,12 @@ async def segment_automatic( 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,6 +859,7 @@ 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( @@ -802,7 +867,12 @@ async def segment_predict( 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 diff --git a/samgeo/samgeo.py b/samgeo/samgeo.py index d0879e54..a5be720d 100644 --- a/samgeo/samgeo.py +++ b/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,6 +336,7 @@ 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: @@ -341,6 +344,8 @@ def save_masks( 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): diff --git a/samgeo/samgeo2.py b/samgeo/samgeo2.py index 9fcb9ef6..0ba010cc 100644 --- a/samgeo/samgeo2.py +++ b/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,6 +338,7 @@ 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: @@ -343,6 +346,8 @@ def save_masks( 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: diff --git a/samgeo/samgeo3.py b/samgeo/samgeo3.py index a10ad041..8a60d04d 100644 --- a/samgeo/samgeo3.py +++ b/samgeo/samgeo3.py @@ -246,6 +246,7 @@ def __init__( self.masks = None self.boxes = None self.scores = None + self.mask_scores: Optional[Dict[int, float]] = None self.logits = None self.objects = None self.prediction = None @@ -2530,6 +2531,26 @@ def show_points( plt.axis(axis) plt.show() + def _score_at(self, index: int) -> Optional[float]: + """Return the confidence score of the mask at ``index`` as a float. + + Args: + index: Position of the mask in ``self.masks``. + + Returns: + The score as a Python float, or None when no score is available. + """ + if self.scores is None: + return None + try: + value = self.scores[index] + except (IndexError, TypeError, KeyError): + return None + try: + return float(value.item()) if hasattr(value, "item") else float(value) + except (TypeError, ValueError): + return None + def save_masks( self, output: Optional[str] = None, @@ -2578,6 +2599,11 @@ def save_masks( (self.image_height, self.image_width), dtype=np.float32 ) + # Map each unique raster value to its confidence score so downstream + # consumers (e.g. the REST API's GeoJSON output) can attach scores to + # vectorized masks without a second inference run. + self.mask_scores: Optional[Dict[int, float]] = {} if unique else None + # Process each mask valid_mask_count = 0 mask_index = 0 @@ -2609,23 +2635,21 @@ def save_masks( continue # Get confidence score for this mask - if save_scores is not None: - if hasattr(self.scores[mask_index], "item"): - score = self.scores[mask_index].item() - else: - score = float(self.scores[mask_index]) + score = self._score_at(mask_index) # Add mask to array if unique: # Assign unique value to each mask (starting from 1) mask_value = valid_mask_count + 1 mask_array[mask_bool] = mask_value + if score is not None and self.mask_scores is not None: + self.mask_scores[mask_value] = score else: # Binary mask: all foreground pixels are 255 mask_array[mask_bool] = 255 # Add score to scores array - if save_scores is not None: + if save_scores is not None and score is not None: scores_array[mask_bool] = score valid_mask_count += 1 @@ -2635,12 +2659,30 @@ def save_masks( print("No masks met the size criteria.") return - # Convert to requested dtype - if dtype == "uint8": - if unique and valid_mask_count > 255: + # Convert to requested dtype. Unique mask ids must survive the cast + # unchanged, or the raster values no longer match ``mask_scores`` + # (and masks beyond the range wrap onto earlier ids), so promote an + # unsigned dtype that is too narrow for the number of masks. + if unique: + requested = np.dtype(dtype) + if requested.kind not in "iu": + raise ValueError( + f"Unique mask ids need an integer dtype, got {dtype!r}." + ) + if valid_mask_count > np.iinfo(requested).max: + for candidate in ("uint16", "uint32"): + if valid_mask_count <= np.iinfo(candidate).max: + break + else: + raise ValueError( + f"{valid_mask_count} masks exceed what uint32 can represent." + ) print( - f"Warning: {valid_mask_count} masks found, but uint8 can only represent 255 unique values. Consider using dtype='uint16'." + f"Warning: {valid_mask_count} masks found, more than {dtype} " + f"can represent; saving as {candidate} instead." ) + dtype = candidate + if dtype == "uint8": mask_array = mask_array.astype(np.uint8) elif dtype == "uint16": mask_array = mask_array.astype(np.uint16) diff --git a/tests/test_api.py b/tests/test_api.py index 3a637699..77f1897b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -330,3 +330,65 @@ def test_set_image_cached_skips_for_sam2_but_not_sam3(): assert _set_image_cached(sam3_model, sam3_key, "/tmp/y.tif", "hash-b") is True assert _set_image_cached(sam3_model, sam3_key, "/tmp/y.tif", "hash-b") is True assert sam3_model.set_image.call_count == 2 + + +def test_attach_mask_scores_joins_by_value(): + """Each feature gains the score recorded for its mask value.""" + from samgeo.api import _attach_mask_scores + + data = { + "type": "FeatureCollection", + "features": [ + {"type": "Feature", "geometry": None, "properties": {"value": 1}}, + {"type": "Feature", "geometry": None, "properties": {"value": 2.0}}, + {"type": "Feature", "geometry": None, "properties": {"value": 3}}, + {"type": "Feature", "geometry": None, "properties": None}, + ], + } + out = _attach_mask_scores(data, {1: 0.9, 2: 0.5}) + props = [f["properties"] for f in out["features"]] + assert props[0] == {"value": 1, "score": 0.9} + assert props[1] == {"value": 2.0, "score": 0.5} + assert props[2] == {"value": 3} + assert props[3] == {} + # No scores: the collection is returned untouched. + assert _attach_mask_scores(data, None) is data + assert _attach_mask_scores(data, {}) is data + + +@patch("samgeo.api.get_model") +def test_text_geojson_includes_scores(mock_get, sample_image_path): + """/segment/text geojson carries a ``score`` per mask in one inference run. + + The model records ``mask_scores`` (raster value -> confidence) when it + saves the masks, and the API joins that onto the vectorized output so + clients do not need a second ``detections`` request. + """ + from threading import Lock + + mock_model = MagicMock() + mock_model.masks = [np.ones((64, 64), dtype=np.uint8)] * 2 + + def fake_save_masks(output, **kwargs): + from PIL import Image + + arr = np.zeros((64, 64), dtype=np.uint8) + arr[:32, :] = 1 + arr[32:, :] = 2 + Image.fromarray(arr).save(output) + mock_model.mask_scores = {1: 0.91, 2: 0.47} + + mock_model.save_masks = fake_save_masks + mock_get.return_value = (mock_model, Lock(), ("sam3", "facebook/sam3", True)) + + with open(sample_image_path, "rb") as f: + response = client.post( + "/segment/text", + files={"file": ("test.tif", f, "image/tiff")}, + data={"prompt": "building", "output_format": "geojson"}, + ) + assert response.status_code == 200 + features = response.json()["features"] + scores = {f["properties"]["value"]: f["properties"].get("score") for f in features} + assert scores == {1: 0.91, 2: 0.47} + mock_model.save_masks = None # ensure no second inference path was needed diff --git a/tests/test_samgeo3.py b/tests/test_samgeo3.py index b7368686..f42e552c 100644 --- a/tests/test_samgeo3.py +++ b/tests/test_samgeo3.py @@ -262,3 +262,73 @@ def test_sam31_missing_checkpoint_helper_has_clear_error(monkeypatch, tmp_path, ) builder.assert_not_called() + + +def test_save_masks_records_mask_scores(monkeypatch, samgeo3, tmp_path): + """``save_masks`` maps each unique raster value to its confidence score, + skipping masks filtered by size so the values stay aligned.""" + import numpy as np + + model = samgeo3.SamGeo3.__new__(samgeo3.SamGeo3) + model.image_height = 4 + model.image_width = 4 + model.source = None + model.scores = [0.9, 0.2, 0.6] + small = np.zeros((4, 4), dtype=bool) + small[0, 0] = True + big_a = np.zeros((4, 4), dtype=bool) + big_a[:2, :] = True + big_b = np.zeros((4, 4), dtype=bool) + big_b[2:, :] = True + model.masks = [big_a, small, big_b] + monkeypatch.setattr( + samgeo3.common, "array_to_image", lambda *a, **k: None, raising=False + ) + + model.save_masks(output=None, min_size=2) + + assert model.mask_scores == {1: 0.9, 2: 0.6} + assert model.objects.max() == 2 + + model.save_masks(output=None, unique=False) + assert model.mask_scores is None + + +def test_save_masks_promotes_dtype_so_scores_stay_aligned(monkeypatch, samgeo3): + """More masks than uint8 can hold promote the dtype instead of wrapping + ids, so every raster value still matches its ``mask_scores`` key.""" + import numpy as np + + n = 300 + model = samgeo3.SamGeo3.__new__(samgeo3.SamGeo3) + model.image_height = 2 + model.image_width = n + model.source = None + model.scores = [i / n for i in range(n)] + masks = [] + for i in range(n): + m = np.zeros((2, n), dtype=bool) + m[:, i] = True + masks.append(m) + model.masks = masks + monkeypatch.setattr( + samgeo3.common, "array_to_image", lambda *a, **k: None, raising=False + ) + + model.save_masks(output=None, dtype="uint8") + + assert model.objects.dtype == np.uint16 + np.testing.assert_array_equal( + model.objects[0], np.arange(1, n + 1, dtype=np.uint16) + ) + assert model.mask_scores == {i + 1: i / n for i in range(n)} + + # Signed dtypes are promoted the same way once ids would wrap. + model.save_masks(output=None, dtype="int8") + assert model.objects.dtype == np.uint16 + np.testing.assert_array_equal( + model.objects[0], np.arange(1, n + 1, dtype=np.uint16) + ) + + with pytest.raises(ValueError, match="integer dtype"): + model.save_masks(output=None, dtype="float32")