from __future__ import annotations import logging from pathlib import Path from typing import Any from PIL import Image from tqdm import tqdm try: from _bootstrap import add_project_root_to_path except ModuleNotFoundError: from scripts._bootstrap import add_project_root_to_path add_project_root_to_path() from src.config import PROJECT_ROOT, load_settings from src.dataset import load_metadata from src.vector_store import VectorStore logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s") logger = logging.getLogger(__name__) BATCH_SIZE = 32 def _absolute_image_path(path_value: str) -> Path: path = Path(path_value) return path if path.is_absolute() else PROJECT_ROOT / path def _flush_batch(vector_store: VectorStore, batch: list[dict[str, Any]]) -> int: if not batch: return 0 vector_store.upsert_many(batch) indexed = len(batch) batch.clear() return indexed def main() -> None: settings = load_settings() metadata = load_metadata(settings.metadata_csv) if metadata.empty: logger.warning("metadata.csv is empty. Add images and run `python scripts/build_metadata.py` first.") return from src.clip_embedder import ClipEmbedder vector_store = VectorStore(settings) embedder = ClipEmbedder(settings.clip_model_name) batch: list[dict[str, Any]] = [] indexed = 0 failed = 0 for row in tqdm(metadata.to_dict(orient="records"), desc="Indexing images"): image_path = _absolute_image_path(str(row["path"])) if not image_path.exists(): logger.warning("Skipping missing image: %s", image_path) failed += 1 continue try: with Image.open(image_path) as opened: image = opened.convert("RGB") vector = embedder.encode_image(image) except Exception as exc: logger.warning("Skipping damaged or unreadable image %s: %s", image_path, exc) failed += 1 continue metadata_payload = { "path": str(row["path"]), "filename": str(row["filename"]), "category": str(row["category"]), } batch.append({"id": str(row["id"]), "vector": vector, "metadata": metadata_payload}) if len(batch) >= BATCH_SIZE: indexed += _flush_batch(vector_store, batch) indexed += _flush_batch(vector_store, batch) logger.info("Indexing finished. Successful: %d, failed: %d", indexed, failed) if __name__ == "__main__": main()