fix(api): make REST segmentation API robust for SAM3 (#535)

* fix(api): make REST segmentation API robust for SAM3

Validated the REST API (samgeo/api.py) end-to-end against real example
imagery and fixed several bugs found in the process:

- Model cache keyed only on (model_version, model_id), so an automatic
  mask-generation instance was reused for a prompt-predict request and
  crashed with "'SamGeo2' object has no attribute 'predictor'". The cache
  key now includes the automatic flag; get_model returns the key so the
  image-hash cache stays consistent.
- SAM3 mutates its encoded-image state during generate_*, so the image-hash
  optimization (skipping set_image on repeat requests) made the 2nd+ SAM3
  request fail with "expected scalar type BFloat16 but found Float".
  _set_image_cached now always re-encodes for SAM3 (still caches SAM/SAM2).
- raster_to_vector crashed on an empty mask (no geometry column). It now
  emits a valid empty FeatureCollection.
- /segment/automatic gained the empty-masks guard the other endpoints have
  (404 instead of a 500 "No masks found").
- /segment/automatic now defaults to SAM3.

Verified on docs/examples/tree_image.tif: text/box/point/detections all
return GeoJSON across consecutive requests. Added regression tests.

* fix(api): include model-shaping kwargs in the model cache key

Address CodeRabbit review on #535: the cache key only considered
(model_version, model_id, automatic), so two requests differing only by
backend/confidence_threshold (SAM3) or points_per_side/pred_iou_thresh/
stability_score_thresh (SAM/SAM2) would reuse the first, differently
configured instance and silently run with the wrong settings.

The key now folds in all constructor kwargs via a small order-independent
_freeze_kwargs() helper. Added regression tests covering distinct configs
(SAM3 confidence_threshold, SAM2 points_per_side) and the freezing helper.

Author: giswqs

samgeo/api.py

@@ -103,6 +103,30 @@ def _normalize_max_size(max_size: Optional[int]) -> Optional[int]:
 }
 
 
+def _freeze_kwargs(kwargs: dict):
+    """Return a hashable, order-independent representation of model kwargs.
+
+    Used to build the model cache key. Nested dicts (e.g. ``sam_kwargs``) become
+    sorted tuples and lists become tuples so two calls with the same
+    configuration map to the same key regardless of argument order.
+
+    Args:
+        kwargs: The keyword arguments passed to a model constructor.
+
+    Returns:
+        A hashable value derived from ``kwargs``.
+    """
+
+    def freeze(value):
+        if isinstance(value, dict):
+            return tuple(sorted((k, freeze(v)) for k, v in value.items()))
+        if isinstance(value, (list, tuple)):
+            return tuple(freeze(v) for v in value)
+        return value
+
+    return freeze(kwargs)
+
+
 def get_model(model_version: str, model_id: Optional[str] = None, **kwargs):
     """Get or create a cached model instance.
 
@@ -112,7 +136,12 @@ def get_model(model_version: str, model_id: Optional[str] = None, **kwargs):
         **kwargs: Additional keyword arguments for model initialization.
 
     Returns:
-        tuple: (model_instance, threading.Lock)
+        tuple: (model_instance, threading.Lock, cache_key). The cache key
+            includes the model-shaping kwargs (e.g. ``automatic``, ``backend``,
+            ``confidence_threshold``, ``points_per_side``) so that instances
+            built with different configurations -- including automatic
+            mask-generation vs. prompt prediction -- are cached separately and a
+            request never silently reuses a differently-configured instance.
 
     Raises:
         HTTPException: If model_version or model_id is invalid, or
@@ -140,11 +169,20 @@ def get_model(model_version: str, model_id: Optional[str] = None, **kwargs):
             ),
         )
 
-    key = (model_version, model_id)
+    # Include the model-shaping kwargs in the cache key. The same model class is
+    # built differently for automatic mask generation (a mask generator, via the
+    # default automatic=True) vs. prompt prediction (an image predictor, via
+    # automatic=False), and SAM3 takes per-request knobs like ``backend`` and
+    # ``confidence_threshold``. Keying on these keeps configurations from
+    # colliding -- which previously reused an automatic instance for a predict
+    # request ("no attribute 'predictor'") and would also silently reuse a
+    # stale threshold/points_per_side from an earlier request.
+    key = (model_version, model_id, _freeze_kwargs(kwargs))
     with _model_cache_lock:
         if key in _model_cache:
             logger.info("Model cache hit for %s", key)
-            return _model_cache[key]
+            model, lock = _model_cache[key]
+            return model, lock, key
 
         logger.info("Loading model %s", key)
         extra = _EXTRAS_MAP.get(model_version, model_version)
@@ -172,7 +210,8 @@ def get_model(model_version: str, model_id: Optional[str] = None, **kwargs):
                 ),
             )
         _model_cache[key] = (model, threading.Lock())
-        return _model_cache[key]
+        model, lock = _model_cache[key]
+        return model, lock, key
 
 
 def _set_image_cached(
@@ -183,16 +222,23 @@ def _set_image_cached(
     Compares the provided hash to the last image encoded on this model.
     Skips the expensive image-encoder forward pass when the hash matches.
 
+    SAM3 is exempt from the skip: its ``generate_*`` calls mutate the encoded
+    image state in place, so reusing it on a second request without re-encoding
+    fails with an ``expected scalar type BFloat16 but found Float`` dtype
+    mismatch. For SAM3 we always re-encode (correctness over the small speedup);
+    SAM/SAM2 still benefit from the cache.
+
     Args:
         model: The SAM model instance.
-        model_key: Cache key for the model, e.g. ("sam3", "facebook/sam3").
+        model_key: Cache key for the model, e.g. ("sam3", "facebook/sam3", True).
         image_path: Path to the saved image file.
         image_hash: SHA-256 hex digest of the uploaded file bytes.
 
     Returns:
         True if set_image was called (new image), False if skipped (cache hit).
     """
-    if _image_hash_cache.get(model_key) == image_hash:
+    is_sam3 = bool(model_key) and model_key[0] == "sam3"
+    if not is_sam3 and _image_hash_cache.get(model_key) == image_hash:
         logger.info("Image cache hit for model %s, skipping set_image()", model_key)
         # Update source path so save_masks() can find GeoTIFF metadata
         # even though we skip the expensive encoding step.
@@ -545,7 +591,7 @@ def clear_models():
 @app.post("/segment/automatic")
 async def segment_automatic(
     file: UploadFile = File(...),
-    model_version: str = Form("sam2"),
+    model_version: str = Form("sam3"),
     model_id: Optional[str] = Form(None),
     output_format: str = Form("geojson"),
     foreground: bool = Form(True),
@@ -583,15 +629,20 @@ async def segment_automatic(
 
         t_start = time.time()
         if model_version == "sam3":
-            model, lock = get_model(model_version, model_id)
-            model_key = (model_version, model_id or _DEFAULT_MODEL_IDS[model_version])
+            model, lock, model_key = get_model(model_version, model_id)
             with lock:
                 _set_image_cached(model, model_key, input_path, image_hash)
                 model.generate_masks(
                     prompt="everything",
                     min_size=min_size,
                     max_size=max_size,
                 )
+                if model.masks is None or len(model.masks) == 0:
+                    _cleanup_tmpdir(tmpdir)
+                    raise HTTPException(
+                        status_code=404,
+                        detail="No objects found for automatic segmentation.",
+                    )
                 model.save_masks(output=output_path, unique=unique)
         else:
             sam_kwargs = {
@@ -600,11 +651,11 @@ async def segment_automatic(
                 "stability_score_thresh": stability_score_thresh,
             }
             if model_version == "sam":
-                model, lock = get_model(
+                model, lock, _ = get_model(
                     model_version, model_id, sam_kwargs=sam_kwargs
                 )
             else:
-                model, lock = get_model(model_version, model_id, **sam_kwargs)
+                model, lock, _ = get_model(model_version, model_id, **sam_kwargs)
 
             with lock:
                 model.generate(
@@ -707,11 +758,7 @@ async def segment_predict(
         t_start = time.time()
 
         if model_version == "sam3":
-            model, lock = get_model(model_version, model_id)
-            model_key = (
-                model_version,
-                model_id or _DEFAULT_MODEL_IDS[model_version],
-            )
+            model, lock, model_key = get_model(model_version, model_id)
             with lock:
                 _set_image_cached(model, model_key, input_path, image_hash)
                 if parsed_boxes is not None:
@@ -745,10 +792,8 @@ async def segment_predict(
                     max_size=max_size,
                 )
         else:
-            model, lock = get_model(model_version, model_id, automatic=False)
-            model_key = (
-                model_version,
-                model_id or _DEFAULT_MODEL_IDS[model_version],
+            model, lock, model_key = get_model(
+                model_version, model_id, automatic=False
             )
             with lock:
                 _set_image_cached(model, model_key, input_path, image_hash)
@@ -818,14 +863,13 @@ async def segment_text(
         input_path, image_hash = await _save_upload(file, tmpdir)
         output_path = os.path.join(tmpdir, "mask.tif")
 
-        model, lock = get_model(
+        model, lock, model_key = get_model(
             "sam3",
             model_id,
             backend=backend,
             confidence_threshold=confidence_threshold,
         )
         t_start = time.time()
-        model_key = ("sam3", model_id or _DEFAULT_MODEL_IDS["sam3"])
         with lock:
             _set_image_cached(model, model_key, input_path, image_hash)
             model.generate_masks(

samgeo/common.py

@@ -1444,7 +1444,14 @@ def raster_to_vector(source, output, simplify_tolerance=None, dst_crs=None, **kw
         for i in fc:
             i["geometry"] = i["geometry"].simplify(tolerance=simplify_tolerance)
 
-    gdf = gpd.GeoDataFrame.from_features(fc)
+    if fc:
+        gdf = gpd.GeoDataFrame.from_features(fc)
+    else:
+        # No foreground pixels in the mask. Build an empty GeoDataFrame that
+        # still carries a geometry column so set_crs/to_crs/to_file work and we
+        # emit a valid (empty) layer instead of raising "no active geometry
+        # column" from geopandas.
+        gdf = gpd.GeoDataFrame({"value": []}, geometry=[])
     if src.crs is not None:
         gdf.set_crs(crs=src.crs, inplace=True)
 

tests/test_api.py

@@ -107,20 +107,41 @@ def test_predict_missing_prompts(sample_image_path):
     assert "point_coords or boxes" in response.json()["detail"]
 
 
-def test_predict_sam3_rejected(sample_image_path):
-    """Test that predict endpoint rejects SAM3."""
+@patch("samgeo.api.get_model")
+def test_predict_sam3_points_accepted(mock_get, sample_image_path):
+    """SAM3 point prompts are accepted by /segment/predict (added in #508).
+
+    SAM3 support was added to the predict endpoint, so a point prompt should
+    be dispatched to ``predict_inst`` and produce a mask rather than being
+    rejected. The model is mocked so the test stays fast and offline.
+    """
+    from threading import Lock
+
+    mock_model = MagicMock()
+    mock_model.masks = [np.ones((64, 64), dtype=np.uint8)]
+
+    def fake_save_masks(output, **kwargs):
+        from PIL import Image
+
+        Image.fromarray(np.zeros((64, 64), dtype=np.uint8)).save(output)
+
+    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/predict",
             files={"file": ("test.tif", f, "image/tiff")},
             data={
                 "model_version": "sam3",
-                "point_coords": "[[100, 200]]",
+                "point_coords": "[[10, 20]]",
                 "point_labels": "[1]",
+                "output_format": "png",
             },
         )
-    assert response.status_code == 400
-    assert "SAM3" in response.json()["detail"]
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "image/png"
+    mock_model.predict_inst.assert_called_once()
 
 
 def test_invalid_output_format(sample_image_path):
@@ -149,7 +170,7 @@ def fake_generate(source, output, **kwargs):
         img.save(output)
 
     mock_model.generate = fake_generate
-    mock_get.return_value = (mock_model, Lock())
+    mock_get.return_value = (mock_model, Lock(), ("sam2", "sam2-hiera-large", True))
 
     with open(sample_image_path, "rb") as f:
         response = client.post(
@@ -159,3 +180,95 @@ def fake_generate(source, output, **kwargs):
         )
     assert response.status_code == 200
     assert response.headers["content-type"] == "image/png"
+
+
+def test_get_model_caches_automatic_and_predict_separately():
+    """Regression: automatic and predict instances must not share a cache slot.
+
+    A model built for automatic mask generation has no image ``predictor``;
+    reusing it for a predict request crashed with
+    ``'SamGeo2' object has no attribute 'predictor'``. The cache key now
+    includes the ``automatic`` flag so the two are distinct instances.
+    """
+    import samgeo.api as api
+
+    with patch("samgeo.samgeo2.SamGeo2") as mock_cls:
+        mock_cls.side_effect = lambda **kw: MagicMock()
+        auto_model, _, auto_key = api.get_model("sam2", "sam2-hiera-tiny")
+        pred_model, _, pred_key = api.get_model(
+            "sam2", "sam2-hiera-tiny", automatic=False
+        )
+
+    assert auto_key != pred_key
+    assert auto_model is not pred_model
+
+    # A repeat request for the same key is served from cache (no new instance).
+    with patch("samgeo.samgeo2.SamGeo2") as mock_again:
+        cached_model, _, _ = api.get_model("sam2", "sam2-hiera-tiny")
+        mock_again.assert_not_called()
+    assert cached_model is auto_model
+
+
+def test_get_model_caches_distinct_configs_separately():
+    """Requests differing only by model-shaping kwargs get separate instances.
+
+    Regression: the cache key once ignored constructor kwargs, so a second
+    SAM3 text request with a different ``confidence_threshold`` (or a SAM2
+    request with a different ``points_per_side``) silently reused the first,
+    differently-configured instance. The key now folds in the kwargs.
+    """
+    import samgeo.api as api
+
+    # SAM3: two text requests differing only by confidence_threshold.
+    with patch("samgeo.samgeo3.SamGeo3") as mock_sam3:
+        mock_sam3.side_effect = lambda **kw: MagicMock()
+        m_a, _, key_a = api.get_model("sam3", backend="meta", confidence_threshold=0.4)
+        m_b, _, key_b = api.get_model("sam3", backend="meta", confidence_threshold=0.6)
+        m_a2, _, _ = api.get_model("sam3", backend="meta", confidence_threshold=0.4)
+    assert key_a != key_b
+    assert m_a is not m_b
+    assert m_a2 is m_a  # identical config is a cache hit
+
+    # SAM2: two automatic requests differing only by points_per_side.
+    with patch("samgeo.samgeo2.SamGeo2") as mock_sam2:
+        mock_sam2.side_effect = lambda **kw: MagicMock()
+        s_a, _, k_a = api.get_model("sam2", "sam2-hiera-tiny", points_per_side=16)
+        s_b, _, k_b = api.get_model("sam2", "sam2-hiera-tiny", points_per_side=32)
+    assert k_a != k_b
+    assert s_a is not s_b
+
+
+def test_freeze_kwargs_is_order_independent():
+    """_freeze_kwargs yields the same hashable key regardless of arg order."""
+    from samgeo.api import _freeze_kwargs
+
+    a = _freeze_kwargs({"backend": "meta", "sam_kwargs": {"x": 1, "y": 2}})
+    b = _freeze_kwargs({"sam_kwargs": {"y": 2, "x": 1}, "backend": "meta"})
+    assert a == b
+    assert hash(a) == hash(b)
+    assert _freeze_kwargs({"a": 1}) != _freeze_kwargs({"a": 2})
+
+
+def test_set_image_cached_skips_for_sam2_but_not_sam3():
+    """The image-encoder skip is safe for SAM/SAM2 but must be off for SAM3.
+
+    Regression: SAM3 mutates its encoded-image state during ``generate_*``, so
+    reusing it on a second request without re-encoding fails with
+    ``expected scalar type BFloat16 but found Float``. ``_set_image_cached``
+    therefore always re-encodes for SAM3, while still caching for SAM/SAM2.
+    """
+    from samgeo.api import _set_image_cached, _image_hash_cache
+
+    # SAM2: second call with the same hash is skipped (set_image not re-run).
+    sam2_model = MagicMock()
+    sam2_key = ("sam2", "sam2-hiera-tiny", False)
+    assert _set_image_cached(sam2_model, sam2_key, "/tmp/x.tif", "hash-a") is True
+    assert _set_image_cached(sam2_model, sam2_key, "/tmp/x.tif", "hash-a") is False
+    assert sam2_model.set_image.call_count == 1
+
+    # SAM3: every call re-encodes even when the hash matches.
+    sam3_model = MagicMock()
+    sam3_key = ("sam3", "facebook/sam3", True)
+    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