Add REST API for image segmentation (#505)

giswqs · web-flow · commit 87a2dd65ca7e · 2026-03-14T00:12:36.000-04:00
* Add REST API for image segmentation

Add a FastAPI-based REST API module (samgeo/api.py) that exposes
endpoints for automatic, prompt-based, and text-based segmentation
using SAM, SAM2, and SAM3 models. Includes model caching, multiple
output formats (GeoJSON, GeoTIFF, PNG), and a console script entry
point (samgeo-api). Also adds an 'api' optional dependency group
and tests.

* Fix max_size=0 being treated as a size filter

When max_size=0 is sent via the API (e.g., from Swagger UI defaults),
it filtered out all masks since every mask has size &gt; 0. Normalize
max_size &lt;= 0 to None (no limit) in all segmentation endpoints.

* Cache image encoding to skip redundant set_image calls

Compute a SHA-256 hash of uploaded file bytes and track the last
image encoded per model. When the same image is sent again (e.g.,
trying different prompts or parameters), the expensive image encoder
forward pass is skipped entirely.

* Add logging for model and image cache hits

Log messages indicate whether a model was loaded fresh or served
from cache, and whether image encoding was performed or skipped
due to a cache hit.

* Fix cache log messages not appearing in uvicorn output

Use uvicorn.error logger instead of module-level logger so cache
hit/miss messages are visible in the server console output.

* Log inference time for each segmentation request

Show elapsed time in seconds for automatic, prompt-based, and
text segmentation endpoints. For text segmentation, the prompt
is also included in the log line.

* Address Copilot review comments

- Validate model_id against _AVAILABLE_MODELS early (400 error)
- Stream uploads to disk in chunks to reduce memory pressure
- Validate output_format before running segmentation
- Clean up temp directories on error and after JSON responses
- Fix pip install hint to use correct extras (samgeo/samgeo2/samgeo3)
- Add global lock around model cache to prevent duplicate loading
- Validate --preload CLI format with friendly error message
- Move importorskip before samgeo.api import in tests
- Use pytest tmp_path fixture for automatic temp cleanup
- Fix test name/docstring mismatch (GeoTIFF -&gt; PNG)
- Make invalid output_format test deterministic (assert 400)
- Add test for invalid model_id validation

* Fix image cache hit failing when source path is stale

When set_image() is skipped due to cache hit, model.source still
references the previous request's temp file (now cleaned up).
Update model.source to the current request's file path even when
skipping the encoding step.

* Add REST API documentation and update feature lists

- Add docs/api.md with full API documentation (endpoints, caching, examples)
- Add REST API entry to mkdocs.yml navigation and API Reference
- Add REST API feature bullet to README.md and docs/index.md
- Add segment-geospatial[api] extra to install lists
diff --git a/README.md b/README.md
@@ -38,6 +38,7 @@ The **SamGeo** package draws its inspiration from [segment-anything-eo](https://
 -   Save input prompts as GeoJSON files
 -   Visualize segmentation results on interactive maps
 -   Segment objects from timeseries remote sensing imagery
+-   REST API for serving segmentation over HTTP (see [API docs](https://samgeo.gishub.org/api))
 
 ## QGIS Plugin
 
@@ -86,6 +87,7 @@ Depending on what tools you need to use, you might want to do:
 -   `segment-geospatial[text]`: Installs Grounding DINO to use SAMGeo 1 and 2 with text prompts
 -   `segment-geospatial[fer]`: Installs the dependencies to run the feature
     edge reconstruction algorithm
+-   `segment-geospatial[api]`: Installs FastAPI and Uvicorn for serving segmentation as a REST API
 
 Additionally, these other two optional imports are defined:
 
diff --git a/docs/api.md b/docs/api.md
@@ -0,0 +1,224 @@
+# REST API
+
+segment-geospatial includes a built-in REST API powered by [FastAPI](https://fastapi.tiangolo.com/) that allows you to run image segmentation over HTTP. This is useful for integrating segmentation into web applications, pipelines, and non-Python clients.
+
+## Installation
+
+Install the API dependencies with the `api` extra:
+
+```bash
+pip install "segment-geospatial[api]"
+```
+
+To also install a specific SAM model backend, combine extras:
+
+```bash
+pip install "segment-geospatial[api,samgeo3]"
+```
+
+## Starting the Server
+
+Use the `samgeo-api` command:
+
+```bash
+samgeo-api
+```
+
+Options:
+
+```bash
+samgeo-api --host 0.0.0.0 --port 8000        # Custom host/port
+samgeo-api --preload sam2:sam2-hiera-large     # Preload a model at startup
+samgeo-api --reload                            # Auto-reload for development
+```
+
+Alternatively, use `uvicorn` directly:
+
+```bash
+uvicorn samgeo.api:app --host 0.0.0.0 --port 8000
+```
+
+Once running, interactive API docs (Swagger UI) are available at [http://localhost:8000/docs](http://localhost:8000/docs).
+
+## Endpoints
+
+### Health Check
+
+```
+GET /health
+```
+
+Returns the server status and version.
+
+```bash
+curl http://localhost:8000/health
+```
+
+```json
+{"status": "ok", "version": "1.2.3"}
+```
+
+### List Models
+
+```
+GET /models
+```
+
+Returns available model versions/IDs and which models are currently loaded in memory.
+
+```bash
+curl http://localhost:8000/models
+```
+
+### Clear Models
+
+```
+DELETE /models
+```
+
+Clears the model cache and frees GPU memory.
+
+```bash
+curl -X DELETE http://localhost:8000/models
+```
+
+### Automatic Segmentation
+
+```
+POST /segment/automatic
+```
+
+Runs automatic mask generation on an uploaded image. Supports SAM, SAM2, and SAM3.
+
+**Parameters (multipart form):**
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `file` | file | required | Image file (TIFF, PNG, JPEG) |
+| `model_version` | string | `sam2` | One of `sam`, `sam2`, `sam3` |
+| `model_id` | string | auto | Model identifier (e.g., `sam2-hiera-large`) |
+| `output_format` | string | `geojson` | One of `geojson`, `geotiff`, `png` |
+| `foreground` | bool | `true` | Extract foreground objects only |
+| `unique` | bool | `true` | Assign unique ID to each object |
+| `min_size` | int | `0` | Minimum mask size in pixels |
+| `max_size` | int | none | Maximum mask size in pixels |
+| `points_per_side` | int | `32` | Points sampled per side (SAM/SAM2) |
+| `pred_iou_thresh` | float | `0.8` | IoU threshold for filtering |
+| `stability_score_thresh` | float | `0.95` | Stability score threshold |
+
+**Example:**
+
+```bash
+curl -X POST http://localhost:8000/segment/automatic \
+  -F "file=@image.tif" \
+  -F "model_version=sam2" \
+  -F "output_format=geojson"
+```
+
+### Prompt-based Segmentation
+
+```
+POST /segment/predict
+```
+
+Runs segmentation with point or bounding box prompts. Supports SAM and SAM2.
+
+**Parameters (multipart form):**
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `file` | file | required | Image file (TIFF, PNG, JPEG) |
+| `model_version` | string | `sam2` | One of `sam`, `sam2` |
+| `model_id` | string | auto | Model identifier |
+| `output_format` | string | `geojson` | One of `geojson`, `geotiff`, `png` |
+| `point_coords` | string | none | JSON array of `[[x, y], ...]` |
+| `point_labels` | string | none | JSON array of `[1, 0, ...]` (1=foreground, 0=background) |
+| `boxes` | string | none | JSON array of `[[xmin, ymin, xmax, ymax], ...]` |
+| `point_crs` | string | none | CRS string (e.g., `EPSG:4326`) |
+| `multimask_output` | bool | `false` | Return multiple masks per prompt |
+
+**Example with point prompts:**
+
+```bash
+curl -X POST http://localhost:8000/segment/predict \
+  -F "file=@image.tif" \
+  -F "point_coords=[[100, 200]]" \
+  -F "point_labels=[1]" \
+  -F "output_format=geojson"
+```
+
+**Example with box prompts:**
+
+```bash
+curl -X POST http://localhost:8000/segment/predict \
+  -F "file=@image.tif" \
+  -F "boxes=[[10, 20, 300, 400]]" \
+  -F "output_format=geotiff"
+```
+
+### Text-prompt Segmentation
+
+```
+POST /segment/text
+```
+
+Runs text-prompt segmentation using SAM3.
+
+**Parameters (multipart form):**
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `file` | file | required | Image file (TIFF, PNG, JPEG) |
+| `prompt` | string | required | Text description (e.g., `building`, `tree`) |
+| `model_id` | string | auto | SAM3 model identifier |
+| `backend` | string | `meta` | One of `meta`, `transformers` |
+| `output_format` | string | `geojson` | One of `geojson`, `geotiff`, `png` |
+| `confidence_threshold` | float | `0.5` | Detection confidence threshold |
+| `min_size` | int | `0` | Minimum mask size in pixels |
+| `max_size` | int | none | Maximum mask size in pixels |
+
+**Example:**
+
+```bash
+curl -X POST http://localhost:8000/segment/text \
+  -F "file=@image.tif" \
+  -F "prompt=building" \
+  -F "output_format=geojson"
+```
+
+## Caching
+
+The API automatically caches models and image encodings for better performance:
+
+- **Model cache**: Models are loaded once and reused across requests. Use `DELETE /models` to free GPU memory.
+- **Image cache**: When the same image is sent multiple times (e.g., with different prompts), the expensive image encoding step is skipped. This makes subsequent requests significantly faster.
+
+Example timing with a 13 MB GeoTIFF:
+
+| Request | Description | Time |
+|---------|-------------|------|
+| 1st | Model load + image encoding | ~7s |
+| 2nd | Same image, different prompt | ~0.4s |
+| 3rd | Same image, another prompt | ~0.2s |
+
+## Python Client Example
+
+```python
+import requests
+
+url = "http://localhost:8000/segment/text"
+
+with open("image.tif", "rb") as f:
+    response = requests.post(
+        url,
+        files={"file": ("image.tif", f, "image/tiff")},
+        data={"prompt": "building", "output_format": "geojson"},
+    )
+
+geojson = response.json()
+print(f"Found {len(geojson['features'])} features")
+```
+
+## API Reference
+
+::: samgeo.api
diff --git a/docs/index.md b/docs/index.md
@@ -38,6 +38,7 @@ The **SamGeo** package draws its inspiration from [segment-anything-eo](https://
 -   Save input prompts as GeoJSON files
 -   Visualize segmentation results on interactive maps
 -   Segment objects from timeseries remote sensing imagery
+-   REST API for serving segmentation over HTTP (see [API docs](https://samgeo.gishub.org/api))
 
 ## QGIS Plugin
 
@@ -61,6 +62,7 @@ Depending on what tools you need to use, you might want to do:
 -   `segment-geospatial[text]`: Installs Grounding DINO to use SAMGeo 1 and 2 with text prompts
 -   `segment-geospatial[fer]`: Installs the dependencies to run the feature
     edge reconstruction algorithm
+-   `segment-geospatial[api]`: Installs FastAPI and Uvicorn for serving segmentation as a REST API
 
 Additionally, these other two optional imports are defined:
 
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -44,6 +44,7 @@ nav:
     - Home: index.md
     - Installation: installation.md
     - Usage: usage.md
+    - REST API: api.md
     - Contributing: contributing.md
     - FAQ: faq.md
     - Changelog: changelog.md
@@ -98,4 +99,5 @@ nav:
           - hq_sam module: hq_sam.md
           - text_sam module: text_sam.md
           - detectree2 module: detectree2.md
+          - api module: api.md
           #   - fer module: fer.md
diff --git a/pyproject.toml b/pyproject.toml
@@ -86,8 +86,14 @@ text = [
     "segment_geospatial[samgeo2]",
     "groundingdino-py",
 ]
+api = [
+    "segment_geospatial[core]",
+    "fastapi>=0.100.0",
+    "uvicorn[standard]>=0.20.0",
+    "python-multipart>=0.0.6",
+]
 all = [
-    "segment-geospatial[fast, hq, fer, samgeo, samgeo2, samgeo3, text]",
+    "segment-geospatial[fast, hq, fer, samgeo, samgeo2, samgeo3, text, api]",
 ]
 extra = [
     "leafmap",
@@ -133,5 +139,8 @@ max-line-length = 88
 
 [tool.setuptools_scm]
 
+[project.scripts]
+samgeo-api = "samgeo.api:main"
+
 [project.urls]
 Homepage = "https://github.com/opengeos/segment-geospatial"
diff --git a/samgeo/api.py b/samgeo/api.py
diff --git a/tests/test_api.py b/tests/test_api.py
