Spaces:
Running on Zero
Running on Zero
| """Lazy local remote-sensing image captioning specialist.""" | |
| from __future__ import annotations | |
| import importlib.util | |
| import os | |
| import threading | |
| import time | |
| from typing import Any, Optional | |
| from PIL import Image | |
| from ..model_cache import model_cache_dir | |
| from ..models import ( | |
| CaptionResult, | |
| Confidence, | |
| ConfidenceLevel, | |
| ImageMetadata, | |
| Modality, | |
| ModelProvenance, | |
| SpecialistHealth, | |
| ) | |
| CHECKPOINT = "Gurveer05/blip-image-captioning-base-rscid-finetuned" | |
| CHECKPOINT_REVISION = "57c4f0ff06264d374131637ce7c1e8dea079d589" | |
| BASE_ARCHITECTURE = "Salesforce BLIP image-captioning base" | |
| ADAPTATION_DATASET = "RSICD (Remote Sensing Image Caption Dataset)" | |
| MODEL_LICENSE = "Apache-2.0" | |
| MODEL_SOURCE = f"https://huggingface.co/{CHECKPOINT}/tree/{CHECKPOINT_REVISION}" | |
| SUPPORTED_MODALITIES = {Modality.OPTICAL, Modality.MULTISPECTRAL} | |
| MAX_NEW_TOKENS = 40 | |
| NUM_BEAMS = 3 | |
| LIMITATIONS = [ | |
| "The caption is model-generated and may omit or misidentify small objects.", | |
| "The checkpoint was adapted on 224×224 optical aerial imagery from RSICD.", | |
| "The result is not a substitute for expert remote-sensing validation.", | |
| ] | |
| class CaptionerError(RuntimeError): | |
| def __init__(self, code: str, message: str) -> None: | |
| super().__init__(message) | |
| self.code = code | |
| self.message = message | |
| def captioner_enabled() -> bool: | |
| return os.getenv("SATQUERY_CAPTIONER_ENABLED", "1").strip().lower() not in {"0", "false", "no", "off"} | |
| def captioner_dependencies_available() -> bool: | |
| return all(importlib.util.find_spec(name) is not None for name in ("torch", "transformers", "safetensors")) | |
| def captioner_configured() -> bool: | |
| return captioner_enabled() and bool(os.getenv("SATQUERY_CAPTIONER_CHECKPOINT", CHECKPOINT).strip()) and captioner_dependencies_available() | |
| class RemoteSensingCaptioner: | |
| def __init__(self) -> None: | |
| self._lock = threading.Lock() | |
| self._state = "unloaded" | |
| self._device: Optional[str] = None | |
| self._processor: Optional[Any] = None | |
| self._model: Optional[Any] = None | |
| self._safe_error: Optional[str] = None | |
| self._load_count = 0 | |
| def checkpoint(self) -> str: | |
| return os.getenv("SATQUERY_CAPTIONER_CHECKPOINT", CHECKPOINT).strip() | |
| def cache_dir(self) -> str: | |
| return str(model_cache_dir()) | |
| def checkpoint_revision(self) -> str | None: | |
| return CHECKPOINT_REVISION if self.checkpoint == CHECKPOINT else None | |
| def limitations(self) -> list[str]: | |
| return list(LIMITATIONS) | |
| def is_available(self) -> bool: | |
| return captioner_configured() and self._state != "failed" | |
| def health(self) -> SpecialistHealth: | |
| if not captioner_configured(): | |
| return SpecialistHealth(status="disabled", device=None, error=None) | |
| return SpecialistHealth(status=self._state, device=self._device, error=self._safe_error) | |
| def _select_device(self, torch: Any) -> str: | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| mps = getattr(torch.backends, "mps", None) | |
| if mps is not None and mps.is_available(): | |
| return "mps" | |
| return "cpu" | |
| def load(self) -> None: | |
| if self._state == "ready": | |
| return | |
| if not captioner_configured(): | |
| raise CaptionerError("CAPTIONER_UNAVAILABLE", "The remote-sensing captioner is not configured in this local environment.") | |
| with self._lock: | |
| if self._state == "ready": | |
| return | |
| if self._state == "failed": | |
| raise CaptionerError("CAPTIONER_LOAD_FAILED", "The local captioning checkpoint could not be loaded.") | |
| self._state = "loading" | |
| self._safe_error = None | |
| try: | |
| import torch | |
| from transformers import BlipForConditionalGeneration, BlipProcessor | |
| device = self._select_device(torch) | |
| load_options: dict[str, Any] = {"cache_dir": self.cache_dir} | |
| if self.checkpoint_revision is not None: | |
| load_options.update(revision=self.checkpoint_revision, token=False) | |
| processor = BlipProcessor.from_pretrained(self.checkpoint, **load_options) | |
| model = BlipForConditionalGeneration.from_pretrained(self.checkpoint, **load_options) | |
| model.eval() | |
| model.to(device) | |
| self._processor = processor | |
| self._model = model | |
| self._device = device | |
| self._load_count += 1 | |
| self._state = "ready" | |
| except Exception as error: | |
| self._processor = None | |
| self._model = None | |
| self._device = None | |
| self._state = "failed" | |
| self._safe_error = "Checkpoint loading failed." | |
| raise CaptionerError("CAPTIONER_LOAD_FAILED", "The local captioning checkpoint could not be loaded.") from error | |
| def describe( | |
| self, | |
| image: Image.Image, | |
| metadata: ImageMetadata, | |
| modality: Modality, | |
| bands_used: list[str], | |
| image_representation: str, | |
| ) -> CaptionResult: | |
| if modality not in SUPPORTED_MODALITIES: | |
| raise CaptionerError( | |
| "UNSUPPORTED_CAPTION_MODALITY", | |
| "The connected RSICD captioner supports optical and multispectral RGB representations only; SAR captioning is not implemented.", | |
| ) | |
| if metadata.band_count < 3: | |
| raise CaptionerError( | |
| "UNSUPPORTED_CAPTION_BANDS", | |
| "The connected captioner requires an RGB or RGB-like optical representation with at least three bands.", | |
| ) | |
| if metadata.band_count > 3 and len(metadata.selected_visual_bands) < 3: | |
| raise CaptionerError( | |
| "UNSUPPORTED_CAPTION_BANDS", | |
| "Captioning requires an explicit RGB band mapping for multispectral rasters; no mapping was guessed.", | |
| ) | |
| was_ready = self._state == "ready" | |
| load_started = time.perf_counter() | |
| self.load() | |
| model_load_ms = 0 if was_ready else max(0, round((time.perf_counter() - load_started) * 1000)) | |
| if self._processor is None or self._model is None or self._device is None: | |
| raise CaptionerError("CAPTIONER_UNAVAILABLE", "The local captioning specialist is unavailable.") | |
| import torch | |
| inference_started = time.perf_counter() | |
| try: | |
| inputs = self._processor(images=image.convert("RGB"), return_tensors="pt") | |
| inputs = {name: value.to(self._device) for name, value in inputs.items()} | |
| with torch.inference_mode(): | |
| tokens = self._model.generate( | |
| **inputs, | |
| max_new_tokens=MAX_NEW_TOKENS, | |
| num_beams=NUM_BEAMS, | |
| do_sample=False, | |
| ) | |
| caption = self._processor.decode(tokens[0], skip_special_tokens=True).strip() | |
| except Exception as error: | |
| raise CaptionerError("CAPTION_INFERENCE_FAILED", "The local captioning model could not generate a description safely.") from error | |
| runtime_ms = max(0, round((time.perf_counter() - inference_started) * 1000)) | |
| if not caption: | |
| raise CaptionerError("EMPTY_CAPTION", "The captioning model returned an empty description.") | |
| warnings = list(metadata.warnings) | |
| confidence_level = ConfidenceLevel.MODERATE | |
| if modality == Modality.MULTISPECTRAL and bands_used == ["band_1", "band_2", "band_3"]: | |
| confidence_level = ConfidenceLevel.LOW | |
| warnings.append("Multispectral band roles were unavailable; the model used the first three display bands as an RGB-like representation.") | |
| confidence = Confidence( | |
| level=confidence_level, | |
| score=None, | |
| reason=( | |
| "The captioning model does not provide calibrated probability estimates. " | |
| "Confidence reflects supported optical input and band-mapping certainty only." | |
| ), | |
| ) | |
| return CaptionResult( | |
| caption=caption, | |
| confidence=confidence, | |
| model=ModelProvenance( | |
| tool_id="rs_captioner", | |
| checkpoint=self.checkpoint, | |
| base_architecture=BASE_ARCHITECTURE, | |
| adaptation_dataset=ADAPTATION_DATASET, | |
| remote_sensing_adapted=True, | |
| license=MODEL_LICENSE, | |
| source=MODEL_SOURCE, | |
| ), | |
| warnings=list(dict.fromkeys(warnings)), | |
| runtime_ms=runtime_ms, | |
| device=self._device, | |
| image_representation=image_representation, | |
| bands_used=bands_used, | |
| model_load_ms=model_load_ms, | |
| reused_model=was_ready, | |
| ) | |
| def reset_for_tests(self) -> None: | |
| with self._lock: | |
| self._state = "unloaded" | |
| self._device = None | |
| self._processor = None | |
| self._model = None | |
| self._safe_error = None | |
| self._load_count = 0 | |
| _CAPTIONER = RemoteSensingCaptioner() | |
| def get_captioner() -> RemoteSensingCaptioner: | |
| return _CAPTIONER | |