embeddings/image_encoder.py

"""Image embedding client for the local embedding HTTP service."""
import os
import logging
from typing import Any, List, Optional, Union
import numpy as np
import requests
from PIL import Image
logger = logging.getLogger(__name__)
from config.services_config import get_embedding_image_base_url
from config.env_config import REDIS_CONFIG
from embeddings.cache_keys import build_image_cache_key
from embeddings.redis_embedding_cache import RedisEmbeddingCache
class CLIPImageEncoder:
    """
    Image Encoder for generating image embeddings using network service.
    This client is stateless and safe to instantiate per caller.
    """
    def __init__(self, service_url: Optional[str] = None):
        resolved_url = (
            service_url
            or os.getenv("EMBEDDING_IMAGE_SERVICE_URL")
            or os.getenv("EMBEDDING_SERVICE_URL")
            or get_embedding_image_base_url()
        )
        self.service_url = str(resolved_url).rstrip("/")
        self.endpoint = f"{self.service_url}/embed/image"
        # Reuse embedding cache prefix, but separate namespace for images to avoid collisions.
        self.cache_prefix = str(REDIS_CONFIG.get("embedding_cache_prefix", "embedding")).strip() or "embedding"
        logger.info("Creating CLIPImageEncoder instance with service URL: %s", self.service_url)
        self.cache = RedisEmbeddingCache(
            key_prefix=self.cache_prefix,
            namespace="image",
        )
    def _call_service(self, request_data: List[str], normalize_embeddings: bool = True) -> List[Any]:
        """
        Call the embedding service API.
        Args:
            request_data: List of image URLs / local file paths
        Returns:
            List of embeddings (list[float]) or nulls (None), aligned to input order
        """
        try:
            response = requests.post(
                self.endpoint,
                params={"normalize": "true" if normalize_embeddings else "false"},
                json=request_data,
                timeout=60
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"CLIPImageEncoder service request failed: {e}", exc_info=True)
            raise
    def encode_image(self, image: Image.Image) -> np.ndarray:
        """
        Encode image to embedding vector using network service.
        Note: This method is kept for compatibility but the service only works with URLs.
        """
        raise NotImplementedError("encode_image with PIL Image is not supported by embedding service")
    def encode_image_from_url(self, url: str, normalize_embeddings: bool = True) -> np.ndarray:
        """
        Generate image embedding via network service using URL.
        Args:
            url: Image URL to process
        Returns:
            Embedding vector
        """
        cache_key = build_image_cache_key(url, normalize=normalize_embeddings)
        cached = self.cache.get(cache_key)
        if cached is not None:
            return cached
        response_data = self._call_service([url], normalize_embeddings=normalize_embeddings)
        if not response_data or len(response_data) != 1 or response_data[0] is None:
            raise RuntimeError(f"No image embedding returned for URL: {url}")
        vec = np.array(response_data[0], dtype=np.float32)
        if vec.ndim != 1 or vec.size == 0 or not np.isfinite(vec).all():
            raise RuntimeError(f"Invalid image embedding returned for URL: {url}")
        self.cache.set(cache_key, vec)
        return vec
    def encode_batch(
        self,
        images: List[Union[str, Image.Image]],
        batch_size: int = 8,
        normalize_embeddings: bool = True,
    ) -> List[np.ndarray]:
        """
        Encode a batch of images efficiently via network service.
        Args:
            images: List of image URLs or PIL Images
            batch_size: Batch size for processing (used for service requests)
        Returns:
            List of embeddings
        """
        for i, img in enumerate(images):
            if isinstance(img, Image.Image):
                raise NotImplementedError(f"PIL Image at index {i} is not supported by service")
            if not isinstance(img, str) or not img.strip():
                raise ValueError(f"Invalid image URL/path at index {i}: {img!r}")
        results: List[np.ndarray] = []
        pending_urls: List[str] = []
        pending_positions: List[int] = []
        normalized_urls = [str(u).strip() for u in images]  # type: ignore[list-item]
        for pos, url in enumerate(normalized_urls):
            cache_key = build_image_cache_key(url, normalize=normalize_embeddings)
            cached = self.cache.get(cache_key)
            if cached is not None:
                results.append(cached)
                continue
            results.append(np.array([], dtype=np.float32))  # placeholder
            pending_positions.append(pos)
            pending_urls.append(url)
        for i in range(0, len(pending_urls), batch_size):
            batch_urls = pending_urls[i : i + batch_size]
            response_data = self._call_service(batch_urls, normalize_embeddings=normalize_embeddings)
            if not response_data or len(response_data) != len(batch_urls):
                raise RuntimeError(
                    f"Image embedding response length mismatch: expected {len(batch_urls)}, "
                    f"got {0 if response_data is None else len(response_data)}"
                )
            for j, url in enumerate(batch_urls):
                embedding = response_data[j]
                if embedding is None:
                    raise RuntimeError(f"No image embedding returned for URL: {url}")
                vec = np.array(embedding, dtype=np.float32)
                if vec.ndim != 1 or vec.size == 0 or not np.isfinite(vec).all():
                    raise RuntimeError(f"Invalid image embedding returned for URL: {url}")
                self.cache.set(build_image_cache_key(url, normalize=normalize_embeddings), vec)
                pos = pending_positions[i + j]
                results[pos] = vec
        return results
    def encode_image_urls(
        self,
        urls: List[str],
        batch_size: Optional[int] = None,
        normalize_embeddings: bool = True,
    ) -> List[np.ndarray]:
        """
        与 ClipImageModel / ClipAsServiceImageEncoder 一致的接口，供索引器 document_transformer 调用。
        Args:
            urls: 图片 URL 列表
            batch_size: 批大小（默认 8）
        Returns:
            与 urls 等长的向量列表
        """
        return self.encode_batch(
            urls,
            batch_size=batch_size or 8,
            normalize_embeddings=normalize_embeddings,
        )