scripts/index_images.py

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()