Spaces:
Sleeping
Sleeping
| from pathlib import Path | |
| import hashlib | |
| import numpy as np | |
| from PIL import Image | |
| from .config import Settings | |
| class ImageEmbedder: | |
| def __init__(self, settings: Settings): | |
| self.settings = settings | |
| self._model = None | |
| self._text_model = None | |
| self.mode = "fastembed" | |
| def _load_model(self): | |
| if self._model is not None: | |
| return self._model | |
| try: | |
| from fastembed import ImageEmbedding | |
| self._model = ImageEmbedding(model_name=self.settings.model_name) | |
| return self._model | |
| except Exception: | |
| if not self.settings.allow_fake_embeddings: | |
| raise | |
| self.mode = "deterministic-fallback" | |
| self._model = False | |
| return None | |
| def _load_text_model(self): | |
| if self._text_model is not None: | |
| return self._text_model | |
| from fastembed import TextEmbedding | |
| self._text_model = TextEmbedding(model_name=self.settings.text_model_name) | |
| return self._text_model | |
| def embed_text(self, query: str) -> list[float]: | |
| # CLIP text tower — shares the image vectors' 512-d space, so a text | |
| # query can search the visual memory directly. | |
| model = self._load_text_model() | |
| vector = next(iter(model.embed([query]))) | |
| return self._normalize(np.asarray(vector, dtype=np.float32)) | |
| def embed_path(self, path: Path) -> list[float]: | |
| if not path.exists(): | |
| raise FileNotFoundError(f"Image not found: {path}") | |
| model = self._load_model() | |
| if model is None: | |
| return self._fallback_embedding(path) | |
| image = Image.open(path).convert("RGB") | |
| vector = next(iter(model.embed([image]))) | |
| return self._normalize(np.asarray(vector, dtype=np.float32)) | |
| def _fallback_embedding(self, path: Path) -> list[float]: | |
| image = Image.open(path).convert("RGB").resize((64, 64)) | |
| arr = np.asarray(image, dtype=np.float32) / 255.0 | |
| means = arr.mean(axis=(0, 1)) | |
| stds = arr.std(axis=(0, 1)) | |
| hist = [] | |
| for channel in range(3): | |
| values, _ = np.histogram(arr[:, :, channel], bins=32, range=(0, 1), density=True) | |
| hist.extend(values.tolist()) | |
| digest = hashlib.sha256(path.read_bytes()).digest() | |
| noise = np.frombuffer(digest * 16, dtype=np.uint8)[:410].astype(np.float32) / 255.0 | |
| vector = np.concatenate([means, stds, np.array(hist, dtype=np.float32), noise]) | |
| return self._normalize(vector[: self.settings.vector_size]) | |
| def _normalize(self, vector: np.ndarray) -> list[float]: | |
| if vector.shape[0] != self.settings.vector_size: | |
| raise ValueError(f"Expected {self.settings.vector_size}-dim vector, got {vector.shape[0]}") | |
| norm = np.linalg.norm(vector) | |
| if norm == 0: | |
| return vector.tolist() | |
| return (vector / norm).astype(float).tolist() | |