Support JPEG and PNG image inputs

ac44754 17 days ago

15 kB

#!/usr/bin/env python3
"""Run standalone GeoText-1652 image/text retrieval on an RGB GeoTIFF.
Example:
    IMAGE_URL="https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif"
    python infer.py \
      --image-url "$IMAGE_URL" \
      --query "a building" \
      --query "a sports field"
"""

from __future__ import annotations

import argparse
from io import BytesIO
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import requests
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from rasterio.io import MemoryFile
from ruamel.yaml import YAML
from torchvision import transforms

MODEL_REPOSITORY = "truemanv5666/GeoText1652_model"
CHECKPOINT_NAME = "geotext_official_checkpoint.pth"
DEFAULT_IMAGE_URL = (
    "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/"
    "geoembeddings-demo/building-detection_sm.tif"
)
CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--image-url", default=DEFAULT_IMAGE_URL, help="Public RGB GeoTIFF/COG URL."
    )
    parser.add_argument(
        "--query",
        action="append",
        required=True,
        help="Text query to score. Pass this option repeatedly to rank multiple queries.",
    )
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=Path(
            os.environ.get(
                "GEOTEXT_SOURCE_DIR", Path(__file__).resolve().parent / "GeoText-1652"
            )
        ),
        help="Checkout of https://github.com/MultimodalGeo/GeoText-1652.",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path(os.environ.get("GEOTEXT_MODEL_DIR", Path(__file__).resolve().parent)),
        help="Directory holding the Hugging Face checkpoint. Downloaded on first use.",
    )
    parser.add_argument("--device", default="auto", choices=("auto", "cuda", "cpu"))
    parser.add_argument("--output", type=Path, help="Optional JSON output file.")
    return parser.parse_args()


def resolve_device(requested: str) -> torch.device:
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    if requested == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but no CUDA device is available.")
    return torch.device(requested)


def ensure_checkpoint(model_dir: Path) -> Path:
    checkpoint = model_dir / CHECKPOINT_NAME
    if not checkpoint.exists():
        model_dir.mkdir(parents=True, exist_ok=True)
        snapshot_download(repo_id=MODEL_REPOSITORY, local_dir=model_dir)
    if not checkpoint.exists():
        raise FileNotFoundError(f"Expected checkpoint was not downloaded: {checkpoint}")
    return checkpoint


def fetch_rgb_geotiff(url: str) -> Tuple[Image.Image, Dict[str, Any]]:
    """Fetch RGB GeoTIFF/COG, JPEG, or PNG imagery from a public URL.
    The historical function name is kept so existing server imports remain
    compatible. JPEG and PNG inputs have no CRS or bounds, so their geospatial
    metadata fields are returned as ``None``.
    """
    response = requests.get(url, timeout=180)
    response.raise_for_status()

    try:
        with MemoryFile(response.content) as memory_file:
            with memory_file.open() as dataset:
                if dataset.count < 3:
                    raise ValueError(
                        f"GeoText requires RGB imagery; input has {dataset.count} band(s)."
                    )
                rgb = dataset.read([1, 2, 3], masked=True).astype(np.float32).filled(np.nan)
                metadata = {
                    "crs": str(dataset.crs),
                    "bounds": list(dataset.bounds),
                    "width": dataset.width,
                    "height": dataset.height,
                    "bands_used": [1, 2, 3],
                    "format": "geotiff",
                    "is_geospatial": True,
                }
        low, high = np.nanpercentile(rgb, (2, 98))
        if not np.isfinite(low) or not np.isfinite(high) or high <= low:
            raise ValueError("Unable to create an RGB image from the GeoTIFF pixel values.")
        stretched = np.nan_to_num((rgb - low) / (high - low), nan=0.0, posinf=1.0, neginf=0.0)
        stretched = np.clip(stretched, 0, 1)
        return Image.fromarray((stretched.transpose(1, 2, 0) * 255).astype(np.uint8)), metadata
    except Exception:
        pass

    try:
        with Image.open(BytesIO(response.content)) as source:
            image_format = source.format
            if image_format not in ("JPEG", "PNG"):
                raise ValueError(f"unsupported image format: {image_format}")
            image = source.convert("RGB").copy()
    except Exception as error:
        raise ValueError("image_url could not be decoded as GeoTIFF, JPEG, or PNG") from error
    return image, {
        "crs": None,
        "bounds": None,
        "width": image.width,
        "height": image.height,
        "bands_used": ["red", "green", "blue"],
        "format": image_format.lower(),
        "is_geospatial": False,
    }


def tile_image(
    image: Image.Image, tile_size: int, overlap: int
) -> List[Dict[str, Any]]:
    """Split an image into overlapping source-pixel tiles without padding."""
    if tile_size <= 0:
        raise ValueError("tile_size must be greater than zero")
    if overlap < 0 or overlap >= tile_size:
        raise ValueError("tile_overlap must satisfy 0 <= tile_overlap < tile_size")

    width, height = image.size
    stride = tile_size - overlap
    tiles = []
    tile_index = 0
    for row in range(0, height, stride):
        for column in range(0, width, stride):
            tile_width = min(tile_size, width - column)
            tile_height = min(tile_size, height - row)
            tiles.append(
                {
                    "column": column,
                    "height": tile_height,
                    "image": image.crop((column, row, column + tile_width, row + tile_height)),
                    "index": tile_index,
                    "row": row,
                    "width": tile_width,
                }
            )
            tile_index += 1
    return tiles


def load_model(source_dir: Path, model_dir: Path, device: torch.device):
    method_dir = source_dir / "Method"
    if not method_dir.is_dir():
        raise FileNotFoundError(f"GeoText source was not found at {method_dir}")
    sys.path.insert(0, str(method_dir))

    from models.model_re_bbox import Bench  # noqa: PLC0415
    from models.tokenization_bert import BertTokenizer  # noqa: PLC0415

    with (method_dir / "configs" / "re_bbox.yaml").open() as file:
        config = YAML(typ="safe").load(file)
    config["vision_config"] = str(method_dir / "configs" / "config_swinB_384.json")
    config["text_config"] = str(method_dir / "configs" / "config_bert.json")
    config["text_encoder"] = str(model_dir / "bert")

    tokenizer = BertTokenizer.from_pretrained(config["text_encoder"])
    model = Bench(config)
    model.load_pretrained(str(ensure_checkpoint(model_dir)), config, is_eval=True)
    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,
    tokenizer,
    config: Dict[str, Any],
    image: Image.Image,
    queries: List[str],
    device: torch.device,
    include_embeddings: bool = False,
    include_bboxes: bool = False,
    include_patch_embeddings: bool = False,
) -> Dict[str, Any]:
    image_tensor = _image_tensor(image, config, device)
    text_input = tokenizer(
        queries,
        padding="max_length",
        truncation=True,
        max_length=config["max_tokens"],
        return_tensors="pt",
    ).to(device)

    image_embeds, _ = model.get_vision_embeds(image_tensor)
    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()
    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_patch_embeddings:
        # GeoText's vision encoder returns one CLS token followed by its spatial
        # patch tokens.  Project every patch with the same retrieval head used
        # for the global CLS embedding, so their 256-D vectors share the text
        # embedding space.
        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)
            x1 = image.width * column / columns
            y1 = image.height * row / rows
            x2 = image.width * (column + 1) / columns
            y2 = image.height * (row + 1) / rows
            patches.append(
                {
                    "patch_index": index,
                    "row": row,
                    "column": column,
                    "source_pixel_xyxy": [x1, y1, x2, y2],
                    "embedding": embedding,
                }
            )
        result["patch_embeddings"] = {
            "embedding_dimension": patch_features.shape[-1],
            "grid": {"rows": rows, "columns": columns},
            "patches": patches,
        }

    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:
    args = parse_args()
    device = resolve_device(args.device)
    image, image_metadata = fetch_rgb_geotiff(args.image_url)
    model, tokenizer, config = load_model(args.source_dir, args.model_dir, device)
    ranked_queries = rank_queries(model, tokenizer, config, image, args.query, device)
    result = {
        "model": MODEL_REPOSITORY,
        "image_url": args.image_url,
        "image_metadata": image_metadata,
        "device": str(device),
        "ranked_queries": ranked_queries,
    }
    output = json.dumps(result, indent=2)
    print(output)
    if args.output:
        args.output.write_text(output + "\n")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Run standalone GeoText-1652 image/text retrieval on an RGB GeoTIFF.

	Example:
	IMAGE_URL="https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif"
	python infer.py \
	--image-url "$IMAGE_URL" \
	--query "a building" \
	--query "a sports field"
	"""

	from __future__ import annotations

	import argparse
	from io import BytesIO
	import json
	import os
	import sys
	from pathlib import Path
	from typing import Any, Dict, List, Tuple

	import numpy as np
	import requests
	import torch
	from huggingface_hub import snapshot_download
	from PIL import Image
	from rasterio.io import MemoryFile
	from ruamel.yaml import YAML
	from torchvision import transforms

	MODEL_REPOSITORY = "truemanv5666/GeoText1652_model"
	CHECKPOINT_NAME = "geotext_official_checkpoint.pth"
	DEFAULT_IMAGE_URL = (
	"https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/"
	"geoembeddings-demo/building-detection_sm.tif"
	)
	CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
	CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument(
	"--image-url", default=DEFAULT_IMAGE_URL, help="Public RGB GeoTIFF/COG URL."
	)
	parser.add_argument(
	"--query",
	action="append",
	required=True,
	help="Text query to score. Pass this option repeatedly to rank multiple queries.",
	)
	parser.add_argument(
	"--source-dir",
	type=Path,
	default=Path(
	os.environ.get(
	"GEOTEXT_SOURCE_DIR", Path(__file__).resolve().parent / "GeoText-1652"
	)
	),
	help="Checkout of https://github.com/MultimodalGeo/GeoText-1652.",
	)
	parser.add_argument(
	"--model-dir",
	type=Path,
	default=Path(os.environ.get("GEOTEXT_MODEL_DIR", Path(__file__).resolve().parent)),
	help="Directory holding the Hugging Face checkpoint. Downloaded on first use.",
	)
	parser.add_argument("--device", default="auto", choices=("auto", "cuda", "cpu"))
	parser.add_argument("--output", type=Path, help="Optional JSON output file.")
	return parser.parse_args()


	def resolve_device(requested: str) -> torch.device:
	if requested == "auto":
	requested = "cuda" if torch.cuda.is_available() else "cpu"
	if requested == "cuda" and not torch.cuda.is_available():
	raise RuntimeError("CUDA was requested but no CUDA device is available.")
	return torch.device(requested)


	def ensure_checkpoint(model_dir: Path) -> Path:
	checkpoint = model_dir / CHECKPOINT_NAME
	if not checkpoint.exists():
	model_dir.mkdir(parents=True, exist_ok=True)
	snapshot_download(repo_id=MODEL_REPOSITORY, local_dir=model_dir)
	if not checkpoint.exists():
	raise FileNotFoundError(f"Expected checkpoint was not downloaded: {checkpoint}")
	return checkpoint


	def fetch_rgb_geotiff(url: str) -> Tuple[Image.Image, Dict[str, Any]]:
	"""Fetch RGB GeoTIFF/COG, JPEG, or PNG imagery from a public URL.

	The historical function name is kept so existing server imports remain
	compatible. JPEG and PNG inputs have no CRS or bounds, so their geospatial
	metadata fields are returned as ``None``.
	"""
	response = requests.get(url, timeout=180)
	response.raise_for_status()

	try:
	with MemoryFile(response.content) as memory_file:
	with memory_file.open() as dataset:
	if dataset.count < 3:
	raise ValueError(
	f"GeoText requires RGB imagery; input has {dataset.count} band(s)."
	)
	rgb = dataset.read([1, 2, 3], masked=True).astype(np.float32).filled(np.nan)
	metadata = {
	"crs": str(dataset.crs),
	"bounds": list(dataset.bounds),
	"width": dataset.width,
	"height": dataset.height,
	"bands_used": [1, 2, 3],
	"format": "geotiff",
	"is_geospatial": True,
	}
	low, high = np.nanpercentile(rgb, (2, 98))
	if not np.isfinite(low) or not np.isfinite(high) or high <= low:
	raise ValueError("Unable to create an RGB image from the GeoTIFF pixel values.")
	stretched = np.nan_to_num((rgb - low) / (high - low), nan=0.0, posinf=1.0, neginf=0.0)
	stretched = np.clip(stretched, 0, 1)
	return Image.fromarray((stretched.transpose(1, 2, 0) * 255).astype(np.uint8)), metadata
	except Exception:
	pass

	try:
	with Image.open(BytesIO(response.content)) as source:
	image_format = source.format
	if image_format not in ("JPEG", "PNG"):
	raise ValueError(f"unsupported image format: {image_format}")
	image = source.convert("RGB").copy()
	except Exception as error:
	raise ValueError("image_url could not be decoded as GeoTIFF, JPEG, or PNG") from error
	return image, {
	"crs": None,
	"bounds": None,
	"width": image.width,
	"height": image.height,
	"bands_used": ["red", "green", "blue"],
	"format": image_format.lower(),
	"is_geospatial": False,
	}


	def tile_image(
	image: Image.Image, tile_size: int, overlap: int
	) -> List[Dict[str, Any]]:
	"""Split an image into overlapping source-pixel tiles without padding."""
	if tile_size <= 0:
	raise ValueError("tile_size must be greater than zero")
	if overlap < 0 or overlap >= tile_size:
	raise ValueError("tile_overlap must satisfy 0 <= tile_overlap < tile_size")

	width, height = image.size
	stride = tile_size - overlap
	tiles = []
	tile_index = 0
	for row in range(0, height, stride):
	for column in range(0, width, stride):
	tile_width = min(tile_size, width - column)
	tile_height = min(tile_size, height - row)
	tiles.append(
	{
	"column": column,
	"height": tile_height,
	"image": image.crop((column, row, column + tile_width, row + tile_height)),
	"index": tile_index,
	"row": row,
	"width": tile_width,
	}
	)
	tile_index += 1
	return tiles


	def load_model(source_dir: Path, model_dir: Path, device: torch.device):
	method_dir = source_dir / "Method"
	if not method_dir.is_dir():
	raise FileNotFoundError(f"GeoText source was not found at {method_dir}")
	sys.path.insert(0, str(method_dir))

	from models.model_re_bbox import Bench # noqa: PLC0415
	from models.tokenization_bert import BertTokenizer # noqa: PLC0415

	with (method_dir / "configs" / "re_bbox.yaml").open() as file:
	config = YAML(typ="safe").load(file)
	config["vision_config"] = str(method_dir / "configs" / "config_swinB_384.json")
	config["text_config"] = str(method_dir / "configs" / "config_bert.json")
	config["text_encoder"] = str(model_dir / "bert")

	tokenizer = BertTokenizer.from_pretrained(config["text_encoder"])
	model = Bench(config)
	model.load_pretrained(str(ensure_checkpoint(model_dir)), config, is_eval=True)
	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,
	tokenizer,
	config: Dict[str, Any],
	image: Image.Image,
	queries: List[str],
	device: torch.device,
	include_embeddings: bool = False,
	include_bboxes: bool = False,
	include_patch_embeddings: bool = False,
	) -> Dict[str, Any]:
	image_tensor = _image_tensor(image, config, device)
	text_input = tokenizer(
	queries,
	padding="max_length",
	truncation=True,
	max_length=config["max_tokens"],
	return_tensors="pt",
	).to(device)

	image_embeds, _ = model.get_vision_embeds(image_tensor)
	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()
	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_patch_embeddings:
	# GeoText's vision encoder returns one CLS token followed by its spatial
	# patch tokens. Project every patch with the same retrieval head used
	# for the global CLS embedding, so their 256-D vectors share the text
	# embedding space.
	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)
	x1 = image.width * column / columns
	y1 = image.height * row / rows
	x2 = image.width * (column + 1) / columns
	y2 = image.height * (row + 1) / rows
	patches.append(
	{
	"patch_index": index,
	"row": row,
	"column": column,
	"source_pixel_xyxy": [x1, y1, x2, y2],
	"embedding": embedding,
	}
	)
	result["patch_embeddings"] = {
	"embedding_dimension": patch_features.shape[-1],
	"grid": {"rows": rows, "columns": columns},
	"patches": patches,
	}

	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:
	args = parse_args()
	device = resolve_device(args.device)
	image, image_metadata = fetch_rgb_geotiff(args.image_url)
	model, tokenizer, config = load_model(args.source_dir, args.model_dir, device)
	ranked_queries = rank_queries(model, tokenizer, config, image, args.query, device)
	result = {
	"model": MODEL_REPOSITORY,
	"image_url": args.image_url,
	"image_metadata": image_metadata,
	"device": str(device),
	"ranked_queries": ranked_queries,
	}
	output = json.dumps(result, indent=2)
	print(output)
	if args.output:
	args.output.write_text(output + "\n")


	if __name__ == "__main__":
	main()