REST API¶

@app.delete("/models")
def clear_models():
    """Clear the model cache and free GPU memory."""
    _model_cache.clear()
    _image_hash_cache.clear()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass
    return {"status": "cleared"}

def get_model(model_version: str, model_id: Optional[str] = None, **kwargs):
    """Get or create a cached model instance.

    Args:
        model_version: One of "sam", "sam2", "sam3".
        model_id: Specific model identifier. Uses default if None.
        **kwargs: Additional keyword arguments for model initialization.

    Returns:
        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
            dependencies are missing.
    """
    if model_version not in DEFAULT_MODEL_IDS:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Invalid model_version '{model_version}'. "
                f"Must be one of: {list(DEFAULT_MODEL_IDS.keys())}"
            ),
        )

    if not model_id:
        model_id = DEFAULT_MODEL_IDS[model_version]

    valid_ids = AVAILABLE_MODELS[model_version]
    if model_id not in valid_ids:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Invalid model_id '{model_id}' for {model_version}. "
                f"Must be one of: {valid_ids}"
            ),
        )

    # 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)
            model, lock = _model_cache[key]
            return model, lock, key

        logger.info("Loading model %s", key)
        extra = EXTRAS_MAP.get(model_version, model_version)
        try:
            if model_version == "sam":
                from samgeo.samgeo import SamGeo

                model = SamGeo(model_type=model_id, **kwargs)
            elif model_version == "sam2":
                from samgeo.samgeo2 import SamGeo2

                model = SamGeo2(model_id=model_id, **kwargs)
            elif model_version == "sam3":
                from samgeo.samgeo3 import SamGeo3

                kwargs.setdefault("enable_inst_interactivity", True)
                model = SamGeo3(model_id=model_id, **kwargs)
        except ValueError as e:
            msg = str(e)
            is_user_config_error = model_version == "sam3" and (
                msg.startswith("Invalid backend")
                or "only supported with backend='meta'" in msg
            )
            if is_user_config_error:
                raise HTTPException(status_code=400, detail=msg) from e
            raise
        except ImportError as e:
            raise HTTPException(
                status_code=503,
                detail=(
                    f"Dependencies for {model_version} are not installed. "
                    f"Install with: pip install segment-geospatial[{extra}]. "
                    f"Error: {e}"
                ),
            )
        _model_cache[key] = (model, threading.Lock())
        model, lock = _model_cache[key]
        return model, lock, key

@app.get("/health")
def health():
    """Health check endpoint."""
    return {"status": "ok", "version": __version__}

@app.get("/models")
def list_models():
    """List available and currently loaded models."""
    loaded = [list(key) for key in _model_cache]
    return {"models": AVAILABLE_MODELS, "loaded": loaded}

def main():
    """Entry point for the samgeo-api console script."""
    parser = argparse.ArgumentParser(
        description="Run the samgeo REST API server."
    )
    parser.add_argument(
        "--host", default="0.0.0.0", help="Host to bind to (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port", type=int, default=8000, help="Port to listen on (default: 8000)"
    )
    parser.add_argument(
        "--reload", action="store_true", help="Enable auto-reload for development"
    )
    parser.add_argument(
        "--preload",
        type=str,
        default=None,
        help="Preload a model at startup, e.g. 'sam2:sam2-hiera-large'",
    )
    args = parser.parse_args()

    if args.preload:
        if ":" not in args.preload:
            parser.error(
                "Invalid --preload format. "
                "Expected 'model_version:model_id', e.g. 'sam2:sam2-hiera-large'"
            )
        version, mid = args.preload.split(":", 1)
        get_model(version, mid)

    uvicorn.run("samgeo.api:app", host=args.host, port=args.port, reload=args.reload)

@app.post("/segment/automatic")
async def segment_automatic(
    file: UploadFile = File(...),
    model_version: str = Form("sam3"),
    model_id: Optional[str] = Form(None),
    output_format: str = Form("geojson"),
    foreground: bool = Form(True),
    unique: bool = Form(True),
    min_size: int = Form(0),
    max_size: Optional[int] = Form(None),
    points_per_side: int = Form(32),
    pred_iou_thresh: float = Form(0.8),
    stability_score_thresh: float = Form(0.95),
):
    """Run automatic mask generation on an uploaded image.

    Args:
        file: Image file (TIFF, PNG, JPEG).
        model_version: One of "sam", "sam2", "sam3".
        model_id: Specific model identifier.
        output_format: One of "geojson", "geotiff", "png".
        foreground: Whether to extract foreground objects only.
        unique: Whether to assign unique IDs to each object.
        min_size: Minimum mask size in pixels.
        max_size: Maximum mask size in pixels.
        points_per_side: Number of points sampled per side (SAM/SAM2).
        pred_iou_thresh: IoU threshold for filtering masks.
        stability_score_thresh: Stability score threshold for filtering.

    Returns:
        Segmentation result in the requested format.
    """
    _validate_output_format(output_format)
    max_size = _normalize_max_size(max_size)
    tmpdir = tempfile.mkdtemp()
    try:
        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)
            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)
                mask_scores = _snapshot_mask_scores(model)
        else:
            sam_kwargs = {
                "points_per_side": points_per_side,
                "pred_iou_thresh": pred_iou_thresh,
                "stability_score_thresh": stability_score_thresh,
            }
            if model_version == "sam":
                model, lock, _ = get_model(
                    model_version, model_id, sam_kwargs=sam_kwargs
                )
            else:
                model, lock, _ = get_model(model_version, model_id, **sam_kwargs)

            with lock:
                model.generate(
                    source=input_path,
                    output=output_path,
                    foreground=foreground,
                    unique=unique,
                    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,
            mask_scores=mask_scores,
        )
    except HTTPException:
        _cleanup_tmpdir(tmpdir)
        raise
    except Exception as e:
        _cleanup_tmpdir(tmpdir)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/segment/predict")
async def segment_predict(
    file: UploadFile = File(...),
    model_version: str = Form("sam3"),
    model_id: Optional[str] = Form(None),
    output_format: str = Form("geojson"),
    point_coords: Optional[str] = Form(None),
    point_labels: Optional[str] = Form(None),
    boxes: Optional[str] = Form(None),
    point_crs: Optional[str] = Form(None),
    multimask_output: bool = Form(False),
    min_size: int = Form(0),
    max_size: Optional[int] = Form(None),
):
    """Run prompt-based segmentation with points or bounding boxes.

    For SAM3 with bounding box prompts, the model finds all similar objects
    in the image (not just the object inside the box). Point prompts with
    SAM3 segment the specific object at the point location.

    Args:
        file: Image file (TIFF, PNG, JPEG).
        model_version: One of "sam", "sam2", "sam3".
        model_id: Specific model identifier.
        output_format: One of "geojson", "geotiff", "png", "json", "detections".
        point_coords: JSON string of [[x, y], ...] coordinate pairs.
        point_labels: JSON string of [1, 0, ...] labels (1=foreground,
            0=background).
        boxes: JSON string of [[xmin, ymin, xmax, ymax], ...] bounding boxes.
        point_crs: CRS string (e.g., "EPSG:4326") for point/box coordinates.
        multimask_output: Whether to return multiple masks per prompt.
        min_size: Minimum mask size in pixels.
        max_size: Maximum mask size in pixels.

    Returns:
        Segmentation result in the requested format.
    """
    _validate_output_format(output_format)

    # Swagger UI sends empty strings for unfilled optional fields
    if not point_coords:
        point_coords = None
    if not point_labels:
        point_labels = None
    if not boxes:
        boxes = None
    if not point_crs:
        point_crs = None

    if point_coords is None and boxes is None:
        raise HTTPException(
            status_code=400,
            detail="At least one of point_coords or boxes must be provided.",
        )

    max_size = _normalize_max_size(max_size)
    tmpdir = tempfile.mkdtemp()
    try:
        input_path, image_hash = await _save_upload(file, tmpdir)
        output_path = os.path.join(tmpdir, "mask.tif")

        # Parse JSON prompt fields
        parsed_coords = None
        parsed_labels = None
        parsed_boxes = None

        if point_coords is not None:
            parsed_coords = np.array(json.loads(point_coords))
        if point_labels is not None:
            parsed_labels = np.array(json.loads(point_labels))
        if boxes is not None:
            parsed_boxes = np.array(json.loads(boxes))

        mask_scores = None
        t_start = time.time()

        if model_version == "sam3":
            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:
                    # Use generate_masks_by_boxes to find all similar objects
                    box_list = parsed_boxes.tolist()
                    if parsed_boxes.ndim == 1:
                        box_list = [box_list]
                    model.generate_masks_by_boxes(
                        boxes=box_list,
                        box_crs=point_crs,
                        min_size=min_size,
                        max_size=max_size,
                    )
                else:
                    # Use predict_inst for point-only prompts
                    model.predict_inst(
                        point_coords=parsed_coords,
                        point_labels=parsed_labels,
                        multimask_output=multimask_output,
                        point_crs=point_crs,
                    )
                if model.masks is None or len(model.masks) == 0:
                    _cleanup_tmpdir(tmpdir)
                    raise HTTPException(
                        status_code=404,
                        detail="No objects found for the given prompts.",
                    )
                model.save_masks(
                    output=output_path,
                    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
            )
            with lock:
                _set_image_cached(model, model_key, input_path, image_hash)
                model.predict(
                    point_coords=parsed_coords,
                    point_labels=parsed_labels,
                    boxes=parsed_boxes,
                    point_crs=point_crs,
                    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,
            mask_scores=mask_scores,
        )
    except HTTPException:
        _cleanup_tmpdir(tmpdir)
        raise
    except json.JSONDecodeError as e:
        _cleanup_tmpdir(tmpdir)
        raise HTTPException(
            status_code=400, detail=f"Invalid JSON in prompt fields: {e}"
        )
    except Exception as e:
        _cleanup_tmpdir(tmpdir)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/segment/text")
async def segment_text(
    file: UploadFile = File(...),
    prompt: str = Form(...),
    model_id: Optional[str] = Form(None),
    backend: str = Form("meta"),
    output_format: str = Form("geojson"),
    confidence_threshold: float = Form(0.5),
    min_size: int = Form(0),
    max_size: Optional[int] = Form(None),
):
    """Run text-prompt segmentation using SAM3.

    Args:
        file: Image file (TIFF, PNG, JPEG).
        prompt: Text description of objects to segment (e.g., "building").
        model_id: SAM3 model identifier.
        backend: SAM3 backend, one of "meta" or "transformers".
        output_format: One of "geojson", "geotiff", "png", "detections", "json".
            Use "detections" to get a GeoJSON FeatureCollection of bounding
            box polygons in geographic coordinates with confidence scores.
            Use "json" for a plain JSON array of bounding boxes in pixel
            coordinates, suitable for non-georeferenced images.
        confidence_threshold: Confidence threshold for detections.
        min_size: Minimum mask size in pixels.
        max_size: Maximum mask size in pixels.

    Returns:
        Segmentation result in the requested format.
    """
    _validate_output_format(output_format)
    max_size = _normalize_max_size(max_size)
    tmpdir = tempfile.mkdtemp()
    try:
        input_path, image_hash = await _save_upload(file, tmpdir)
        output_path = os.path.join(tmpdir, "mask.tif")

        model, lock, model_key = get_model(
            "sam3",
            model_id,
            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)
            model.generate_masks(
                prompt=prompt,
                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 the given prompt. "
                        "Please try a different prompt or adjust parameters."
                    ),
                )
            if output_format in ("detections", "json"):
                if output_format == "detections":
                    det_result = _build_detections_geojson(model, input_path)
                else:
                    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(
            "Text segmentation completed in %.2fs (prompt: '%s')",
            t_inference,
            prompt,
        )
        if output_format in ("detections", "json"):
            _cleanup_tmpdir(tmpdir)
            return JSONResponse(content=det_result)
        return _format_response(
            output_path,
            output_format,
            tmpdir,
            mask_scores=mask_scores,
        )
    except HTTPException:
        _cleanup_tmpdir(tmpdir)
        raise
    except Exception as e:
        _cleanup_tmpdir(tmpdir)
        raise HTTPException(status_code=500, detail=str(e))