Add get_patch_embeddings endpoint for ROI / patch-selection search

Additive endpoint returning per-patch token embeddings (grid x grid) plus a
mean-pooled dense vector, for region-of-interest late-interaction search.
The existing /get_embedding endpoint is left untouched for backward compatibility.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -355,6 +355,118 @@ def get_embedding(image_pil: Image.Image, model_name: str, embedding_method: str
             "message": error_msg
         }
 
+# --- Per-Patch Embedding Function (ROI / patch-selection search) ---
+def get_patch_embeddings(image_pil: Image.Image, model_name: str = MULTI_FPS_MODEL_NAME) -> dict:
+    """Return per-patch token embeddings plus a mean-pooled dense embedding.
+
+    This is an additive endpoint used by the region-of-interest (ROI) search
+    flow: the client lets the user brush a grid of patches and uses the selected
+    patch tokens directly for a late-interaction (MaxSim) query. The existing
+    `/get_embedding` endpoint is intentionally left untouched for backward
+    compatibility.
+
+    Returns a dict with:
+      - model_name: the model used (must match the collection's document tokens)
+      - embedding_method: always 'mean pooling' (used for first-stage dense retrieval)
+      - data: [D] L2-normalized mean-pooled dense embedding (first stage)
+      - patch_tokens: [num_patches, D] L2-normalized per-patch tokens in row-major
+                      order (grid x grid), for user-selected second-stage query
+      - grid: int side length of the patch grid (sqrt(num_patches), e.g. 14)
+      - num_patches: int total number of patch tokens
+      - message: 'Success' or an error description
+    """
+    if image_pil is None:
+        return {
+            "model_name": model_name,
+            "embedding_method": "mean pooling",
+            "data": None,
+            "patch_tokens": None,
+            "grid": None,
+            "num_patches": None,
+            "message": "Error: Please upload an image."
+        }
+
+    # Fall back to the known multi-token model if an unknown name is requested so
+    # the returned patch tokens stay comparable with the collection's doc tokens.
+    if model_name not in MODEL_CONFIGS:
+        print(f"Unknown model '{model_name}' for patch embeddings; falling back to '{MULTI_FPS_MODEL_NAME}'.")
+        model_name = MULTI_FPS_MODEL_NAME
+
+    if model_name not in LOADED_MODELS:
+        try:
+            print(f"Loading model {model_name} on demand (patch embeddings)...")
+            LOADED_MODELS[model_name] = load_model(model_name)
+        except Exception as e:
+            print(f"Error loading model {model_name}: {e}")
+            return {
+                "model_name": model_name,
+                "embedding_method": "mean pooling",
+                "data": None,
+                "patch_tokens": None,
+                "grid": None,
+                "num_patches": None,
+                "message": f"Error loading model '{model_name}'. Check logs."
+            }
+
+    model = LOADED_MODELS[model_name]
+    preprocess = get_preprocess(model_name)
+
+    try:
+        with torch.no_grad():
+            img_tensor = preprocess(image_pil).unsqueeze(0)  # [1, C, 224, 224]
+            features = model.forward_features(img_tensor)
+            if isinstance(features, tuple):
+                features = features[0]
+
+            if len(features.shape) != 3 or features.shape[1] <= 1:
+                return {
+                    "model_name": model_name,
+                    "embedding_method": "mean pooling",
+                    "data": None,
+                    "patch_tokens": None,
+                    "grid": None,
+                    "num_patches": None,
+                    "message": f"Error: Unexpected feature output shape from model '{model_name}'. Check logs."
+                }
+
+            patch_tokens = features[:, 1:]  # (1, P, D) drop CLS token
+            num_patches = patch_tokens.shape[1]
+
+            # First-stage dense vector: mean pool over patch tokens, L2-normalized
+            dense = F.normalize(patch_tokens.mean(dim=1), p=2, dim=1)  # (1, D)
+
+            # Per-patch tokens, L2-normalized so cosine == dot for the MaxSim query
+            patch_norm = F.normalize(patch_tokens, p=2, dim=-1)  # (1, P, D)
+
+            grid = int(round(num_patches ** 0.5))
+
+            data = dense.squeeze(0).cpu().numpy().tolist()
+            patch_list = patch_norm.squeeze(0).cpu().numpy().tolist()
+
+        return {
+            "model_name": model_name,
+            "embedding_method": "mean pooling",
+            "data": data,
+            "patch_tokens": patch_list,
+            "grid": grid,
+            "num_patches": num_patches,
+            "message": "Success"
+        }
+
+    except Exception as e:
+        print(f"Error computing patch embeddings with model {model_name}: {e}")
+        import traceback
+        traceback.print_exc()
+        return {
+            "model_name": model_name,
+            "embedding_method": "mean pooling",
+            "data": None,
+            "patch_tokens": None,
+            "grid": None,
+            "num_patches": None,
+            "message": f"Error processing image with model '{model_name}'. Check logs."
+        }
+
 # --- Gradio Interface ---
 EXAMPLE_DIR = "examples"
 EXAMPLE_IMAGE = os.path.join(EXAMPLE_DIR, "sample_image.png")
@@ -410,6 +522,25 @@ with gr.Blocks() as iface:
         outputs=output_json
     )
 
+    # --- Per-patch (ROI) embeddings: additive endpoint, existing flow untouched ---
+    gr.Markdown("### Patch / Region-of-Interest Embeddings")
+    gr.Markdown(
+        "Returns per-patch token embeddings (grid x grid) plus a mean-pooled dense "
+        "vector, for region-of-interest late-interaction search."
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            patch_button = gr.Button("Calculate Patch Embeddings (ROI)")
+        with gr.Column(scale=2):
+            patch_output_json = gr.JSON(label="Patch Embeddings (JSON)")
+
+    patch_button.click(
+        fn=get_patch_embeddings,
+        inputs=[input_image, model_selector],
+        outputs=patch_output_json,
+        api_name="get_patch_embeddings",
+    )
+
 # --- Launch the Gradio App ---
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0")