Expose GeoText embeddings and region boxes

infer.py

@@ -136,14 +136,16 @@ def load_model(source_dir: Path, model_dir: Path, device: torch.device):
 
 
 @torch.inference_mode()
-def rank_queries(
+def model_outputs(
     model,
     tokenizer,
     config: Dict[str, Any],
     image: Image.Image,
     queries: List[str],
     device: torch.device,
-):
+    include_embeddings: bool = False,
+    include_bboxes: bool = False,
+) -> Dict[str, Any]:
     image_transform = transforms.Compose(
         [
             transforms.Resize(
@@ -166,14 +168,50 @@ def rank_queries(
     text_embeds = model.get_text_embeds(text_input.input_ids, text_input.attention_mask)
     image_features, text_features = model.get_features(image_embeds, text_embeds)
     scores = (image_features @ text_features.T).flatten().float().cpu().tolist()
-    return sorted(
-        (
-            {"query": query, "similarity": score}
-            for query, score in zip(queries, scores)
-        ),
-        key=lambda item: item["similarity"],
-        reverse=True,
-    )
+    result = {
+        "ranked_queries": sorted(
+            (
+                {"query": query, "similarity": score}
+                for query, score in zip(queries, scores)
+            ),
+            key=lambda item: item["similarity"],
+            reverse=True,
+        )
+    }
+
+    if include_embeddings:
+        result["image_embedding"] = image_features[0].float().cpu().tolist()
+        result["text_embeddings"] = [
+            {"query": query, "embedding": embedding.float().cpu().tolist()}
+            for query, embedding in zip(queries, text_features)
+        ]
+
+    if include_bboxes:
+        repeated_image_embeds = image_embeds.expand(len(queries), -1, -1)
+        boxes = model.predict_bbox(
+            repeated_image_embeds, text_embeds, text_input.attention_mask
+        ).float().cpu().tolist()
+        result["text_conditioned_boxes"] = [
+            {
+                "query": query,
+                "format": "normalized_cx_cy_width_height",
+                "box": box,
+            }
+            for query, box in zip(queries, boxes)
+        ]
+
+    return result
+
+
+def rank_queries(
+    model,
+    tokenizer,
+    config: Dict[str, Any],
+    image: Image.Image,
+    queries: List[str],
+    device: torch.device,
+):
+    return model_outputs(model, tokenizer, config, image, queries, device)["ranked_queries"]
 
 
 def main() -> None:

server.py

@@ -9,13 +9,17 @@ from typing import Any, Dict, List
 
 import torch
 from fastapi import FastAPI, HTTPException
-from infer import fetch_rgb_geotiff, load_model, rank_queries
+from infer import fetch_rgb_geotiff, load_model, model_outputs
 from pydantic import BaseModel, Field
 
 
 class InferenceRequest(BaseModel):
     image_url: str = Field(..., description="Public RGB GeoTIFF/COG URL")
     queries: List[str] = Field(..., min_items=1, description="Text queries to rank")
+    include_embeddings: bool = Field(
+        False, description="Return 256-dimensional image and text vectors"
+    )
+    include_bboxes: bool = Field(False, description="Return text-conditioned normalized boxes")
 
 
 @asynccontextmanager
@@ -47,18 +51,23 @@ def infer(request: InferenceRequest) -> Dict[str, Any]:
     try:
         image, image_metadata = fetch_rgb_geotiff(request.image_url)
         runtime = app.state.runtime
-        return {
+        result = {
             "image_url": request.image_url,
             "image_metadata": image_metadata,
             "device": str(runtime["device"]),
-            "ranked_queries": rank_queries(
+        }
+        result.update(
+            model_outputs(
                 runtime["model"],
                 runtime["tokenizer"],
                 runtime["config"],
                 image,
                 request.queries,
                 runtime["device"],
-            ),
-        }
+                include_embeddings=request.include_embeddings,
+                include_bboxes=request.include_bboxes,
+            )
+        )
+        return result
     except Exception as error:
         raise HTTPException(status_code=422, detail=str(error)) from error