Add dedicated text and image embedding endpoints

EXAMPLES.md

@@ -0,0 +1,211 @@
+# GeoText-1652 request examples
+
+All image routes require a public RGB GeoTIFF/COG URL. Use this image in the
+examples below:
+
+```text
+https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif
+```
+
+## Inputs and outcomes at a glance
+
+| Workflow | Inputs | Primary returned fields | What to do next |
+| --- | --- | --- | --- |
+| Embed text | `queries` | `text_embeddings` | Search global or patch vector collections. |
+| Embed one image | `image_url` | `image_embedding` | Store in an image index for image-to-image retrieval. |
+| Embed image patches | `image_url`, `include_patch_embeddings` | `patch_embeddings.patches` | Search each footprint for text-guided localization. |
+| Rank text descriptions | `image_url`, `queries` | `ranked_queries` | Determine which description best matches the image. |
+| Rank a large raster | image, queries, `tile_size` | `ranked_queries` with `tile_index` | Find the best tile for each query. |
+| Propose matching regions | image, queries, `include_bboxes` | `text_conditioned_boxes` / `stitched_text_conditioned_boxes` | Inspect semantic candidate regions. |
+| Persist patch index data | image, `store_patch_embeddings` | `patch_embeddings_storage` | Download Parquet and run vector similarity search. |
+
+## 1. Health check
+
+```bash
+curl https://YOUR-ENDPOINT.endpoints.huggingface.cloud/health
+```
+
+## 2. Text-only embeddings
+
+`POST /embed/text`
+
+```json
+{
+  "queries": ["a parking lot", "a large building", "a sports field"]
+}
+```
+
+Returns one normalized 256-dimensional embedding per query. This is the route
+to use when searching a previously stored image or patch-vector collection.
+
+## 3. Global image embedding
+
+`POST /embed/image`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "include_global_embedding": true
+}
+```
+
+## 4. Image patch embeddings
+
+`POST /embed/image`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "include_global_embedding": false,
+  "include_patch_embeddings": true
+}
+```
+
+Each patch includes `source_pixel_xyxy` and a normalized 256-dimensional
+embedding. The model returns a 12 by 12 patch grid for one 384 px model input.
+
+## 5. Tiled global and patch image embeddings
+
+`POST /embed/image`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "tile_size": 384,
+  "tile_overlap": 64,
+  "include_global_embedding": true,
+  "include_patch_embeddings": true,
+  "return_tile_results": false
+}
+```
+
+The global vector is an area-weighted, L2-normalized mean of tile vectors.
+Patch vectors keep their `tile_index` and original-image pixel footprints.
+
+## 6. Store tiled patch embeddings in the Hub
+
+`POST /embed/image`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "tile_size": 384,
+  "tile_overlap": 64,
+  "include_global_embedding": true,
+  "store_patch_embeddings": true,
+  "output_prefix": "geotext/demo-image"
+}
+```
+
+Endpoint secrets/configuration required:
+
+```text
+HF_TOKEN=<write-capable token>
+HF_BUCKET=your-namespace/your-embedding-dataset
+```
+
+This stores a compressed Parquet file but leaves the large patch collection out
+of the response. Add `"include_patch_embeddings": true` to return it too.
+
+## 7. Image-level text ranking
+
+`POST /infer` (or `POST /` in the Inference Endpoints Playground)
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "queries": ["a parking lot", "a building", "a sports field"]
+}
+```
+
+Returns `ranked_queries`, sorted by descending cosine similarity. This answers
+questions such as: *which of these descriptions best matches this image?*
+
+## 8. Search with global image and text vectors
+
+`POST /infer`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "queries": ["a parking lot", "a building"],
+  "include_embeddings": true
+}
+```
+
+## 9. Large-raster ranking, localization, and stitched region proposals
+
+`POST /infer`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "queries": ["a parking lot", "a building"],
+  "tile_size": 384,
+  "tile_overlap": 64,
+  "include_embeddings": true,
+  "include_bboxes": true,
+  "return_tile_results": false
+}
+```
+
+Per-query similarity is the maximum matching tile score. Region proposals are
+converted to original-image pixels and deduplicated with non-maximum
+suppression.
+
+The `tile_index` in `ranked_queries` identifies the strongest tile for each
+query. Add `return_tile_results: true` to receive all tile scores and inspect
+the full ranking surface.
+
+## 10. Inspect every tile
+
+`POST /infer`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "queries": ["a parking lot"],
+  "tile_size": 384,
+  "tile_overlap": 64,
+  "include_embeddings": true,
+  "include_patch_embeddings": true,
+  "include_bboxes": true,
+  "return_tile_results": true
+}
+```
+
+This is for debugging only: it returns all tile results and can be very large.
+
+## 11. Text-to-patch localization from stored embeddings
+
+This is a two-step workflow. First create or retrieve a text vector:
+
+`POST /embed/text`
+
+```json
+{
+  "queries": ["a parking lot"]
+}
+```
+
+Then compute dot-product similarity between that 256-D vector and the
+`embedding` column in `patch_embeddings.parquet`. Sort descending and use each
+result's `source_pixel_xyxy` field to draw the matching locations on the
+original image. This is the scalable text-guided localization workflow.
+
+## 12. Image-to-image retrieval
+
+Create a global vector for a query image:
+
+`POST /embed/image`
+
+```json
+{
+  "image_url": "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif",
+  "include_global_embedding": true
+}
+```
+
+Store `image_embedding` in a vector database or local matrix with global
+vectors from other images. Rank candidates by dot product to find visually and
+semantically similar images.

INFERENCE.md

@@ -5,8 +5,53 @@ the official [GeoText-1652 source](https://github.com/MultimodalGeo/GeoText-1652
 and checkpoint.
 
 The image input must be an RGB GeoTIFF/COG with at least three bands. GeoText
-ranks text descriptions against the image; it is not an object-detection model
-and does not return building polygons or bounding boxes.
+is a text-image retrieval model. It returns similarity scores and embeddings;
+its optional text-conditioned boxes are semantic region proposals, not
+pixel-accurate object-detection polygons.
+
+## Use cases, inputs, and outcomes
+
+| Use case | Route | Required inputs | Optional inputs | Outcome |
+| --- | --- | --- | --- | --- |
+| Text embedding creation | `/embed/text` | `queries` | — | One normalized 256-D vector per text query. |
+| Global image embedding creation | `/embed/image` | `image_url` | `tile_size`, `tile_overlap` | One normalized 256-D vector representing the full image. |
+| Patch embedding creation | `/embed/image` | `image_url` | `include_patch_embeddings`, tiling | One 256-D vector per spatial patch plus its source-pixel footprint. |
+| Image-level text ranking | `/infer` | `image_url`, `queries` | `include_embeddings` | Every submitted query ranked by image-text similarity. |
+| Tile-level text ranking | `/infer` | `image_url`, `queries`, `tile_size` | `return_tile_results` | Best score and winning tile for each query; optionally every tile score. |
+| Text-guided localization | `/embed/text` + stored patch vectors | query text and a Parquet/vector index | top-k | Highest-similarity patch footprints showing where the text concept occurs. |
+| Online text-guided region proposals | `/infer` | `image_url`, `queries` | `include_bboxes`, tiling | Text-conditioned candidate boxes in source-image pixels; not segmentation polygons. |
+| Image-to-image retrieval | `/embed/image` + global-vector index | query `image_url` | top-k in downstream index | A global image vector for nearest-neighbor retrieval outside the endpoint. |
+| Patch-to-patch / change exploration | `/embed/image` + patch-vector index | query image or patch | tiling, top-k | Patch vectors for downstream spatial nearest-neighbor comparison. |
+| Persistent embedding dataset | `/embed/image` | `image_url`, `store_patch_embeddings` | `output_prefix`, tiling | Compressed Parquet stored in the configured Hugging Face dataset repository. |
+
+### Outcome terminology
+
+- **Embedding:** a normalized 256-dimensional vector. Compare vectors with dot
+  product, which equals cosine similarity because GeoText normalizes them.
+- **Ranking:** a descending list of supplied text queries scored against one
+  image or its tiles. Scores are relative similarity, not probabilities.
+- **Localization:** downstream similarity search between a text embedding and
+  stored patch embeddings. The returned `source_pixel_xyxy` footprint locates
+  the match in the original raster.
+- **Region proposal:** the optional `include_bboxes` output. It is a
+  text-conditioned semantic box, not a detected object instance or a polygon.
+- **Stitching:** combines tiled outputs. Search uses each query's maximum tile
+  score; global image embeddings use an area-weighted mean and L2
+  normalization; overlapping proposal boxes are deduplicated by NMS.
+
+## API routes
+
+| Route | Purpose | Required input |
+| --- | --- | --- |
+| `POST /embed/text` | Text-only 256-D embeddings | `queries` |
+| `POST /embed/image` | Image-only global and/or patch embeddings | `image_url` |
+| `POST /infer` or `POST /` | Combined text-guided search | `image_url`, `queries` |
+| `GET /health` | Endpoint readiness probe | none |
+
+`/embed/text` does not fetch an image or run the vision encoder. `/embed/image`
+does not tokenize text or run the text encoder. The combined search route is
+kept for Playground compatibility. See [EXAMPLES.md](EXAMPLES.md) for every
+request payload.
 
 ## Build a checkpoint-free image
 

infer.py

@@ -166,6 +166,91 @@ def load_model(source_dir: Path, model_dir: Path, device: torch.device):
     return model.to(device).eval(), tokenizer, config
 
 
+def _image_tensor(image: Image.Image, config: Dict[str, Any], device: torch.device):
+    transform = transforms.Compose(
+        [
+            transforms.Resize(
+                (config["image_res"], config["image_res"]), interpolation=Image.BICUBIC
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
+        ]
+    )
+    return transform(image).unsqueeze(0).to(device)
+
+
+@torch.inference_mode()
+def text_embedding_outputs(
+    model, tokenizer, config: Dict[str, Any], queries: List[str], device: torch.device
+) -> Dict[str, Any]:
+    """Encode text without reading an image or running the vision encoder."""
+    text_input = tokenizer(
+        queries,
+        padding="max_length",
+        truncation=True,
+        max_length=config["max_tokens"],
+        return_tensors="pt",
+    ).to(device)
+    text_embeds = model.get_text_embeds(text_input.input_ids, text_input.attention_mask)
+    text_features = model.get_features(text_embeds=text_embeds)
+    return {
+        "embedding_dimension": text_features.shape[-1],
+        "text_embeddings": [
+            {"query": query, "embedding": embedding.float().cpu().tolist()}
+            for query, embedding in zip(queries, text_features)
+        ],
+    }
+
+
+@torch.inference_mode()
+def image_embedding_outputs(
+    model,
+    config: Dict[str, Any],
+    image: Image.Image,
+    device: torch.device,
+    include_global_embedding: bool = True,
+    include_patch_embeddings: bool = False,
+) -> Dict[str, Any]:
+    """Encode an image without tokenizing text or running the text encoder."""
+    image_embeds, _ = model.get_vision_embeds(_image_tensor(image, config, device))
+    result = {"embedding_dimension": model.vision_proj.out_features}
+    if include_global_embedding:
+        image_feature = model.get_features(image_embeds=image_embeds)[0]
+        result["image_embedding"] = image_feature.float().cpu().tolist()
+    if include_patch_embeddings:
+        patch_features = torch.nn.functional.normalize(
+            model.vision_proj(image_embeds[:, 1:, :]), dim=-1
+        )[0]
+        patch_count = patch_features.shape[0]
+        columns = int(round(patch_count**0.5))
+        if columns <= 0 or patch_count % columns:
+            columns = patch_count
+        rows = patch_count // columns
+        patches = []
+        for index, embedding in enumerate(patch_features.float().cpu().tolist()):
+            row, column = divmod(index, columns)
+            patches.append(
+                {
+                    "patch_index": index,
+                    "row": row,
+                    "column": column,
+                    "source_pixel_xyxy": [
+                        image.width * column / columns,
+                        image.height * row / rows,
+                        image.width * (column + 1) / columns,
+                        image.height * (row + 1) / rows,
+                    ],
+                    "embedding": embedding,
+                }
+            )
+        result["patch_embeddings"] = {
+            "embedding_dimension": patch_features.shape[-1],
+            "grid": {"rows": rows, "columns": columns},
+            "patches": patches,
+        }
+    return result
+
+
 @torch.inference_mode()
 def model_outputs(
     model,
@@ -178,16 +263,7 @@ def model_outputs(
     include_bboxes: bool = False,
     include_patch_embeddings: bool = False,
 ) -> Dict[str, Any]:
-    image_transform = transforms.Compose(
-        [
-            transforms.Resize(
-                (config["image_res"], config["image_res"]), interpolation=Image.BICUBIC
-            ),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
-        ]
-    )
-    image_tensor = image_transform(image).unsqueeze(0).to(device)
+    image_tensor = _image_tensor(image, config, device)
     text_input = tokenizer(
         queries,
         padding="max_length",

server.py

@@ -11,7 +11,14 @@ import numpy as np
 import torch
 from fastapi import FastAPI, HTTPException
 from hub_storage import persist_patch_embeddings
-from infer import fetch_rgb_geotiff, load_model, model_outputs, tile_image
+from infer import (
+    fetch_rgb_geotiff,
+    image_embedding_outputs,
+    load_model,
+    model_outputs,
+    text_embedding_outputs,
+    tile_image,
+)
 from pydantic import BaseModel, Field
 
 
@@ -41,6 +48,33 @@ class InferenceRequest(BaseModel):
     )
 
 
+class TextEmbeddingRequest(BaseModel):
+    queries: List[str] = Field(..., min_items=1, description="Text strings to embed")
+
+
+class ImageEmbeddingRequest(BaseModel):
+    image_url: str = Field(..., description="Public RGB GeoTIFF/COG URL")
+    include_global_embedding: bool = Field(
+        True, description="Return the normalized 256-dimensional image vector"
+    )
+    include_patch_embeddings: bool = Field(
+        False, description="Return normalized 256-dimensional spatial patch vectors"
+    )
+    store_patch_embeddings: bool = Field(
+        False, description="Persist patch embeddings as Parquet in HF_BUCKET"
+    )
+    output_prefix: Optional[str] = Field(
+        None, description="Optional Hugging Face dataset path prefix"
+    )
+    tile_overlap: int = Field(64, ge=0, description="Source-pixel overlap between tiles")
+    tile_size: Optional[int] = Field(
+        None, ge=64, description="Optional source-pixel tile size"
+    )
+    return_tile_results: bool = Field(
+        False, description="Include per-tile embeddings in addition to stitched output"
+    )
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -218,6 +252,89 @@ def _tiled_outputs(
     return result
 
 
+def _image_outputs(
+    runtime: Dict[str, Any], image, request: ImageEmbeddingRequest
+) -> Dict[str, Any]:
+    """Run vision-only embedding inference, optionally over overlapping tiles."""
+    include_patches = request.include_patch_embeddings or request.store_patch_embeddings
+    if request.tile_size is None:
+        return image_embedding_outputs(
+            runtime["model"],
+            runtime["config"],
+            image,
+            runtime["device"],
+            include_global_embedding=request.include_global_embedding,
+            include_patch_embeddings=include_patches,
+        )
+
+    tiles = tile_image(image, request.tile_size, request.tile_overlap)
+    tile_outputs = []
+    for tile in tiles:
+        output = image_embedding_outputs(
+            runtime["model"],
+            runtime["config"],
+            tile["image"],
+            runtime["device"],
+            include_global_embedding=request.include_global_embedding,
+            include_patch_embeddings=include_patches,
+        )
+        tile_outputs.append((tile, output))
+
+    result = {
+        "embedding_dimension": tile_outputs[0][1]["embedding_dimension"],
+        "tiling": {
+            "tile_size": request.tile_size,
+            "tile_overlap": request.tile_overlap,
+            "tile_count": len(tiles),
+            "global_embedding_stitching": "area-weighted mean, then L2 normalization",
+        },
+    }
+    if request.include_global_embedding:
+        weights = np.asarray(
+            [tile["width"] * tile["height"] for tile, _ in tile_outputs], dtype=np.float32
+        )
+        vectors = np.asarray(
+            [output["image_embedding"] for _, output in tile_outputs], dtype=np.float32
+        )
+        embedding = np.average(vectors, axis=0, weights=weights)
+        embedding /= max(float(np.linalg.norm(embedding)), 1e-8)
+        result["image_embedding"] = embedding.tolist()
+    if include_patches:
+        patches = []
+        for tile, output in tile_outputs:
+            for patch in output["patch_embeddings"]["patches"]:
+                x1, y1, x2, y2 = patch["source_pixel_xyxy"]
+                patches.append(
+                    {
+                        **patch,
+                        "tile_index": tile["index"],
+                        "source_pixel_xyxy": [
+                            x1 + tile["column"],
+                            y1 + tile["row"],
+                            x2 + tile["column"],
+                            y2 + tile["row"],
+                        ],
+                    }
+                )
+        result["patch_embeddings"] = {
+            "embedding_dimension": result["embedding_dimension"],
+            "grid": "per_tile",
+            "patches": patches,
+        }
+    if request.return_tile_results:
+        result["tile_results"] = [
+            {
+                "tile_index": tile["index"],
+                "source_pixel_window": [
+                    tile["column"], tile["row"], tile["width"], tile["height"]
+                ],
+                **output,
+            }
+            for tile, output in tile_outputs
+        ]
+    return result
+
+
 @app.get("/health")
 def health() -> Dict[str, str]:
     return {"status": "ok"}
@@ -226,6 +343,7 @@ def health() -> Dict[str, str]:
 @app.post("/")
 @app.post("/infer")
 def infer(request: InferenceRequest) -> Dict[str, Any]:
+    """Backward-compatible combined image/text search endpoint."""
     try:
         image, image_metadata = fetch_rgb_geotiff(request.image_url)
         runtime = app.state.runtime
@@ -262,3 +380,45 @@ def infer(request: InferenceRequest) -> Dict[str, Any]:
         return result
     except Exception as error:
         raise HTTPException(status_code=422, detail=str(error)) from error
+
+
+@app.post("/embed/text")
+def embed_text(request: TextEmbeddingRequest) -> Dict[str, Any]:
+    """Return query vectors without fetching imagery or running the vision encoder."""
+    try:
+        runtime = app.state.runtime
+        result = text_embedding_outputs(
+            runtime["model"],
+            runtime["tokenizer"],
+            runtime["config"],
+            request.queries,
+            runtime["device"],
+        )
+        return {"device": str(runtime["device"]), **result}
+    except Exception as error:
+        raise HTTPException(status_code=422, detail=str(error)) from error
+
+
+@app.post("/embed/image")
+def embed_image(request: ImageEmbeddingRequest) -> Dict[str, Any]:
+    """Return image or patch vectors without tokenizing text or running search."""
+    try:
+        image, image_metadata = fetch_rgb_geotiff(request.image_url)
+        runtime = app.state.runtime
+        result = {
+            "image_url": request.image_url,
+            "image_metadata": image_metadata,
+            "device": str(runtime["device"]),
+        }
+        result.update(_image_outputs(runtime, image, request))
+        if request.store_patch_embeddings:
+            result["patch_embeddings_storage"] = persist_patch_embeddings(
+                image_url=request.image_url,
+                patches=result["patch_embeddings"]["patches"],
+                output_prefix=request.output_prefix,
+            )
+            if not request.include_patch_embeddings:
+                result.pop("patch_embeddings")
+        return result
+    except Exception as error:
+        raise HTTPException(status_code=422, detail=str(error)) from error