| """Hugging Face Inference Endpoints custom handler for GeoText-1652.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any, Dict |
|
|
| import torch |
| from infer import fetch_rgb_geotiff, load_model, rank_queries |
|
|
|
|
| class EndpointHandler: |
| """Rank natural-language descriptions against one public RGB GeoTIFF/COG.""" |
|
|
| def __init__(self, path: str = "") -> None: |
| self.repository_dir = Path(path) if path else Path(__file__).resolve().parent |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model, self.tokenizer, self.config = load_model( |
| self.repository_dir / "GeoText-1652", self.repository_dir, self.device |
| ) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| inputs = data.get("inputs", data) |
| if not isinstance(inputs, dict): |
| raise ValueError("Request inputs must be a JSON object.") |
|
|
| image_url = inputs.get("image_url") or inputs.get("image") |
| if not isinstance(image_url, str) or not image_url: |
| raise ValueError("inputs.image_url must be a public GeoTIFF/COG URL.") |
|
|
| queries = inputs.get("queries", inputs.get("query")) |
| if isinstance(queries, str): |
| queries = [queries] |
| if not isinstance(queries, list) or not queries or not all( |
| isinstance(query, str) and query for query in queries |
| ): |
| raise ValueError("inputs.query or inputs.queries must contain one or more strings.") |
|
|
| image, image_metadata = fetch_rgb_geotiff(image_url) |
| return { |
| "model": "truemanv5666/GeoText1652_model", |
| "image_url": image_url, |
| "image_metadata": image_metadata, |
| "device": str(self.device), |
| "ranked_queries": rank_queries( |
| self.model, self.tokenizer, self.config, image, queries, self.device |
| ), |
| } |
|
|
